roi 视频版

import cv2
import numpy as np
import json
import os


class ROIDrawer:
    def __init__(self, image_o, label="beizi_5"):
        self.drawing = False
        self.ix, self.iy = -1, -1
        self.rois = []  # 存储多个ROI
        self.image_o = image_o
        self.image = self.image_o.copy()
        self.temp_image = self.image.copy()
        self.ok = False
        self.label = label  # 目标标签

    def draw_crosshair(self, event, x, y, flags, param):
        self.temp_image = self.image.copy()  # 每次更新临时图像

        # 关键修改：无论是否有已框选的ROI，只要不在绘制中就显示十字星
        if not self.drawing:  # 只有当不处于拖拽框选状态时，才显示十字星
            cv2.line(self.temp_image, (x, 0), (x, self.temp_image.shape[0]), (0, 255, 0), 1)
            cv2.line(self.temp_image, (0, y), (self.temp_image.shape[1], y), (0, 255, 0), 1)

        # 鼠标按下：开始框选
        if event == cv2.EVENT_LBUTTONDOWN:
            self.drawing = True
            self.ix, self.iy = x, y

        # 鼠标移动：实时绘制矩形（此时不显示十字星，因为drawing=True）
        elif event == cv2.EVENT_MOUSEMOVE:
            if self.drawing:
                # 绘制当前正在拖拽的矩形
                cv2.rectangle(self.temp_image, (self.ix, self.iy), (x, y), (255, 0, 0), 2)

        # 鼠标左键释放：确认当前ROI（继续框选）
        elif event == cv2.EVENT_LBUTTONUP:
            self.drawing = False  # 结束绘制，恢复十字星显示
            # 计算规范的坐标（确保x1 < x2, y1 < y2）
            x1, y1 = min(self.ix, x), min(self.iy, y)
            x2, y2 = max(self.ix, x), max(self.iy, y)
            # 绘制最终矩形到原始图像
            cv2.rectangle(self.image, (x1, y1), (x2, y2), (255, 0, 0), 2)
            # 保存ROI坐标
            self.rois.append([[x1, y1], [x2, y2]])
            print(f"已添加ROI: {[x1, y1]} - {[x2, y2]} (共{len(self.rois)}个)")

        # 鼠标右键释放：完成框选
        elif event == cv2.EVENT_RBUTTONUP:
            self.drawing = False  # 结束绘制，恢复十字星显示
            x1, y1 = min(self.ix, x), min(self.iy, y)
            x2, y2 = max(self.ix, x), max(self.iy, y)
            cv2.rectangle(self.image, (x1, y1), (x2, y2), (255, 0, 0), 2)
            self.rois.append([[x1, y1], [x2, y2]])
            print(f"已添加ROI: {[x1, y1]} - {[x2, y2]} (共{len(self.rois)}个)")
            self.ok = True

    def save_json(self, image_path, output_json_path=None):
        if not self.rois:
            print("没有框选任何目标，不保存JSON")
            return

        if not output_json_path:
            img_dir, img_name = os.path.split(image_path)
            img_base = os.path.splitext(img_name)[0]
            output_json_path = os.path.join(img_dir, f"{img_base}.json")

        # 构建JSON结构
        json_data = {"version": "1.0.0", "flags": {}, "shapes": [], "imagePath": image_path, "imageHeight": self.image_o.shape[0], "imageWidth": self.image_o.shape[1]}

        for points in self.rois:
            (x1, y1), (x2, y2) = points
            bbox = [x1, y1, x2, y2] 

            shape = {"label": self.label, "shape_type": "rectangle", "points": points,  # 仍然保留原来的 [[x1,y1],[x2,y2]]
                "bbox": bbox,
                "description": "", "flags": {}}
            json_data["shapes"].append(shape)

        with open(output_json_path, 'w', encoding='utf-8') as f:
            json.dump(json_data, f, ensure_ascii=False, indent=4)
        print(f"JSON已保存至: {output_json_path}")

    def run(self, image_path, output_json=None):
        cv2.namedWindow('Draw ROI')
        cv2.setMouseCallback('Draw ROI', self.draw_crosshair)

        print("操作说明:")
        print("1. 左键拖拽框选目标（松开后继续框选下一个，十字星始终显示）")
        print("2. 右键拖拽框选最后一个目标（松开后结束框选）")
        print("3. 按Esc键取消操作，按Enter键保存并退出")

        while True:
            cv2.imshow('Draw ROI', self.temp_image)
            key = cv2.waitKey(1) & 0xFF

            if key == 27:  # Esc键：取消操作
                print("已取消操作")
                self.rois = []
                break
            elif key == 13:  # Enter键：保存并退出
                break
            elif self.ok:  # 右键结束框选
                break

        cv2.destroyAllWindows()
        self.save_json(image_path, output_json)
        return self.rois


if __name__ == '__main__':

    video_path=r"B:\data\tiaosheng\data_1116\test_1115.mp4"
    video_path=r"B:\data\tiaosheng\data_1116\label_test.mp4"
    json_path=os.path.splitext(video_path)[0]+".json"
    cap = cv2.VideoCapture(video_path)
    frame_rate = cap.get(cv2.CAP_PROP_FPS)
    loaded_frames = []
    prompts={}
    roi_box=True
    image_path="roi_box.jpg"
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        if len(prompts) < 1:
            if roi_box:
                roi_drawer = ROIDrawer(frame)
                selected_rois = roi_drawer.run(image_path, output_json=json_path)
                if len(selected_rois)>0:
                    selected_roi =selected_rois[0]
                    prompts[0] = ([selected_roi[0][0], selected_roi[0][1], selected_roi[1][0] + selected_roi[1][1]], 0)
                    print('boxes', prompts)
        else:
            break
            # else:
            #     mp = MultiPointDrawer(frame)
            #     points, labels = mp.run()
            #     prompts[0] = (points, labels)