直接上代码
import sys
import cv2
import numpy as np
from PyQt5.QtWidgets import (QApplication, QMainWindow, QFileDialog,
QLabel, QPushButton, QVBoxLayout, QWidget,
QHBoxLayout, QComboBox, QSlider)
from PyQt5.QtGui import QPixmap, QImage, QPainter, QPen, QColor
from PyQt5.QtCore import Qt, QPoint
class ImageMattingApp(QMainWindow):
def __init__(self):
super().__init__()
self.setWindowTitle("Python交互式抠图软件")
self.setGeometry(100, 100, 1000, 800)
# 初始化变量
self.image = None
self.mask = None
self.display_image = None
self.drawing = False
self.last_point = QPoint()
self.brush_size = 5
self.current_tool = "foreground" # foreground/background/eraser
self.grabcut_iterations = 5
self.initUI()
def initUI(self):
# 创建主部件和布局
main_widget = QWidget()
layout = QVBoxLayout()
# 图像显示标签
self.image_label = QLabel()
self.image_label.setAlignment(Qt.AlignCenter)
self.image_label.setMinimumSize(800, 600)
self.image_label.mousePressEvent = self.mousePressEvent
self.image_label.mouseMoveEvent = self.mouseMoveEvent
self.image_label.mouseReleaseEvent = self.mouseReleaseEvent
# 工具选择区域
tool_layout = QHBoxLayout()
self.tool_combo = QComboBox()
self.tool_combo.addItems(["标记前景", "标记背景", "橡皮擦"])
self.tool_combo.currentIndexChanged.connect(self.changeTool)
self.brush_slider = QSlider(Qt.Horizontal)
self.brush_slider.setRange(1, 30)
self.brush_slider.setValue(self.brush_size)
self.brush_slider.valueChanged.connect(self.changeBrushSize)
tool_layout.addWidget(QLabel("工具:"))
tool_layout.addWidget(self.tool_combo)
tool_layout.addWidget(QLabel("笔刷大小:"))
tool_layout.addWidget(self.brush_slider)
# 按钮区域
button_layout = QHBoxLayout()
self.load_btn = QPushButton("加载图片")
self.load_btn.clicked.connect(self.load_image)
self.reset_btn = QPushButton("重置标记")
self.reset_btn.clicked.connect(self.resetMarking)
self.reset_btn.setEnabled(False)
self.matting_btn = QPushButton("执行抠图")
self.matting_btn.clicked.connect(self.applyGrabCut)
self.matting_btn.setEnabled(False)
self.save_btn = QPushButton("保存结果")
self.save_btn.clicked.connect(self.save_result)
self.save_btn.setEnabled(False)
button_layout.addWidget(self.load_btn)
button_layout.addWidget(self.reset_btn)
button_layout.addWidget(self.matting_btn)
button_layout.addWidget(self.save_btn)
# 添加到主布局
layout.addWidget(self.image_label)
layout.addLayout(tool_layout)
layout.addLayout(button_layout)
main_widget.setLayout(layout)
self.setCentralWidget(main_widget)
def load_image(self):
file_path, _ = QFileDialog.getOpenFileName(
self, "选择图片", "", "图片文件 (*.png *.jpg *.jpeg *.bmp)"
)
if file_path:
self.image = cv2.imread(file_path)
if self.image is not None:
self.image = cv2.cvtColor(self.image, cv2.COLOR_BGR2RGB)
self.display_image = self.image.copy()
# 初始化mask和GrabCut工作区
self.mask = np.zeros(self.image.shape[:2], np.uint8)
self.bgd_model = np.zeros((1, 65), np.float64)
self.fgd_model = np.zeros((1, 65), np.float64)
# 显示图片
self.updateDisplay()
self.reset_btn.setEnabled(True)
self.matting_btn.setEnabled(True)
def changeTool(self, index):
tools = ["foreground", "background", "eraser"]
self.current_tool = tools[index]
def changeBrushSize(self, value):
self.brush_size = value
def resetMarking(self):
if self.image is not None:
self.mask = np.zeros(self.image.shape[:2], np.uint8)
self.display_image = self.image.copy()
self.updateDisplay()
def mousePressEvent(self, event):
if event.button() == Qt.LeftButton and self.image is not None:
pos = self.get_image_position(event.pos())
if pos:
self.drawing = True
self.last_point = QPoint(pos[0], pos[1])
self.drawOnMask(self.last_point)
def mouseMoveEvent(self, event):
if self.drawing and self.image is not None:
pos = self.get_image_position(event.pos())
if pos:
current_point = QPoint(pos[0], pos[1])
self.drawOnMask(current_point)
self.last_point = current_point
def mouseReleaseEvent(self, event):
if event.button() == Qt.LeftButton:
self.drawing = False
def drawOnMask(self, point):
if self.image is None:
return
# 转换为图像坐标
x = point.x()
y = point.y()
# 确保坐标在图像范围内
x = max(0, min(x, self.image.shape[1] - 1))
y = max(0, min(y, self.image.shape[0] - 1))
# 根据当前工具设置mask值
if self.current_tool == "foreground":
cv2.circle(self.mask, (x, y), self.brush_size, cv2.GC_FGD, -1)
cv2.circle(self.display_image, (x, y), self.brush_size, (0, 255, 0), -1)
elif self.current_tool == "background":
cv2.circle(self.mask, (x, y), self.brush_size, cv2.GC_BGD, -1)
cv2.circle(self.display_image, (x, y), self.brush_size, (255, 0, 0), -1)
elif self.current_tool == "eraser":
cv2.circle(self.mask, (x, y), self.brush_size, 0, -1)
# 恢复原始图像区域
self.display_image[y - self.brush_size:y + self.brush_size,
x - self.brush_size:x + self.brush_size] = \
self.image[y - self.brush_size:y + self.brush_size,
x - self.brush_size:x + self.brush_size]
self.updateDisplay()
def applyGrabCut(self):
if self.image is not None:
# 如果没有用户标记,使用自动矩形选择
if np.all(self.mask == 0):
height, width = self.image.shape[:2]
rect = (width // 4, height // 4, width // 2, height // 2)
cv2.grabCut(self.image, self.mask, rect,
self.bgd_model, self.fgd_model,
self.grabcut_iterations, cv2.GC_INIT_WITH_RECT)
else:
# 使用用户标记
cv2.grabCut(self.image, self.mask, None,
self.bgd_model, self.fgd_model,
self.grabcut_iterations, cv2.GC_INIT_WITH_MASK)
# 创建最终掩码
final_mask = np.where((self.mask == 2) | (self.mask == 0), 0, 1).astype('uint8')
# 应用掩码
self.display_image = self.image * final_mask[:, :, np.newaxis]
self.updateDisplay()
self.save_btn.setEnabled(True)
def updateDisplay(self):
if self.display_image is not None:
height, width, channel = self.display_image.shape
bytes_per_line = 3 * width
q_img = QImage(
self.display_image.data, width, height, bytes_per_line, QImage.Format_RGB888
)
pixmap = QPixmap.fromImage(q_img)
# 保持宽高比缩放
scaled_pixmap = pixmap.scaled(
self.image_label.size(),
Qt.KeepAspectRatio,
Qt.SmoothTransformation
)
self.image_label.setPixmap(scaled_pixmap)
def save_result(self):
if self.image is not None and self.display_image is not None:
file_path, _ = QFileDialog.getSaveFileName(
self, "保存图片", "", "PNG图片 (*.png);;JPEG图片 (*.jpg *.jpeg)"
)
if file_path:
# 创建透明背景的结果
result = cv2.cvtColor(self.image, cv2.COLOR_RGB2RGBA)
final_mask = np.where((self.mask == 2) | (self.mask == 0), 0, 255).astype('uint8')
result[:, :, 3] = final_mask
cv2.imwrite(file_path, cv2.cvtColor(result, cv2.COLOR_RGBA2BGRA))
def get_image_position(self, event_pos):
"""
精确计算鼠标位置对应的图像坐标
解决鼠标在上但绘制在右下角的问题
"""
if self.image is None or self.image_label.pixmap() is None:
return None
# 获取QLabel和实际图像的尺寸
label_size = self.image_label.size()
pixmap = self.image_label.pixmap()
pixmap_size = pixmap.size()
# 计算图像在QLabel中的显示位置(考虑居中显示)
x_offset = (label_size.width() - pixmap_size.width()) // 2
y_offset = (label_size.height() - pixmap_size.height()) // 2
# 转换为图像坐标
x = event_pos.x() - x_offset
y = event_pos.y() - y_offset
# 检查是否在图像范围内
if x < 0 or y < 0 or x >= pixmap_size.width() or y >= pixmap_size.height():
return None
# 计算缩放比例后的实际图像坐标
scale_x = self.image.shape[1] / pixmap_size.width()
scale_y = self.image.shape[0] / pixmap_size.height()
actual_x = int(x * scale_x)
actual_y = int(y * scale_y)
# 确保不越界
actual_x = max(0, min(actual_x, self.image.shape[1] - 1))
actual_y = max(0, min(actual_y, self.image.shape[0] - 1))
print(f"界面坐标: ({event_pos.x()}, {event_pos.y()})")
print(f"显示偏移: ({x_offset}, {y_offset})")
print(f"显示图像坐标: ({x}, {y})")
print(f"实际图像坐标: ({actual_x}, {actual_y})")
print(f"图像尺寸: {self.image.shape[1]}x{self.image.shape[0]}")
print(f"显示尺寸: {pixmap_size.width()}x{pixmap_size.height()}")
return actual_x, actual_y
if __name__ == "__main__":
app = QApplication(sys.argv)
window = ImageMattingApp()
window.show()
sys.exit(app.exec_())