379 lines
18 KiB
Python
379 lines
18 KiB
Python
|
import os
|
||
|
|
import sys
|
||
|
|
import cv2
|
||
|
|
import numpy as np
|
||
|
|
from PyQt5 import QtCore, QtGui, QtWidgets
|
||
|
|
from segment_anything import sam_model_registry, SamPredictor
|
||
|
|
|
||
|
|
class Ui_MainWindow(object):
|
||
|
|
def setupUi(self, MainWindow):
|
||
|
|
MainWindow.setObjectName("MainWindow")
|
||
|
|
MainWindow.resize(1140, 450)
|
||
|
|
MainWindow.setMinimumSize(QtCore.QSize(1140, 450))
|
||
|
|
MainWindow.setMaximumSize(QtCore.QSize(1140, 450))
|
||
|
|
self.centralwidget = QtWidgets.QWidget(MainWindow)
|
||
|
|
self.centralwidget.setObjectName("centralwidget")
|
||
|
|
self.pushButton_w = QtWidgets.QPushButton(self.centralwidget)
|
||
|
|
self.pushButton_w.setGeometry(QtCore.QRect(10, 90, 151, 51))
|
||
|
|
self.pushButton_w.setObjectName("pushButton_w")
|
||
|
|
self.pushButton_a = QtWidgets.QPushButton(self.centralwidget)
|
||
|
|
self.pushButton_a.setGeometry(QtCore.QRect(10, 160, 71, 51))
|
||
|
|
self.pushButton_a.setObjectName("pushButton_a")
|
||
|
|
self.pushButton_d = QtWidgets.QPushButton(self.centralwidget)
|
||
|
|
self.pushButton_d.setGeometry(QtCore.QRect(90, 160, 71, 51))
|
||
|
|
self.pushButton_d.setObjectName("pushButton_d")
|
||
|
|
self.pushButton_s = QtWidgets.QPushButton(self.centralwidget)
|
||
|
|
self.pushButton_s.setGeometry(QtCore.QRect(10, 360, 151, 51))
|
||
|
|
self.pushButton_s.setObjectName("pushButton_s")
|
||
|
|
self.pushButton_5 = QtWidgets.QPushButton(self.centralwidget)
|
||
|
|
self.pushButton_5.setGeometry(QtCore.QRect(10, 230, 151, 51))
|
||
|
|
self.pushButton_5.setObjectName("pushButton_5")
|
||
|
|
self.label_orign = QtWidgets.QLabel(self.centralwidget)
|
||
|
|
self.label_orign.setGeometry(QtCore.QRect(180, 20, 471, 401))
|
||
|
|
self.label_orign.setStyleSheet("background-color: rgb(255, 255, 255);")
|
||
|
|
self.label_orign.setObjectName("label_orign")
|
||
|
|
self.label_2 = QtWidgets.QLabel(self.centralwidget)
|
||
|
|
self.label_2.setGeometry(QtCore.QRect(660, 20, 471, 401))
|
||
|
|
self.label_2.setStyleSheet("background-color: rgb(255, 255, 255);")
|
||
|
|
self.label_2.setObjectName("label_2")
|
||
|
|
self.pushButton_w_2 = QtWidgets.QPushButton(self.centralwidget)
|
||
|
|
self.pushButton_w_2.setGeometry(QtCore.QRect(10, 20, 151, 51))
|
||
|
|
self.pushButton_w_2.setObjectName("pushButton_w_2")
|
||
|
|
self.lineEdit = QtWidgets.QLineEdit(self.centralwidget)
|
||
|
|
self.lineEdit.setGeometry(QtCore.QRect(50, 290, 81, 21))
|
||
|
|
self.lineEdit.setObjectName("lineEdit")
|
||
|
|
self.horizontalSlider = QtWidgets.QSlider(self.centralwidget)
|
||
|
|
self.horizontalSlider.setGeometry(QtCore.QRect(10, 320, 141, 22))
|
||
|
|
self.horizontalSlider.setSliderPosition(50)
|
||
|
|
self.horizontalSlider.setOrientation(QtCore.Qt.Horizontal)
|
||
|
|
self.horizontalSlider.setTickInterval(0)
|
||
|
|
self.horizontalSlider.setObjectName("horizontalSlider")
|
||
|
|
MainWindow.setCentralWidget(self.centralwidget)
|
||
|
|
self.menubar = QtWidgets.QMenuBar(MainWindow)
|
||
|
|
self.menubar.setGeometry(QtCore.QRect(0, 0, 1140, 23))
|
||
|
|
self.menubar.setObjectName("menubar")
|
||
|
|
MainWindow.setMenuBar(self.menubar)
|
||
|
|
self.statusbar = QtWidgets.QStatusBar(MainWindow)
|
||
|
|
self.statusbar.setObjectName("statusbar")
|
||
|
|
MainWindow.setStatusBar(self.statusbar)
|
||
|
|
|
||
|
|
self.retranslateUi(MainWindow)
|
||
|
|
QtCore.QMetaObject.connectSlotsByName(MainWindow)
|
||
|
|
|
||
|
|
def retranslateUi(self, MainWindow):
|
||
|
|
_translate = QtCore.QCoreApplication.translate
|
||
|
|
MainWindow.setWindowTitle(_translate("MainWindow", "MainWindow"))
|
||
|
|
self.pushButton_w.setText(_translate("MainWindow", "Predict"))
|
||
|
|
self.pushButton_a.setText(_translate("MainWindow", "Pre"))
|
||
|
|
self.pushButton_d.setText(_translate("MainWindow", "Next"))
|
||
|
|
self.pushButton_s.setText(_translate("MainWindow", "Save"))
|
||
|
|
self.pushButton_5.setText(_translate("MainWindow", "背景图"))
|
||
|
|
self.label_orign.setText(_translate("MainWindow", "<html><head/><body><p align=\"center\">原始图像</p></body></html>"))
|
||
|
|
self.label_2.setText(_translate("MainWindow", "<html><head/><body><p align=\"center\">预测图像</p></body></html>"))
|
||
|
|
self.pushButton_w_2.setText(_translate("MainWindow", "Openimg"))
|
||
|
|
self.lineEdit.setText(_translate("MainWindow", "改变mask大小"))
|
||
|
|
|
||
|
|
class MyMainWindow(QtWidgets.QMainWindow, Ui_MainWindow):
|
||
|
|
def __init__(self):
|
||
|
|
super().__init__()
|
||
|
|
self.setupUi(self)
|
||
|
|
|
||
|
|
self.image_path = ""
|
||
|
|
self.image_folder = ""
|
||
|
|
self.image_files = []
|
||
|
|
self.current_index = 0
|
||
|
|
self.input_stop = False # 在这里初始化 input_stop
|
||
|
|
|
||
|
|
self.pushButton_w_2.clicked.connect(self.open_image_folder)
|
||
|
|
self.pushButton_a.clicked.connect(self.load_previous_image)
|
||
|
|
self.pushButton_d.clicked.connect(self.load_next_image)
|
||
|
|
|
||
|
|
|
||
|
|
def open_image_folder(self):
|
||
|
|
folder_dialog = QtWidgets.QFileDialog()
|
||
|
|
folder_path = folder_dialog.getExistingDirectory(self, 'Open Image Folder', '')
|
||
|
|
if folder_path:
|
||
|
|
self.image_folder = folder_path
|
||
|
|
self.image_files = self.get_image_files(self.image_folder)
|
||
|
|
if self.image_files:
|
||
|
|
self.show_image_selection_dialog()
|
||
|
|
|
||
|
|
def load_previous_image(self):
|
||
|
|
if self.image_files:
|
||
|
|
if self.current_index > 0:
|
||
|
|
self.current_index -= 1
|
||
|
|
else:
|
||
|
|
self.current_index = len(self.image_files) - 1
|
||
|
|
self.show_image()
|
||
|
|
|
||
|
|
def load_next_image(self):
|
||
|
|
if self.image_files:
|
||
|
|
if self.current_index < len(self.image_files) - 1:
|
||
|
|
self.current_index += 1
|
||
|
|
else:
|
||
|
|
self.current_index = 0
|
||
|
|
self.show_image()
|
||
|
|
|
||
|
|
def get_image_files(self, folder_path):
|
||
|
|
image_files = [file for file in os.listdir(folder_path) if file.endswith(('png', 'jpg', 'jpeg', 'bmp'))]
|
||
|
|
return image_files
|
||
|
|
|
||
|
|
def show_image_selection_dialog(self):
|
||
|
|
dialog = QtWidgets.QDialog(self)
|
||
|
|
dialog.setWindowTitle("Select Image")
|
||
|
|
layout = QtWidgets.QVBoxLayout()
|
||
|
|
|
||
|
|
self.listWidget = QtWidgets.QListWidget()
|
||
|
|
for image_file in self.image_files:
|
||
|
|
item = QtWidgets.QListWidgetItem(image_file)
|
||
|
|
pixmap = QtGui.QPixmap(os.path.join(self.image_folder, image_file)).scaledToWidth(100)
|
||
|
|
item.setIcon(QtGui.QIcon(pixmap))
|
||
|
|
self.listWidget.addItem(item)
|
||
|
|
self.listWidget.itemDoubleClicked.connect(self.image_selected)
|
||
|
|
layout.addWidget(self.listWidget)
|
||
|
|
|
||
|
|
buttonBox = QtWidgets.QDialogButtonBox(QtWidgets.QDialogButtonBox.Ok | QtWidgets.QDialogButtonBox.Cancel)
|
||
|
|
buttonBox.accepted.connect(self.image_selected)
|
||
|
|
buttonBox.rejected.connect(dialog.reject)
|
||
|
|
layout.addWidget(buttonBox)
|
||
|
|
|
||
|
|
dialog.setLayout(layout)
|
||
|
|
|
||
|
|
dialog.exec_()
|
||
|
|
|
||
|
|
def image_selected(self):
|
||
|
|
selected_item = self.listWidget.currentItem()
|
||
|
|
if selected_item:
|
||
|
|
selected_index = self.listWidget.currentRow()
|
||
|
|
if selected_index >= 0 and selected_index < len(self.image_files): # 检查索引是否在有效范围内
|
||
|
|
self.current_index = selected_index
|
||
|
|
self.show_image() # 显示所选图像
|
||
|
|
# 调用OpenCV窗口显示
|
||
|
|
self.call_opencv_interaction(os.path.join(self.image_folder, self.image_files[self.current_index]))
|
||
|
|
|
||
|
|
def show_image(self):
|
||
|
|
if self.image_files and self.current_index < len(self.image_files):
|
||
|
|
file_path = os.path.join(self.image_folder, self.image_files[self.current_index])
|
||
|
|
pixmap = QtGui.QPixmap(file_path)
|
||
|
|
self.label_orign.setPixmap(pixmap.scaled(self.label_orign.size(), QtCore.Qt.KeepAspectRatio))
|
||
|
|
|
||
|
|
def call_opencv_interaction(self, image_path):
|
||
|
|
input_dir = os.path.dirname(image_path)
|
||
|
|
image_orign = cv2.imread(image_path)
|
||
|
|
output_dir = r'D:\Program Files\Pycharm items\segment-anything-model\scripts\output\maskt'
|
||
|
|
crop_mode = True
|
||
|
|
|
||
|
|
print('最好是每加一个点就按w键predict一次')
|
||
|
|
os.makedirs(output_dir, exist_ok=True)
|
||
|
|
image_files = [self.image_files[self.current_index]]
|
||
|
|
|
||
|
|
sam = sam_model_registry["vit_b"](
|
||
|
|
checkpoint=r"D:\Program Files\Pycharm items\segment-anything-model\weights\vit_b.pth")
|
||
|
|
_ = sam.to(device="cuda")
|
||
|
|
predictor = SamPredictor(sam)
|
||
|
|
|
||
|
|
WINDOW_WIDTH = 1280
|
||
|
|
WINDOW_HEIGHT = 720
|
||
|
|
|
||
|
|
def apply_mask(image, mask, alpha_channel=True):
|
||
|
|
if alpha_channel:
|
||
|
|
alpha = np.zeros_like(image[..., 0])
|
||
|
|
alpha[mask == 1] = 255
|
||
|
|
image = cv2.merge((image[..., 0], image[..., 1], image[..., 2], alpha))
|
||
|
|
else:
|
||
|
|
image = np.where(mask[..., None] == 1, image, 0)
|
||
|
|
return image
|
||
|
|
|
||
|
|
def apply_color_mask(image, mask, color, color_dark=0.5):
|
||
|
|
for c in range(3):
|
||
|
|
image[:, :, c] = np.where(mask == 1, image[:, :, c] * (1 - color_dark) + color_dark * color[c],
|
||
|
|
image[:, :, c])
|
||
|
|
return image
|
||
|
|
|
||
|
|
def get_next_filename(base_path, filename):
|
||
|
|
name, ext = os.path.splitext(filename)
|
||
|
|
for i in range(1, 101):
|
||
|
|
new_name = f"{name}_{i}{ext}"
|
||
|
|
if not os.path.exists(os.path.join(base_path, new_name)):
|
||
|
|
return new_name
|
||
|
|
return None
|
||
|
|
|
||
|
|
def save_masked_image(image, mask, output_dir, filename, crop_mode_):
|
||
|
|
# 保存图像到指定路径
|
||
|
|
if crop_mode_:
|
||
|
|
# 如果采用了裁剪模式，则裁剪图像
|
||
|
|
y, x = np.where(mask)
|
||
|
|
y_min, y_max, x_min, x_max = y.min(), y.max(), x.min(), x.max()
|
||
|
|
cropped_mask = mask[y_min:y_max + 1, x_min:x_max + 1]
|
||
|
|
cropped_image = image[y_min:y_max + 1, x_min:x_max + 1]
|
||
|
|
masked_image = apply_mask(cropped_image, cropped_mask)
|
||
|
|
else:
|
||
|
|
masked_image = apply_mask(image, mask)
|
||
|
|
filename = filename[:filename.rfind('.')] + '.png'
|
||
|
|
new_filename = get_next_filename(output_dir, filename)
|
||
|
|
|
||
|
|
if new_filename:
|
||
|
|
if masked_image.shape[-1] == 4:
|
||
|
|
cv2.imwrite(os.path.join(output_dir, new_filename), masked_image,
|
||
|
|
[cv2.IMWRITE_PNG_COMPRESSION, 9])
|
||
|
|
else:
|
||
|
|
cv2.imwrite(os.path.join(output_dir, new_filename), masked_image)
|
||
|
|
print(f"Saved as {new_filename}")
|
||
|
|
|
||
|
|
# 读取保存的图像文件
|
||
|
|
saved_image_path = os.path.join(output_dir, new_filename)
|
||
|
|
saved_image_pixmap = QtGui.QPixmap(saved_image_path)
|
||
|
|
|
||
|
|
# 将保存的图像显示在预测图像区域
|
||
|
|
mainWindow.label_2.setPixmap(
|
||
|
|
saved_image_pixmap.scaled(mainWindow.label_2.size(), QtCore.Qt.KeepAspectRatio))
|
||
|
|
else:
|
||
|
|
print("Could not save the image. Too many variations exist.")
|
||
|
|
|
||
|
|
current_index = 0
|
||
|
|
|
||
|
|
cv2.namedWindow("image", cv2.WINDOW_NORMAL)
|
||
|
|
cv2.resizeWindow("image", WINDOW_WIDTH, WINDOW_HEIGHT)
|
||
|
|
cv2.moveWindow("image", (1920 - WINDOW_WIDTH) // 2, (1080 - WINDOW_HEIGHT) // 2)
|
||
|
|
|
||
|
|
def mouse_click(event, x, y, flags, param):
|
||
|
|
if not self.input_stop:
|
||
|
|
if event == cv2.EVENT_LBUTTONDOWN:
|
||
|
|
input_point.append([x, y])
|
||
|
|
input_label.append(1)
|
||
|
|
elif event == cv2.EVENT_RBUTTONDOWN:
|
||
|
|
input_point.append([x, y])
|
||
|
|
input_label.append(0)
|
||
|
|
else:
|
||
|
|
if event == cv2.EVENT_LBUTTONDOWN or event == cv2.EVENT_RBUTTONDOWN:
|
||
|
|
print('此时不能添加点,按w退出mask选择模式')
|
||
|
|
|
||
|
|
cv2.setMouseCallback("image", mouse_click)
|
||
|
|
input_point = []
|
||
|
|
input_label = []
|
||
|
|
input_stop = False
|
||
|
|
while True:
|
||
|
|
filename = self.image_files[self.current_index]
|
||
|
|
image_orign = cv2.imread(os.path.join(input_dir, filename))
|
||
|
|
image_crop = image_orign.copy()
|
||
|
|
image = cv2.cvtColor(image_orign.copy(), cv2.COLOR_BGR2RGB)
|
||
|
|
selected_mask = None
|
||
|
|
logit_input = None
|
||
|
|
while True:
|
||
|
|
image_display = image_orign.copy()
|
||
|
|
display_info = f'{filename} '
|
||
|
|
cv2.putText(image_display, display_info, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (0, 255, 255),
|
||
|
|
2,
|
||
|
|
cv2.LINE_AA)
|
||
|
|
for point, label in zip(input_point, input_label):
|
||
|
|
color = (0, 255, 0) if label == 1 else (0, 0, 255)
|
||
|
|
cv2.circle(image_display, tuple(point), 5, color, -1)
|
||
|
|
if selected_mask is not None:
|
||
|
|
color = tuple(np.random.randint(0, 256, 3).tolist())
|
||
|
|
selected_image = apply_color_mask(image_display, selected_mask, color)
|
||
|
|
|
||
|
|
cv2.imshow("image", image_display)
|
||
|
|
|
||
|
|
key = cv2.waitKey(1)
|
||
|
|
|
||
|
|
if key == ord(" "):
|
||
|
|
input_point = []
|
||
|
|
input_label = []
|
||
|
|
selected_mask = None
|
||
|
|
logit_input = None
|
||
|
|
elif key == ord("w"):
|
||
|
|
input_stop = True
|
||
|
|
if len(input_point) > 0 and len(input_label) > 0:
|
||
|
|
|
||
|
|
predictor.set_image(image)
|
||
|
|
input_point_np = np.array(input_point)
|
||
|
|
input_label_np = np.array(input_label)
|
||
|
|
|
||
|
|
masks, scores, logits = predictor.predict(
|
||
|
|
point_coords=input_point_np,
|
||
|
|
point_labels=input_label_np,
|
||
|
|
mask_input=logit_input[None, :, :] if logit_input is not None else None,
|
||
|
|
multimask_output=True,
|
||
|
|
)
|
||
|
|
|
||
|
|
mask_idx = 0
|
||
|
|
num_masks = len(masks)
|
||
|
|
|
||
|
|
prediction_window_name = "Prediction"
|
||
|
|
cv2.namedWindow(prediction_window_name, cv2.WINDOW_NORMAL)
|
||
|
|
cv2.resizeWindow(prediction_window_name, WINDOW_WIDTH, WINDOW_HEIGHT)
|
||
|
|
cv2.moveWindow(prediction_window_name, (1920 - WINDOW_WIDTH) // 2,
|
||
|
|
(1080 - WINDOW_HEIGHT) // 2)
|
||
|
|
|
||
|
|
while True:
|
||
|
|
color = tuple(np.random.randint(0, 256, 3).tolist())
|
||
|
|
image_select = image_orign.copy()
|
||
|
|
selected_mask = masks[mask_idx]
|
||
|
|
selected_image = apply_color_mask(image_select, selected_mask, color)
|
||
|
|
mask_info = f'Total: {num_masks} | Current: {mask_idx} | Score: {scores[mask_idx]:.2f} | w 预测 | d 切换下一个 | a 切换上一个 |q 移除最后一个'
|
||
|
|
cv2.putText(selected_image, mask_info, (10, 30), cv2.FONT_HERSHEY_SIMPLEX, 0.7,
|
||
|
|
(0, 255, 255), 2, cv2.LINE_AA)
|
||
|
|
|
||
|
|
cv2.imshow(prediction_window_name, selected_image)
|
||
|
|
|
||
|
|
key = cv2.waitKey(10)
|
||
|
|
if key == ord('q') and len(input_point) > 0:
|
||
|
|
input_point.pop(-1)
|
||
|
|
elif key == ord('s'):
|
||
|
|
save_masked_image(image_crop, selected_mask, output_dir, filename,
|
||
|
|
crop_mode_=crop_mode)
|
||
|
|
elif key == ord('a'):
|
||
|
|
if mask_idx > 0:
|
||
|
|
mask_idx -= 1
|
||
|
|
else:
|
||
|
|
mask_idx = num_masks - 1
|
||
|
|
elif key == ord('d'):
|
||
|
|
if mask_idx < num_masks - 1:
|
||
|
|
mask_idx += 1
|
||
|
|
else:
|
||
|
|
mask_idx = 0
|
||
|
|
elif key == ord('w'):
|
||
|
|
input_stop = False # Allow adding points again
|
||
|
|
break
|
||
|
|
elif key == ord(" "):
|
||
|
|
input_point = []
|
||
|
|
input_label = []
|
||
|
|
selected_mask = None
|
||
|
|
logit_input = None
|
||
|
|
break
|
||
|
|
logit_input = logits[mask_idx, :, :]
|
||
|
|
print('max score:', np.argmax(scores), ' select:', mask_idx)
|
||
|
|
|
||
|
|
elif key == ord('a'):
|
||
|
|
current_index = max(0, current_index - 1)
|
||
|
|
input_point = []
|
||
|
|
input_label = []
|
||
|
|
break
|
||
|
|
elif key == ord('d'):
|
||
|
|
current_index = min(len(image_files) - 1, current_index + 1)
|
||
|
|
input_point = []
|
||
|
|
input_label = []
|
||
|
|
break
|
||
|
|
elif key == 27:
|
||
|
|
break
|
||
|
|
elif key == ord('q') and len(input_point) > 0:
|
||
|
|
input_point.pop(-1)
|
||
|
|
input_label.pop(-1)
|
||
|
|
elif key == ord('s') and selected_mask is not None:
|
||
|
|
save_masked_image(image_crop, selected_mask, output_dir, filename, crop_mode_=crop_mode)
|
||
|
|
|
||
|
|
if key == 27:
|
||
|
|
break
|
||
|
|
|
||
|
|
cv2.destroyAllWindows() # Close all windows before exiting
|
||
|
|
if key == 27:
|
||
|
|
break
|
||
|
|
|
||
|
|
if __name__ == "__main__":
|
||
|
|
app = QtWidgets.QApplication(sys.argv)
|
||
|
|
mainWindow = MyMainWindow()
|
||
|
|
mainWindow.show()
|
||
|
|
sys.exit(app.exec_())
|
||
|
|
|
||
|
|
|