roi_dir.py

import glob

import cv2
import numpy as np
import json
import os

class ROIDrawer:
    def __init__(self, image_o, label="tiaosheng"):
        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]])
        
        # 鼠标右键释放：完成框选
        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)
    
    def run(self, image_path, output_json=None):
        cv2.namedWindow('Draw ROI')
        cv2.setMouseCallback('Draw ROI', self.draw_crosshair)
        

        
        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()
        if output_json:
            self.save_json(image_path, output_json)
        return self.rois


def roi_video(video_path, output_json_path=None):

    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"
    
    video_result = {}
    video_name = os.path.basename(video_path)
    
    while True:
        ret, frame = cap.read()
        if not ret:
            break
        
        h, w = frame.shape[:2]
        target_area = 1000 * 1500
        orig_area = h * w
        if orig_area > target_area:
            scale = np.sqrt(target_area / orig_area)
            new_w = int(w * scale)
            new_h = int(h * scale)
            new_w -= new_w % 2
            new_h -= new_h % 2
            frame = cv2.resize(frame, (new_w, new_h), interpolation=cv2.INTER_LINEAR)
        
        if len(prompts) < 1:
            if roi_box:
                print(f"正在处理视频: {video_name}")
                roi_drawer = ROIDrawer(frame)
                selected_rois = roi_drawer.run(image_path, output_json=json_path)
                if len(selected_rois) > 0:
                    boxes = []
                    for p_i, points in enumerate(selected_rois):
                        (x1, y1), (x2, y2) = points
                        bbox = [x1, y1, x2, y2, p_i]
                        boxes.append(bbox)
                    
                    # 构建返回结果
                    video_result = {'video_name': video_name,  'boxes': boxes, 'frame_resolution': f"{new_w}x{new_h}" if orig_area > target_area else f"{w}x{h}"}
                    break  # 获取到boxes后退出循环
        else:
            break
    
    cap.release()
    return video_name, video_result

def roi_video_directory(directory_path, video_extensions=None):
    print("1. 左键拖拽框选目标（松开后继续框选下一个，十字星始终显示）")
    print("2. 右键拖拽框选最后一个目标（松开后结束框选）")
    print("3. 按Esc键取消操作，按Enter键保存并退出")
    
    video_files=glob.glob(directory_path+"/*.mp4")
    
    if not video_files:
        print(f"警告: 在目录 '{directory_path}' 中未找到视频文件")
        return {}
    
    print(f"找到 {len(video_files)} 个视频文件:")
    for i, video_file in enumerate(video_files, 1):
        print(f"{i}. {video_file}")
    
    results = {}
    
    for video_file in video_files:
        try:
            video_name, video_result = roi_video(str(video_file))
            if video_result:  # 只有当有结果时才添加
                results[video_name] = video_result
                print(results)
            else:
                results[video_name] = {'video_name': video_name, 'video_path': str(video_file), 'boxes': [],'frame_resolution': 'unknown'}
        except Exception as e:
            print(f"处理视频 '{video_file.name}' 时出错: {e}")
            results[video_file.name] = {'video_name': video_file.name, 'video_path': str(video_file), 'boxes': [],  'frame_resolution': 'error', 'error': str(e)}
    
    return results

def print_results_summary(results_dict):

    print("处理结果摘要:")
    print("=" * 60)
    
    total_videos = len(results_dict)
    
    print(f"总视频数: {total_videos}")
    print("-" * 60)
    
    # 详细输出每个视频的boxes
    print("\n详细boxes信息:")
    print("=" * 60)
    for video_name, result in results_dict.items():
        boxes = result.get('boxes', [])
        if boxes:
            print(f"\n{video_name}:")
            for i, box in enumerate(boxes, 1):
                print(f"  ROI{i - 1}: {box}")
        else:
            print(f"\n{video_name}: 无ROI")
    
    print("=" * 60)

def save_results_to_json(results_dict, output_path="video_rois_results.json"):

    # 简化结果以便保存（移除可能包含非JSON序列化的对象）
    simplified_results = {}
    for video_name, result in results_dict.items():
        simplified_results[video_name] = {'video_name': result.get('video_name', video_name), 'boxes': result.get('boxes', []),
           'frame_resolution': result.get('frame_resolution', 'unknown')}
        if 'error' in result:
            simplified_results[video_name]['error'] = result['error']
    
    with open(output_path, 'w', encoding='utf-8') as f:
        json.dump(simplified_results, f, ensure_ascii=False, indent=4)
    
    print(f"\n结果已保存到: {output_path}")


if __name__ == '__main__':
    directory_path = r"B:\data\tiaosheng\20260112"
    results = roi_video_directory(directory_path)
    
    print_results_summary(results)
    
    save_results_to_json(results, "video_rois_results.json")
    
    # 也可以直接打印原始字典
    print("\n原始结果字典:")
    print(json.dumps(results, ensure_ascii=False, indent=2))