import cv2
import numpy as np
import os

class DrawRectangle:

    def __init__(self, div_scale):
        self.mask_save_dir = './mask'
        self.scale = div_scale
        self.current_roi_points = []
        self.rois = []
        self.window_title = "Image - Select ROI"
        self.draw_complete = False
        pass

    def callback(self, event, x, y, flags, param):

        drawing_image = param['Image']
        # 左键添加感兴趣点
        if event == cv2.EVENT_LBUTTONDOWN:
            # 添加感兴趣点
            self.current_roi_points.append((x, y))
            # 绘制标记点
            cv2.circle(drawing_image, (x, y), 4, (0, 255, 0), -1)
            # 绘制多边形
            if len(self.current_roi_points) > 1:
                cv2.line(drawing_image, self.current_roi_points[-2], self.current_roi_points[-1], (0, 255, 0), 2)
            # 显示图像
            cv2.imshow(self.window_title, drawing_image)
            print(f'添加感兴趣点：{y}行, {x}列')
        # 右键闭合感兴趣区域
        if event == cv2.EVENT_RBUTTONDOWN:
            if len(self.current_roi_points) < 3:
                print("[提示] ROI 至少需要 3 个点构成多边形!")
                return
            cv2.line(drawing_image, self.current_roi_points[-1], self.current_roi_points[0], (0, 255, 0), 2)
            cv2.imshow(self.window_title, drawing_image)
            # 清理
            self.rois.append(self.current_roi_points)
            print(f'添加感兴趣区，包含点数:{len(self.current_roi_points)}个')
            self.current_roi_points = []


    def draw(self, img_path: str):
        """在输入图像中绘制多边形区域，然后生成相应的mask图片"""
        # 读取图像
        ori_img = cv2.imread(img_path)
        mask_base_name = os.path.splitext(os.path.basename(img_path))[0] + '.png'
        img = cv2.resize(ori_img, (ori_img.shape[1] // self.scale, ori_img.shape[0] // self.scale))
        if img is None:
            raise RuntimeError('Cannot read the image!')
        param = {'Image': img}
        cv2.namedWindow(self.window_title)
        cv2.setMouseCallback(self.window_title, self.callback, param=param)


        # 显示图像并等待退出
        while True:
            cv2.imshow(self.window_title, img)
            key = cv2.waitKey(1) & 0xFF
            if key == ord('q') or key == 27:  # 按'q'或ESC键退出
                break

        # 为原图生成掩膜
        mask = np.zeros((ori_img.shape[0], ori_img.shape[1]),dtype=np.uint8)  # shape等于原始输入图像
        for roi in self.rois:
            roi_points = np.array(roi, np.int32).reshape((-1, 1, 2)) * self.scale  # 兴趣点的缩放处理
            cv2.fillPoly(mask, [roi_points], 255)

        # 保存掩膜图像
        if not os.path.exists(self.mask_save_dir):
            os.makedirs(self.mask_save_dir)
        cv2.imwrite(os.path.join(self.mask_save_dir, mask_base_name), mask)
        # cv2.imshow("mask", mask)
        # cv2.waitKey(0)
        cv2.destroyAllWindows()

if __name__ == '__main__':
    drawer = DrawRectangle(2)
    drawer.draw(r"D:\code\water_turbidity_det\draw_mask\mask\4_device_capture.jpg")