基于fastSAM的自动标注小程序

import sys
from PyQt5.QtWidgets import QApplication, QLabel, QWidget, QVBoxLayout
from PyQt5.QtCore import Qt
from PyQt5.QtGui import QPixmap, QImage
from PIL import Image
from PIL import Image, ImageDraw
import numpy as np
from fastsam import FastSAM


class ImageLabel(QLabel):
	def __init__(self, parent=None):
		super().__init__(parent)
		self.setMouseTracking(True)
		self.original_pil_image = None
		self.pil_image = None

	def set_image(self, pil_image):
		self.original_pil_image = pil_image.copy()
		self.pil_image = pil_image
		self.update_qimage()

	def update_qimage(self):
		if self.pil_image is not None:
			image_np = np.array(self.pil_image.convert("RGB"))
			height, width, channel = image_np.shape
			bytes_per_line = 3 * width
			q_image = QImage(image_np.data, width, height, bytes_per_line, QImage.Format_RGB888)
			self.setPixmap(QPixmap.fromImage(q_image))

	def mouseMoveEvent(self, event):
		x = event.x()
		y = event.y()
		self.parent().update_mouse_position(x, y)
		self.add_red_dot_to_image(x, y)

	def add_red_dot_to_image(self, x, y):
		if self.original_pil_image is not None:
			self.pil_image = self.original_pil_image.copy()
			self.pil_image = self.pil_image.convert('RGBA')
			if self.parent().display_mask is not None:
				mask = self.parent().display_mask
				# 创建红色图层
				red_layer = np.zeros((mask.shape[0], mask.shape[1], 4), dtype=np.uint8)
				red_layer[:, :, 0] = 255  # 红色通道全满
				red_layer[:, :, 3] = 100 * mask  # Alpha通道，使用mask控制透明度
				red_image = Image.fromarray(red_layer, 'RGBA')
				self.pil_image = Image.alpha_composite(self.pil_image, red_image)
				self.pil_image = self.pil_image.convert('RGB')  # 转换为RGB格式

			draw = ImageDraw.Draw(self.pil_image)
			radius = 5
			left_up_point = (x - radius, y - radius)
			right_down_point = (x + radius, y + radius)
			draw.ellipse([left_up_point, right_down_point], fill=(255, 0, 0))
			self.update_qimage()


class MainWindow(QWidget):
	def __init__(self):
		super().__init__()

		self.model = FastSAM("./weights/FastSAM.pt")
		self.input = Image.open('/home/robotics/FastSAM/daikons.jpg')
		self.input = self.input.convert("RGB")
		self.input = self.input.resize((1280, 720))
		everything_results = self.model(
			self.input,
			device='cuda',
			retina_masks=True,
			imgsz=1024,
			conf=0.4,
			iou=0.9
			)
		self.masks = everything_results[0].masks.data.cpu().numpy()
		self.display_mask = None

		self.initUI()

	def initUI(self):
		self.setWindowTitle('Display Mouse Position in Image')
		self.setGeometry(100, 100, 800, 600)

		# 创建一个标签来显示图像
		self.image_label = ImageLabel(self)
		self.image_label.setAlignment(Qt.AlignCenter)

		# 创建一个标签来显示鼠标位置
		self.position_label = QLabel(self)
		self.position_label.setAlignment(Qt.AlignCenter)

		layout = QVBoxLayout()
		layout.addWidget(self.image_label)
		layout.addWidget(self.position_label)
		self.setLayout(layout)

		self.display_pil_image(self.input)

	def display_pil_image(self, pil_image):
		self.image_label.set_image(pil_image)
		width, height = pil_image.size
		self.image_label.setFixedSize(width, height)

	def update_mouse_position(self, x, y):
		self.position_label.setText(f'Mouse Position: ({x}, {y})')
		self.display_mask = None
		for i in range(self.masks.shape[0]):
			mask = self.masks[i, :, :].astype(np.uint8)
			if mask[y, x] == 1:
				self.display_mask = mask
				break


if __name__ == '__main__':
	app = QApplication(sys.argv)
	main_window = MainWindow()
	main_window.show()
	sys.exit(app.exec_())