R360-UAVmodelTool/detection_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

class DetectionVisualizer:
    """在原始图像上直接绘制目标检测框，并按摄像头ID分组"""
    
    def __init__(self, csv_path: str, output_dir: str = None):
        """初始化检测框可视化器
        
        Args:
            csv_path: 检测报告CSV文件路径
            output_dir: 输出目录，默认为CSV文件所在目录
        """
        self.csv_path = csv_path
        self.data = None
        self.camera_data = {}
        
        # 设置输出目录
        if output_dir is None:
            self.output_dir = os.path.join(os.path.dirname(csv_path), "visualization_results")
        else:
            self.output_dir = output_dir
        
        os.makedirs(self.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) -> str:
        """从文件路径中提取摄像头ID
        
        Args:
            file_path: 图像文件路径
            
        Returns:
            摄像头ID
        """
        # 尝试匹配常见的摄像头ID模式
        # 例如: cam_08_18 或 08_18
        patterns = [
            r'cam_(\d+)_(\d+)',  # 匹配 cam_08_18 格式
            r'(\d+)_(\d+)_cam_(\d+)_(\d+)',  # 匹配 192_168_210_2_cam_08_18 格式
        ]
        
        for pattern in patterns:
            match = re.search(pattern, file_path)
            if match:
                if pattern == patterns[0]:
                    return f"cam_{match.group(1)}_{match.group(2)}"
                elif pattern == patterns[1]:
                    return f"cam_{match.group(3)}_{match.group(4)}"
        
        # 如果无法提取，返回文件名作为ID
        return os.path.basename(file_path).split('.')[0]
    
    def process_data(self) -> None:
        """处理数据，按摄像头ID分组"""
        if self.data is None:
            self.load_data()
        
        # 提取摄像头ID并分组
        self.data['Camera ID'] = self.data['Image File'].apply(self.extract_camera_id)
        
        # 按摄像头ID分组
        for camera_id, group in self.data.groupby('Camera ID'):
            self.camera_data[camera_id] = group
            detection_count = len(group[group['Object Count'] > 0])
            print(f"摄像头 {camera_id}: {len(group)} 个图像, {detection_count} 个检测结果")
    
    def draw_detection_on_image(self, image_path: str, detections: pd.DataFrame, resolution: Tuple[int, int] = None) -> np.ndarray:
        """在单个图像上绘制检测框
        
        Args:
            image_path: 图像文件路径
            detections: 该图像的检测数据
            resolution: 图像分辨率，如果为None则使用图像实际分辨率
            
        Returns:
            绘制了检测框的图像
        """
        try:
            # 读取图像
            image = cv2.imread(image_path)
            if image is None:
                print(f"无法读取图像: {image_path}")
                return None
            
            # 转换为RGB格式（OpenCV默认为BGR）
            image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
            
            # 获取图像实际分辨率
            actual_height, actual_width = image.shape[:2]
            if resolution is None:
                resolution = (actual_width, actual_height)
            
            # 创建一个透明度图层用于叠加检测框
            overlay = image.copy()
            
            # 绘制所有检测框
            detection_count = 0
            for _, row in detections.iterrows():
                if row['Object Count'] <= 0:
                    continue
                    
                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, 255, 255), 3)
                    # 绘制主黄色检测框
                    cv2.rectangle(overlay, (x1, y1), (x2, y2), (0, 255, 255), 3)
                    
                    # 在右下角添加带背景的置信度显示
                    if 'Max Confidence' in row and pd.notna(row['Max Confidence']):
                        confidence = float(row['Max Confidence'])
                        text = f'Conf: {confidence:.2f}'
                        (text_width, text_height), baseline = cv2.getTextSize(text, cv2.FONT_HERSHEY_SIMPLEX, 0.5, 1)
                        text_x = x2 - text_width - 5
                        text_y = y2 - 5
                        # 绘制文本背景
                        cv2.rectangle(overlay, (text_x, text_y - text_height), (text_x + text_width, text_y), (255, 255, 255), -1)
                        # 绘制文本
                        cv2.putText(overlay, text, (text_x, text_y), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 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)
            
            return output_image, detection_count
        
        except Exception as e:
            print(f"处理图像时出错: {e}")
            return None, 0
    
    def visualize_camera_detections(self, camera_id: str) -> None:
        """可视化单个摄像头的所有检测结果
        
        Args:
            camera_id: 摄像头ID
        """
        if not self.camera_data:
            self.process_data()
            
        if camera_id not in self.camera_data:
            print(f"未找到摄像头 {camera_id} 的数据")
            return
        
        # 获取该摄像头的所有数据
        cam_data = self.camera_data[camera_id]
        detection_data = cam_data[cam_data['Object Count'] > 0]
        
        if len(detection_data) == 0:
            print(f"摄像头 {camera_id} 没有检测结果")
            return
        
        print(f"处理摄像头 {camera_id} 的 {len(detection_data)} 个检测结果...")
        
        # 创建一个大图像，用于显示所有检测结果
        # 首先选择一个图像作为背景
        sample_image_path = detection_data.iloc[0]['Image File']
        
        # 检查文件路径是否存在，如果是相对路径则转换为绝对路径
        if not os.path.isabs(sample_image_path):
            # 尝试从CSV文件所在目录解析相对路径
            base_dir = os.path.dirname(self.csv_path)
            sample_image_path = os.path.join(base_dir, sample_image_path)
        
        # 如果文件不存在，尝试修复路径
        if not os.path.exists(sample_image_path):
            # 尝试从CSV文件中提取的路径中获取文件名
            file_name = os.path.basename(sample_image_path)
            # 系统化搜索上级目录
            current_dir = os.path.dirname(self.csv_path)
            max_depth = 3
            for _ in range(max_depth):
                current_dir = os.path.dirname(current_dir)
                for root, _, files in os.walk(current_dir):
                    if file_name in files:
                        sample_image_path = os.path.join(root, file_name)
                        print(f"找到匹配图像: {sample_image_path}")
                        break
                if os.path.exists(sample_image_path):
                    break
        
        # 如果仍然找不到文件，报错并返回
        if not os.path.exists(sample_image_path):
            print(f"无法找到图像文件: {sample_image_path}")
            return
        
        # 读取并验证样本图像
        sample_image = cv2.imread(sample_image_path)
        if sample_image is None:
            print(f"无法读取图像: {sample_image_path}，请检查文件格式（支持JPG/PNG）")
            return
        
        # 转换颜色空间并创建副本
        sample_image = cv2.cvtColor(sample_image, cv2.COLOR_BGR2RGB)
        print(f"成功加载底图图像，分辨率: {sample_image.shape[1]}x{sample_image.shape[0]}")
        
        height, width = sample_image.shape[:2]
        
        # 创建一个合成图像，将所有检测框绘制在同一张图上
        # 使用样本图像作为背景
        composite_image = sample_image.copy()
        
        # 处理每个检测结果
        total_detections = 0
        for _, row in detection_data.iterrows():
            image_path = row['Image File']
            
            # 检查文件路径是否存在，如果是相对路径则转换为绝对路径
            if not os.path.isabs(image_path):
                # 尝试从CSV文件所在目录解析相对路径
                base_dir = os.path.dirname(self.csv_path)
                image_path = os.path.join(base_dir, image_path)
            
            # 如果文件不存在，尝试修复路径
            if not os.path.exists(image_path):
                # 尝试从CSV文件中提取的路径中获取文件名
                file_name = os.path.basename(image_path)
                # 在CSV文件所在目录的上级目录中查找
                parent_dir = os.path.dirname(os.path.dirname(self.csv_path))
                for root, _, files in os.walk(parent_dir):
                    if file_name in files:
                        image_path = os.path.join(root, file_name)
                        break
            
            # 如果仍然找不到文件，跳过此检测结果
            if not os.path.exists(image_path):
                print(f"无法找到图像文件: {image_path}，跳过此检测结果")
                continue
            
            # 绘制检测框
            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_px = row['BBox Width']
                        height_px = row['BBox Height']
                        norm_x = center_x / width
                        norm_y = center_y / height
                        norm_w = width_px / width
                        norm_h = height_px / height
                    else:
                        continue
                
                # 将归一化坐标转换为像素坐标
                center_x = int(norm_x * width)
                center_y = int(norm_y * height)
                width_px = int(norm_w * width)
                height_px = int(norm_h * height)
                
                # 计算左上角和右下角坐标
                x1 = max(0, int(center_x - width_px / 2))
                y1 = max(0, int(center_y - height_px / 2))
                x2 = min(width - 1, int(center_x + width_px / 2))
                y2 = min(height - 1, int(center_y + height_px / 2))
                
                # 在合成图像上绘制检测框
                # 使用不同的颜色表示不同的置信度
                confidence = float(row['Max Confidence']) if pd.notna(row.get('Max Confidence')) else 0.5
                # 颜色从蓝色(低置信度)到红色(高置信度)
                color = (
                    int(255 * (1 - confidence)),  # B
                    0,                           # G
                    int(255 * confidence)        # R
                )
                
                cv2.rectangle(composite_image, (x1, y1), (x2, y2), color, 1)
                
                # 在检测框中心绘制一个小点
                cv2.circle(composite_image, (center_x, center_y), 2, color, -1)
                
                total_detections += 1
            except Exception as e:
                print(f"处理检测框时出错: {e}")
        
        # 保存合成图像
        output_path = os.path.join(self.output_dir, f"{camera_id}_all_detections.png")
        cv2.imwrite(output_path, composite_image)
        
        print(f"已生成摄像头 {camera_id} 的检测框合成图像，包含 {total_detections} 个检测框: {output_path}")
        
        # 使用matplotlib创建热力图
        plt.figure(figsize=(width/100, height/100), dpi=100)
        
        # 提取所有检测框的中心点
        centers_x = []
        centers_y = []
        
        for _, row in detection_data.iterrows():
            try:
                if 'Normalized Coordinates' in row and pd.notna(row['Normalized Coordinates']):
                    norm_coords = row['Normalized Coordinates'].split(',')
                    if len(norm_coords) >= 4:
                        norm_x, norm_y = map(float, norm_coords[:2])
                        centers_x.append(norm_x * width)
                        centers_y.append(norm_y * height)
                elif (pd.notna(row.get('BBox Center X')) and pd.notna(row.get('BBox Center Y'))):
                    centers_x.append(row['BBox Center X'])
                    centers_y.append(row['BBox Center Y'])
            except Exception:
                continue
        
        if centers_x and centers_y:
            # 创建热力图
            plt.hexbin(
                x=centers_x,
                y=centers_y,
                gridsize=50,
                cmap='plasma',
                alpha=0.7,
                mincnt=1
            )
            plt.colorbar(label='检测数量')
            
            # 反转Y轴，使坐标系与图像一致
            plt.gca().invert_yaxis()
            plt.title(f"摄像头 {camera_id} 检测热力图 ({total_detections} 个检测)")
            plt.axis('off')
            
            # 保存热力图
            heatmap_path = os.path.join(self.output_dir, f"{camera_id}_heatmap.png")
            plt.savefig(heatmap_path, bbox_inches='tight', pad_inches=0)
            plt.close()
            
            print(f"已生成摄像头 {camera_id} 的检测热力图: {heatmap_path}")
    
    def visualize_all_cameras(self) -> 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.visualize_camera_detections(camera_id)

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

if __name__ == "__main__":
    main()