R360-UAVmodelTool/src/utils/bbox_visualizer.py

import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
import cv2
import os
import re
import argparse
from typing import Dict, List, Tuple, Optional
from matplotlib.colors import LinearSegmentedColormap

class BBoxVisualizer:
    """在原始图像上直接绘制目标检测框"""
    
    def __init__(self, csv_path: str, image_dir: str, output_dir: str = None):
        """初始化检测框可视化器
        
        Args:
            csv_path: 检测报告CSV文件路径
            image_dir: 原始图像目录
            output_dir: 输出目录，默认为CSV文件所在目录
        """
        self.csv_path = csv_path
        self.image_dir = image_dir
        self.data = None
        self.camera_data = {}
        
        # 设置输出目录
        if output_dir is None:
            self.output_dir = os.path.dirname(csv_path)
        else:
            self.output_dir = output_dir
            os.makedirs(output_dir, exist_ok=True)
    
    def load_data(self) -> None:
        """加载CSV数据"""
        try:
            self.data = pd.read_csv(self.csv_path)
            print(f"成功加载数据，共 {len(self.data)} 条记录")
        except Exception as e:
            print(f"加载CSV文件失败: {e}")
            raise
    
    def extract_camera_id(self, file_path: str) -> Optional[str]:
        """从文件路径中提取摄像头ID
        
        Args:
            file_path: 图像文件路径
            
        Returns:
            摄像头ID，如果无法提取则返回None
        """
        # 尝试匹配常见的摄像头ID模式
        # 例如: cam_08_18 或 08_18
        patterns = [
            r'cam_(\d+)_(\d+)',  # 匹配 cam_08_18 格式
            r'(\d+)_(\d+)_(\d+)',  # 匹配 192_168_210_2_cam_08_18 格式中的 08_18
            r'(\d+)_(\d+)_\d+\.jpg$'  # 匹配 01_07_00000.jpg 格式中的 01_07
        ]
        
        for pattern in patterns:
            match = re.search(pattern, file_path)
            if match:
                if len(match.groups()) >= 2:
                    return f"{match.group(1)}_{match.group(2)}"
                else:
                    return match.group(1)
        
        # 如果无法提取，返回文件名作为ID
        return os.path.basename(file_path).split('.')[0]
    
    def process_data(self) -> None:
        """处理数据，按摄像头ID分组"""
        if self.data is None:
            self.load_data()
        
        # 只处理有检测结果的数据
        detection_data = self.data[self.data['Object Count'] > 0].copy()
        
        if len(detection_data) == 0:
            print("警告: 没有找到任何检测结果")
            return
        
        # 提取摄像头ID并分组
        detection_data['Camera ID'] = detection_data['Image File'].apply(self.extract_camera_id)
        
        # 按摄像头ID分组
        for camera_id, group in detection_data.groupby('Camera ID'):
            self.camera_data[camera_id] = group
            print(f"摄像头 {camera_id}: {len(group)} 个检测结果")
    
    def find_image_file(self, camera_id: str) -> Optional[str]:
        """查找指定摄像头ID的图像文件
        
        Args:
            camera_id: 摄像头ID
            
        Returns:
            图像文件路径，如果未找到则返回None
        """
        # 查找匹配的图像文件
        pattern = f"{camera_id}_\d+\.jpg"
        for root, _, files in os.walk(self.image_dir):
            for file in files:
                if re.match(pattern, file):
                    return os.path.join(root, file)
        
        # 如果没有找到精确匹配，尝试模糊匹配
        pattern = f"{camera_id.split('_')[0]}.*{camera_id.split('_')[1]}.*\.jpg"
        for root, _, files in os.walk(self.image_dir):
            for file in files:
                if re.search(pattern, file):
                    return os.path.join(root, file)
        
        return None
    
    def draw_bboxes(self, camera_id: str, resolution: Tuple[int, int] = (3840, 2160)) -> None:
        """在原始图像上绘制检测框
        
        Args:
            camera_id: 摄像头ID
            resolution: 图像分辨率，默认为3840x2160
        """
        if not self.camera_data:
            self.process_data()
            
        if not self.camera_data:
            print("没有可用的检测数据来绘制检测框")
            return
        
        if camera_id not in self.camera_data:
            print(f"未找到摄像头 {camera_id} 的数据")
            return
        
        # 获取该摄像头的检测数据
        cam_data = self.camera_data[camera_id]
        
        # 查找原始图像
        image_path = self.find_image_file(camera_id)
        if image_path is None:
            print(f"未找到摄像头 {camera_id} 的原始图像")
            return
        
        # 读取原始图像
        try:
            image = cv2.imread(image_path)
            if image is None:
                raise ValueError(f"无法读取图像: {image_path}")
            
            # 转换为RGB格式（OpenCV默认为BGR）
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            
            # 获取图像实际分辨率
            actual_height, actual_width = image.shape[:2]
            resolution = (actual_width, actual_height)
            
            print(f"已加载图像: {image_path}, 分辨率: {resolution}")
        except Exception as e:
            print(f"读取图像时出错: {e}")
            return
        
        # 创建一个透明度图层用于叠加检测框
        overlay = image.copy()
        
        # 绘制所有检测框
        detection_count = 0
        for _, row in cam_data.iterrows():
            try:
                # 获取归一化坐标和宽高
                if 'Normalized Coordinates' in row and pd.notna(row['Normalized Coordinates']):
                    try:
                        norm_coords = row['Normalized Coordinates'].split(',')
                        if len(norm_coords) >= 4:  # 中心点x,y,宽,高
                            norm_x, norm_y, norm_w, norm_h = map(float, norm_coords[:4])
                        else:
                            continue
                    except (ValueError, IndexError):
                        continue
                else:
                    # 如果没有归一化坐标，使用中心点坐标和宽高计算
                    if (pd.notna(row.get('BBox Center X')) and pd.notna(row.get('BBox Center Y')) and
                        pd.notna(row.get('BBox Width')) and pd.notna(row.get('BBox Height'))):
                        center_x = row['BBox Center X']
                        center_y = row['BBox Center Y']
                        width = row['BBox Width']
                        height = row['BBox Height']
                        norm_x = center_x / resolution[0]
                        norm_y = center_y / resolution[1]
                        norm_w = width / resolution[0]
                        norm_h = height / resolution[1]
                    else:
                        continue
                
                # 将归一化坐标转换为像素坐标
                center_x = int(norm_x * resolution[0])
                center_y = int(norm_y * resolution[1])
                width = int(norm_w * resolution[0])
                height = int(norm_h * resolution[1])
                
                # 计算左上角和右下角坐标
                x1 = max(0, int(center_x - width / 2))
                y1 = max(0, int(center_y - height / 2))
                x2 = min(resolution[0] - 1, int(center_x + width / 2))
                y2 = min(resolution[1] - 1, int(center_y + height / 2))
                
                # 绘制检测框，使用半透明的红色
                cv2.rectangle(overlay, (x1, y1), (x2, y2), (255, 0, 0), 2)
                
                # 如果有置信度信息，显示在框上方
                if 'Confidence' in row and pd.notna(row['Confidence']):
                    confidence = float(row['Confidence'])
                    cv2.putText(overlay, f"{confidence:.2f}", (x1, y1 - 5),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 0, 0), 1)
                
                detection_count += 1
            except Exception as e:
                print(f"处理检测框时出错: {e}")
        
        # 设置透明度
        alpha = 0.6
        output_image = cv2.addWeighted(overlay, alpha, image, 1 - alpha, 0)
        
        # 保存结果图像
        output_path = os.path.join(self.output_dir, f"bbox_overlay_{camera_id}.jpg")
        output_image_rgb = cv2.cvtColor(output_image, cv2.COLOR_RGB2BGR)  # 转回BGR格式保存
        cv2.imwrite(output_path, output_image_rgb)
        
        print(f"已在图像上绘制 {detection_count} 个检测框: {output_path}")
        
        # 使用matplotlib显示结果
        plt.figure(figsize=(12, 8))
        plt.imshow(output_image)
        plt.title(f"摄像头 {camera_id} 目标检测框叠加图 ({detection_count} 个检测)")
        plt.axis('off')
        
        # 保存matplotlib版本的图像
        plt_output_path = os.path.join(self.output_dir, f"bbox_overlay_{camera_id}_plt.png")
        plt.savefig(plt_output_path, dpi=300, bbox_inches='tight')
        plt.close()
    
    def draw_all_bboxes(self, resolution: Tuple[int, int] = (3840, 2160)) -> None:
        """为所有摄像头绘制检测框"""
        if not self.camera_data:
            self.process_data()
            
        if not self.camera_data:
            print("没有可用的检测数据来绘制检测框")
            return
        
        for camera_id in self.camera_data.keys():
            self.draw_bboxes(camera_id, resolution)

def main():
    parser = argparse.ArgumentParser(description='目标检测框可视化工具')
    parser.add_argument('--csv', type=str, required=True, help='检测报告CSV文件路径')
    parser.add_argument('--image_dir', type=str, required=True, help='原始图像目录')
    parser.add_argument('--output', type=str, default=None, help='输出目录路径')
    parser.add_argument('--camera', type=str, default=None, help='指定摄像头ID进行分析，不指定则分析所有摄像头')
    
    args = parser.parse_args()
    
    # 创建可视化器并处理数据
    visualizer = BBoxVisualizer(args.csv, args.image_dir, args.output)
    visualizer.load_data()
    visualizer.process_data()
    
    # 绘制检测框
    if args.camera:
        visualizer.draw_bboxes(args.camera)
    else:
        visualizer.draw_all_bboxes()

if __name__ == "__main__":
    main()