R360-UAVmodelTool/report_generator.py

import csv
import xlsxwriter
from openpyxl import load_workbook
from datetime import datetime
from typing import List, Dict

class ReportGenerator:
    def __init__(self, detector):
        self.metadata = {
            'model_name': detector.model_path.split('/')[-1],
            'model_version': detector.model_version,
            'input_size': detector.input_size,
            'test_time': datetime.now().strftime('%Y-%m-%d %H:%M:%S'),
            'total_images': detector.image_count,
            'confidence_threshold': detector.threshold
        }
        self.detections = detector.detection_records

    def generate_csv(self, output_path: str):
        with open(output_path, 'w', newline='', encoding='utf-8') as f:
            writer = csv.writer(f)
            writer.writerow(['Image File', 'Object Count', 'Max Confidence', 'BBox Center X', 'BBox Center Y', 'BBox Width', 'BBox Height', 'Normalized Coordinates'])
            for record in self.detections:
                if record['has_detection']:
                    for detection in record['detections']:
                        x1, y1, x2, y2 = detection['bbox']
                        orig_w = detection.get('orig_w', 0)
                        orig_h = detection.get('orig_h', 0)
                        center_x = round((x1 + x2) / 2, 2)
                        center_y = round((y1 + y2) / 2, 2)
                        width = round(x2 - x1, 2)
                        height = round(y2 - y1, 2)
                        norm_cx = round(center_x / orig_w, 4) if orig_w > 0 else 0
                        norm_cy = round(center_y / orig_h, 4) if orig_h > 0 else 0
                        writer.writerow([
                            record['image_path'],
                            len(record['detections']),
                            detection['confidence'],
                            center_x,
                            center_y,
                            width,
                            height,
                            f"{norm_cx},{norm_cy},{width/orig_w},{height/orig_h}"
                        ])
                else:
                    writer.writerow([record['image_path'], 0, 0])

    def calculate_metrics(self, ground_truth: Dict[str, bool]) -> Dict[str, float]:
        """
        计算模型评估指标
        :param ground_truth: 字典形式的地面真实值 {图像路径: 是否存在目标}
        :return: 包含各项指标的字典
        """
        tp = fp = tn = fn = 0

        for record in self.detections:
            actual = ground_truth.get(record['image_path'], False)
            predicted = record['has_detection']

            if actual and predicted:
                tp += 1
            elif actual and not predicted:
                fn += 1
            elif not actual and predicted:
                fp += 1
            else:
                tn += 1

        # 防止除零错误
        fpr = fp / (fp + tn) if (fp + tn) > 0 else 0
        fnr = fn / (fn + tp) if (fn + tp) > 0 else 0

        return {
            'true_positive': tp,
            'false_positive': fp,
            'true_negative': tn,
            'false_negative': fn,
            'fpr': round(fpr, 4),
            'fnr': round(fnr, 4)
        }

    def generate_excel(self, output_path: str, calculate_metrics: bool = False, ground_truth_file: str = None):
        workbook = xlsxwriter.Workbook(output_path)
        meta_ws = workbook.add_worksheet('Model Info')
        detail_ws = workbook.add_worksheet('Detection Results')

        # 写入元数据
        meta_headers = ['Model Name', 'Version', 'Input Size', 'Test Time', 'Image Count', 'Confidence Threshold']
        meta_ws.write_row(0, 0, meta_headers)
        meta_ws.write_row(1, 0, [
            self.metadata['model_name'],
            self.metadata['model_version'],
            self.metadata['input_size'],
            self.metadata['test_time'],
            self.metadata['total_images'],
            self.metadata['confidence_threshold']
        ])

        # 写入检测明细
        detail_headers = ['Image File', 'Object Count', 'Max Confidence', 'Avg Confidence', 'BBox Center X', 'BBox Center Y', 'BBox Width', 'BBox Height', 'Normalized Coordinates', 'Review Result', 'Remarks']
        detail_ws.write_row(0, 0, detail_headers)

        for row_idx, record in enumerate(self.detections, start=1):
            if not record['has_detection']:
                detail_ws.write_row(row_idx, 0, [
                    record['image_path'],
                    0, 0, 0, '', '', '', '', '', '', ''
                ])
                continue

            # 每个检测框单独生成一行数据
            for detection in record['detections']:
                x1, y1, x2, y2 = detection['bbox']
                orig_w = detection.get('orig_w', 0)
                orig_h = detection.get('orig_h', 0)

                # 计算绝对坐标
                center_x = round((x1 + x2) / 2, 2)
                center_y = round((y1 + y2) / 2, 2)
                width = round(x2 - x1, 2)
                height = round(y2 - y1, 2)

                # 计算归一化坐标
                norm_cx = round(center_x / orig_w, 4) if orig_w > 0 else 0
                norm_cy = round(center_y / orig_h, 4) if orig_h > 0 else 0
                norm_w = round(width / orig_w, 4) if orig_w > 0 else 0
                norm_h = round(height / orig_h, 4) if orig_h > 0 else 0

                detail_ws.write_row(row_idx, 0, [
                    record['image_path'],
                    len(record['detections']),
                    detection['confidence'],
                    '',  # Avg Confidence占位
                    orig_w,  # 新增原始宽度
                    orig_h,  # 新增原始高度
                    center_x,
                    center_y,
                    width,
                    height,
                    f"{norm_cx},{norm_cy},{norm_w},{norm_h}",
                    '',
                    ''
                ])
                row_idx += 1

        # 仅在需要时计算指标
        if calculate_metrics and ground_truth_file:
            metrics = self.calculate_metrics(
                self.parse_excel_ground_truth(ground_truth_file)
            )
            metrics_ws = workbook.add_worksheet('Model Evaluation')
            metrics_headers = ['True Positive', 'False Positive', 'True Negative', 'False Negative', 'False Positive Rate', 'False Negative Rate']
            metrics_ws.write_row(0, 0, metrics_headers)
            metrics_ws.write_row(1, 0, [
                metrics['true_positive'],
                metrics['false_positive'],
                metrics['true_negative'],
                metrics['false_negative'],
                metrics['fpr'],
                metrics['fnr']
            ])
        
        workbook.close()

    def parse_excel_ground_truth(self, excel_path: str) -> Dict[str, bool]:
        """
        从Excel文件解析地面真实值
        :param excel_path: detection_report_0.3.xlsx文件路径
        :return: 包含图像路径和真实检测结果的字典
        """
        wb = load_workbook(excel_path)
        ws = wb.active
        ground_truth = {}
        
        for row in ws.iter_rows(min_row=2, values_only=True):
            image_path = row[0]
            # 假设第二列是真实标签（例如1表示存在目标，0表示不存在）
            actual_value = bool(row[1]) if len(row) > 1 else False
            ground_truth[image_path] = actual_value
        return ground_truth