detection.py

import torch
import torchvision
from torchvision.models.detection import FasterRCNN
from torchvision.models.detection.rpn import AnchorGenerator
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.datasets import CocoDetection, VOCDetection
import torchvision.transforms.v2 as T # 使用新的 v2 transforms
from torch.utils.data import DataLoader, Subset, Dataset
import os
import json
from PIL import Image
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import utils # PyTorch 官方 Faster R-CNN 示例中的辅助文件 (需要下载或实现)
             # https://github.com/pytorch/vision/tree/main/references/detection
import argparse # 用于命令行参数解析
import glob # 用于查找文件
import xml.etree.ElementTree as ET # 用于解析 PASCAL VOC XML
import logging
import numpy as np
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval

# --- 命令行参数解析 ---
parser = argparse.ArgumentParser(description='目标检测训练脚本')
parser.add_argument('--dataset_type', type=str, default='coco', choices=['coco', 'pascal', 'yolo'],
                    help='数据集类型：coco, pascal, 或 yolo')
args = parser.parse_args()
DATASET_TYPE = args.dataset_type
print(f"选择的数据集类型: {DATASET_TYPE}")

# --- 1. 配置和超参数 ---

# --- 数据集路径定义 ---
COCO_BASE_PATH = '/hdd_16T/Zirui/work/CNNa2/dataset/dataset_42028assg2_24902417/Object_Detection/coco'
COCO_TRAIN_IMG_DIR = os.path.join(COCO_BASE_PATH, 'train')
COCO_TRAIN_ANN_FILE = os.path.join(COCO_BASE_PATH, 'train/train_annotations.json')
COCO_VAL_IMG_DIR = os.path.join(COCO_BASE_PATH, 'valid')
COCO_VAL_ANN_FILE = os.path.join(COCO_BASE_PATH, 'valid/valid_annotations.json')
COCO_TEST_IMG_DIR = os.path.join(COCO_BASE_PATH, 'test')
COCO_TEST_ANN_FILE = os.path.join(COCO_BASE_PATH, 'test/test_annotations.json')

PASCAL_BASE_PATH = '/hdd_16T/Zirui/work/CNNa2/dataset/dataset_42028assg2_24902417/Object_Detection/pascal'
PASCAL_YEAR = '2007' # 假设使用 VOC 2007

YOLO_BASE_PATH = '/hdd_16T/Zirui/work/CNNa2/dataset/dataset_42028assg2_24902417/Object_Detection/yolo'


# --- COCO 类别定义 ---
def load_coco_categories(ann_file):
    with open(ann_file, 'r') as f:
        data = json.load(f)
    categories = {cat['id']: cat['name'] for cat in data['categories']}
    return categories, len(categories)

# --- PASCAL VOC 类别定义 ---
# 更新为实际数据集中的类别
PASCAL_CLASSES = [
    "__background__", # 背景类
    "young", "empty_pod",  # 实际数据集中的类别
]
# 创建类别名称到 ID 的映射 (从 1 开始，符合通常习惯)
PASCAL_NAME_TO_ID = {name: i for i, name in enumerate(PASCAL_CLASSES) if name != "__background__"}
PASCAL_ID_TO_NAME = {i: name for name, i in PASCAL_NAME_TO_ID.items()}
PASCAL_NUM_CLASSES = len(PASCAL_CLASSES) - 1 # 实际物体类别数量

# --- YOLO 类别定义 ---
def load_yolo_classes(classes_file=None):
    # 如果有提供classes.txt文件，从文件加载
    if classes_file and os.path.exists(classes_file):
        with open(classes_file, 'r') as f:
            classes = [line.strip() for line in f.readlines()]
        return {i+1: name for i, name in enumerate(classes)}, len(classes)
    
    # 否则使用yaml文件中定义的类别，或者使用默认值
    # 从data.yaml加载
    yaml_file = os.path.join(YOLO_BASE_PATH, 'data.yaml')
    if os.path.exists(yaml_file):
        try:
            import yaml
            with open(yaml_file, 'r') as f:
                data = yaml.safe_load(f)
                if 'names' in data:
                    classes = data['names']
                    return {i+1: name for i, name in enumerate(classes)}, len(classes)
        except Exception as e:
            print(f"从YAML加载类别失败: {e}")
    
    # 如果都失败了，使用默认值
    classes = ['Ready', 'empty_pod', 'germination', 'pod', 'young']
    return {i+1: name for i, name in enumerate(classes)}, len(classes)

# 根据数据集类型加载相应的类别信息
if DATASET_TYPE == 'coco':
    category_id_to_name, NUM_CLASSES = load_coco_categories(COCO_TRAIN_ANN_FILE)
elif DATASET_TYPE == 'pascal':
    category_id_to_name = PASCAL_ID_TO_NAME
    NUM_CLASSES = PASCAL_NUM_CLASSES
elif DATASET_TYPE == 'yolo':
    category_id_to_name, NUM_CLASSES = load_yolo_classes() # 不传入文件，使用默认类别或从yaml加载
else:
    raise ValueError(f"不支持的数据集类型: {DATASET_TYPE}")

print(f"类别数量: {NUM_CLASSES}")

# --- 训练参数 ---
DEVICE = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
print(f"使用的设备: {DEVICE}")
BATCH_SIZE = 32 # 根据你的 GPU 显存调整
NUM_EPOCHS = 10 # 训练轮数
LEARNING_RATE = 0.005
MOMENTUM = 0.9
WEIGHT_DECAY = 0.0005
SAVE_MODEL_PATH = f'faster_rcnn_{DATASET_TYPE}_model.pth'
# 背景类也要算进去，所以模型输出是 NUM_CLASSES + 1
MODEL_NUM_CLASSES = NUM_CLASSES + 1

# --- 2. 数据集和数据加载器 ---

# 更新 get_transform 使其更通用
def get_transform(train):
    transforms = []
    if train:
        # 训练时的数据增强
        transforms.append(T.RandomHorizontalFlip(0.5))
        # 可以添加更多增强...
    
    # 对于 COCO 和 VOC，输入已经是 PIL Image
    # 转换为 tensor 并调整数据类型
    transforms.append(T.PILToTensor())
    transforms.append(T.ConvertImageDtype(torch.float32)) # 归一化到 [0, 1]
    
    # 确保图像格式正确，例如确保3通道
    # transforms.append(lambda x: x if x.shape[0] == 3 else torch.cat([x, x, x], 0) if x.shape[0] == 1 else x[:3])
    
    return T.Compose(transforms)

# 自定义COCO数据集类 (保持不变)
class CustomCocoDetection(CocoDetection):
    def __getitem__(self, index):
        # 注意：确保 CustomCocoDetection 中的 transforms 调用方式正确
        # 如果 get_transform 只接受 img，则调用应为 img = self.transforms(img)
        # 如果 get_transform 设计为 T.Compose([T.PILToTensor(), ...])，则原始调用可能正确
        # 为保持一致性，让 get_transform 只处理图像，并在 Dataset 类中应用
        
        img, coco_targets = super(CustomCocoDetection, self).__getitem__(index) # 获取原始 PIL 图像和 COCO 标注列表

        # 将COCO标注格式转换为Faster R-CNN需要的格式
        image_id = self.ids[index]
        target = {}
        target["boxes"] = []
        target["labels"] = []
        target["image_id"] = torch.tensor([image_id])
        target["area"] = [] # Faster R-CNN 可能需要 area
        target["iscrowd"] = [] # Faster R-CNN 可能需要 iscrowd

        for anno in coco_targets:
            x, y, w, h = anno["bbox"]
            xmin = x
            ymin = y
            xmax = x + w
            ymax = y + h
            if w > 0 and h > 0:
                target["boxes"].append([xmin, ymin, xmax, ymax])
                target["labels"].append(anno["category_id"])
                target["area"].append(anno["area"])
                target["iscrowd"].append(anno["iscrowd"])

        if len(target["boxes"]) > 0:
            target["boxes"] = torch.as_tensor(target["boxes"], dtype=torch.float32)
            target["labels"] = torch.as_tensor(target["labels"], dtype=torch.int64)
            target["area"] = torch.as_tensor(target["area"], dtype=torch.float32)
            target["iscrowd"] = torch.as_tensor(target["iscrowd"], dtype=torch.uint8)
        else:
            # 处理没有标注的情况
            target["boxes"] = torch.zeros((0, 4), dtype=torch.float32)
            target["labels"] = torch.zeros((0,), dtype=torch.int64)
            target["area"] = torch.zeros((0,), dtype=torch.float32)
            target["iscrowd"] = torch.zeros((0,), dtype=torch.uint8)

        # 应用图像变换
        if self.transforms is not None:
            img = self.transforms(img) # 假设 transform 只作用于图像

        return img, target

# 自定义 PASCAL VOC 数据集类
class CustomVOCDetection(Dataset):
    def __init__(self, root, year='2007', image_set='train', transforms=None):
        self.root = root
        self.year = year
        self.image_set = image_set
        self._transforms = transforms
        
        # 针对用户提供的目录结构：图像和XML注释在同一目录下
        self.data_dir = os.path.join(root, image_set)
        
        if not os.path.exists(self.data_dir):
            raise FileNotFoundError(f"can't find Pascal data set directory: {self.data_dir}")
        
        # 加载所有图像文件名（不含扩展名）
        self.images = []
        for f in os.listdir(self.data_dir):
            if f.endswith('.jpg') or f.endswith('.png'):
                img_id = os.path.splitext(f)[0]
                xml_file = os.path.join(self.data_dir, f"{img_id}.xml")
                # 确保每个图像都有对应的XML标注
                if os.path.exists(xml_file):
                    self.images.append(img_id)
        
        # 类别映射
        self.name_to_id = PASCAL_NAME_TO_ID
        
        print(f"found {len(self.images)} Pascal VOC {image_set} images")
    
    def __len__(self):
        return len(self.images)
    
    def __getitem__(self, index):
        img_id = self.images[index]
        
        # 加载图像
        img_path = os.path.join(self.data_dir, f"{img_id}.jpg")
        if not os.path.exists(img_path):
            img_path = os.path.join(self.data_dir, f"{img_id}.png") # 尝试PNG格式
        
        img = Image.open(img_path).convert("RGB")
        img_width, img_height = img.size
        
        # 加载标注XML
        xml_path = os.path.join(self.data_dir, f"{img_id}.xml")
        tree = ET.parse(xml_path)
        root = tree.getroot()
        
        target = {}
        target["boxes"] = []
        target["labels"] = []
        target["image_id"] = torch.tensor([index])
        target["area"] = []
        target["iscrowd"] = []
        
        for obj in root.findall('object'):
            label_name = obj.find('name').text
            if label_name not in self.name_to_id:
                continue
                
            label_id = self.name_to_id[label_name]
            
            bbox = obj.find('bndbox')
            xmin = float(bbox.find('xmin').text)
            ymin = float(bbox.find('ymin').text)
            xmax = float(bbox.find('xmax').text)
            ymax = float(bbox.find('ymax').text)
            
            # 检查边界框有效性
            if xmax > xmin and ymax > ymin:
                target["boxes"].append([xmin, ymin, xmax, ymax])
                target["labels"].append(label_id)
                area = (xmax - xmin) * (ymax - ymin)
                target["area"].append(area)
                target["iscrowd"].append(0) # VOC 没有 crowd 标注
        
        if len(target["boxes"]) > 0:
            target["boxes"] = torch.as_tensor(target["boxes"], dtype=torch.float32)
            target["labels"] = torch.as_tensor(target["labels"], dtype=torch.int64)
            target["area"] = torch.as_tensor(target["area"], dtype=torch.float32)
            target["iscrowd"] = torch.as_tensor(target["iscrowd"], dtype=torch.uint8)
        else:
            target["boxes"] = torch.zeros((0, 4), dtype=torch.float32)
            target["labels"] = torch.zeros((0,), dtype=torch.int64)
            target["area"] = torch.zeros((0,), dtype=torch.float32)
            target["iscrowd"] = torch.zeros((0,), dtype=torch.uint8)
        
        # 应用图像变换
        if self._transforms is not None:
            img = self._transforms(img)
        
        return img, target

# 自定义YOLO数据集类
class CustomYOLODataset(Dataset):
    def __init__(self, img_dir, label_dir, classes_file=None, transforms=None):
        self.img_dir = img_dir
        self.label_dir = label_dir
        self.transforms = transforms
        
        # 获取所有图像文件
        self.img_files = sorted(glob.glob(os.path.join(img_dir, '*.jpg')) + 
                               glob.glob(os.path.join(img_dir, '*.png')))
        
        # 确保每个图像都有对应的标签文件
        self.valid_files = []
        for img_path in self.img_files:
            base_name = os.path.basename(img_path).split('.')[0]
            label_path = os.path.join(label_dir, f"{base_name}.txt")
            if os.path.exists(label_path):
                self.valid_files.append((img_path, label_path))
        
        print(f"found {len(self.valid_files)} YOLO format image-label pairs")
    
    def __len__(self):
        return len(self.valid_files)
    
    def __getitem__(self, idx):
        img_path, label_path = self.valid_files[idx]
        
        # 加载图像
        img = Image.open(img_path).convert("RGB")
        img_width, img_height = img.size
        
        target = {}
        target["boxes"] = []
        target["labels"] = []
        target["image_id"] = torch.tensor([idx])
        target["area"] = []
        target["iscrowd"] = []
        
        # 解析YOLO格式标签 (类别 中心x 中心y 宽 高)
        # YOLO坐标已经被归一化为[0,1]
        with open(label_path, 'r') as f:
            for line in f.readlines():
                if line.strip():
                    parts = line.strip().split()
                    if len(parts) == 5:
                        cls_id = int(parts[0]) + 1  # YOLO类别从0开始，我们从1开始
                        x_center = float(parts[1]) * img_width
                        y_center = float(parts[2]) * img_height
                        width = float(parts[3]) * img_width
                        height = float(parts[4]) * img_height
                        
                        # 将中心坐标和宽高转换为左上角和右下角坐标
                        xmin = max(0, x_center - width/2)
                        ymin = max(0, y_center - height/2)
                        xmax = min(img_width, x_center + width/2)
                        ymax = min(img_height, y_center + height/2)
                        
                        # 确保边界框有效
                        if xmax > xmin and ymax > ymin:
                            target["boxes"].append([xmin, ymin, xmax, ymax])
                            target["labels"].append(cls_id)
                            area = (xmax - xmin) * (ymax - ymin)
                            target["area"].append(area)
                            target["iscrowd"].append(0)
        
        if len(target["boxes"]) > 0:
            target["boxes"] = torch.as_tensor(target["boxes"], dtype=torch.float32)
            target["labels"] = torch.as_tensor(target["labels"], dtype=torch.int64)
            target["area"] = torch.as_tensor(target["area"], dtype=torch.float32)
            target["iscrowd"] = torch.as_tensor(target["iscrowd"], dtype=torch.uint8)
        else:
            target["boxes"] = torch.zeros((0, 4), dtype=torch.float32)
            target["labels"] = torch.zeros((0,), dtype=torch.int64)
            target["area"] = torch.zeros((0,), dtype=torch.float32)
            target["iscrowd"] = torch.zeros((0,), dtype=torch.uint8)
        
        # 应用图像变换
        if self.transforms:
            img = self.transforms(img)
        
        return img, target

# 使用自定义的CocoDetection数据集类
try:
    # 根据数据集类型加载不同的数据集
    if DATASET_TYPE == 'coco':
        dataset_train = CustomCocoDetection(root=COCO_TRAIN_IMG_DIR,
                                      annFile=COCO_TRAIN_ANN_FILE,
                                      transforms=get_transform(train=True))

        dataset_val = CustomCocoDetection(root=COCO_VAL_IMG_DIR,
                                    annFile=COCO_VAL_ANN_FILE,
                                    transforms=get_transform(train=False))

        dataset_test = CustomCocoDetection(root=COCO_TEST_IMG_DIR,
                                     annFile=COCO_TEST_ANN_FILE,
                                     transforms=get_transform(train=False))
    
    elif DATASET_TYPE == 'pascal':
        dataset_train = CustomVOCDetection(root=PASCAL_BASE_PATH, 
                                          year=PASCAL_YEAR, 
                                          image_set='train', 
                                          transforms=get_transform(train=True))
        
        dataset_val = CustomVOCDetection(root=PASCAL_BASE_PATH, 
                                        year=PASCAL_YEAR, 
                                        image_set='valid', 
                                        transforms=get_transform(train=False))
        
        dataset_test = CustomVOCDetection(root=PASCAL_BASE_PATH, 
                                         year=PASCAL_YEAR, 
                                         image_set='test', 
                                         transforms=get_transform(train=False))
        
        # 调试信息: 打印一些样本的标签信息
        print("\nDebug information - Pascal VOC dataset:")
        print(f"Category mapping: {PASCAL_NAME_TO_ID}")
        for i in range(min(3, len(dataset_train))):
            _, target = dataset_train[i]
            print(f"Sample {i}:")
            print(f" Number of bounding boxes: {len(target['boxes'])}")
            print(f" Labels: {target['labels'].tolist()}")
            if len(target['labels']) > 0:
                print(f" Category names: {[PASCAL_ID_TO_NAME[label.item()] for label in target['labels']]}")
            print(f" Bounding boxes: {target['boxes'].shape}")
            
    
    elif DATASET_TYPE == 'yolo':
        # YOLO数据集目录结构
        YOLO_TRAIN_IMG_DIR = os.path.join(YOLO_BASE_PATH, 'train', 'images')
        YOLO_TRAIN_LABEL_DIR = os.path.join(YOLO_BASE_PATH, 'train', 'labels')
        YOLO_VAL_IMG_DIR = os.path.join(YOLO_BASE_PATH, 'valid', 'images')
        YOLO_VAL_LABEL_DIR = os.path.join(YOLO_BASE_PATH, 'valid', 'labels')
        YOLO_TEST_IMG_DIR = os.path.join(YOLO_BASE_PATH, 'test', 'images')
        YOLO_TEST_LABEL_DIR = os.path.join(YOLO_BASE_PATH, 'test', 'labels')
        
        dataset_train = CustomYOLODataset(
            img_dir=YOLO_TRAIN_IMG_DIR,
            label_dir=YOLO_TRAIN_LABEL_DIR,
            classes_file=None,  # 不再使用classes.txt，而是在load_yolo_classes中处理
            transforms=get_transform(train=True)
        )
        
        dataset_val = CustomYOLODataset(
            img_dir=YOLO_VAL_IMG_DIR,
            label_dir=YOLO_VAL_LABEL_DIR,
            classes_file=None,
            transforms=get_transform(train=False)
        )
        
        dataset_test = CustomYOLODataset(
            img_dir=YOLO_TEST_IMG_DIR,
            label_dir=YOLO_TEST_LABEL_DIR,
            classes_file=None,
            transforms=get_transform(train=False)
        )
        
        # 调试信息: 打印一些样本的标签信息
        print("\nDebug information - YOLO dataset:")
        print(f"Category mapping: {category_id_to_name}")
        for i in range(min(3, len(dataset_train))):
            _, target = dataset_train[i]
            print(f"Sample {i}:")
            print(f" Number of bounding boxes: {len(target['boxes'])}")
            print(f" Labels: {target['labels'].tolist()}")
            if len(target['labels']) > 0:
                print(f" Category names: {[category_id_to_name.get(label.item(), f'unknown-{label.item()}') for label in target['labels']]}")
            print(f" Bounding boxes: {target['boxes'].shape}")
    
    else:
        raise ValueError(f"Unsupported dataset type: {DATASET_TYPE}")

    print(f"Training set size: {len(dataset_train)}")
    print(f"Validation set size: {len(dataset_val)}")
    print(f"Test set size: {len(dataset_test)}")

    # # 可选：如果数据集很大，可以只取一部分进行调试
    # dataset_train = Subset(dataset_train, range(100))
    # dataset_val = Subset(dataset_val, range(50))
    # dataset_test = Subset(dataset_test, range(50))

    # 定义数据加载器
    # 需要一个 collate_fn 来处理批次中不同数量的目标
    def collate_fn(batch):
        return tuple(zip(*batch))

    data_loader_train = DataLoader(dataset_train, batch_size=BATCH_SIZE, shuffle=True,
                                   num_workers=4, collate_fn=collate_fn)
    data_loader_val = DataLoader(dataset_val, batch_size=1, shuffle=False, # 验证通常 batch_size=1
                                 num_workers=4, collate_fn=collate_fn)
    data_loader_test = DataLoader(dataset_test, batch_size=1, shuffle=False,
                                  num_workers=4, collate_fn=collate_fn)

except FileNotFoundError as e:
    print(f"Error: can't find training or validation files: {e}")
    print(f"Please ensure the file paths for the selected dataset {DATASET_TYPE} are correct and files exist.")
    exit()
except NotImplementedError as e:
    print(f"Error: {e}")
    print("Please implement the corresponding dataset loading class and try again.")
    exit()
except Exception as e:
    print(f"Error: {e}")
    exit()

# --- 3. 模型定义 (Faster R-CNN) ---

def get_faster_rcnn_model(num_classes):
    # 加载预训练的 Faster R-CNN 模型 (ResNet-50 backbone with FPN)
    weights = torchvision.models.detection.FasterRCNN_ResNet50_FPN_Weights.DEFAULT
    model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights=weights)

    # 获取分类器的输入特征数
    in_features = model.roi_heads.box_predictor.cls_score.in_features

    # 替换预训练的头部为一个新的头部 (适应我们的类别数量)
    model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

    # 可选: 使用 MobileNetV3-Large FPN 作为 backbone (更快但精度可能稍低)
    # weights_mobilenet = torchvision.models.detection.FasterRCNN_MobileNet_V3_Large_FPN_Weights.DEFAULT
    # model = torchvision.models.detection.fasterrcnn_mobilenet_v3_large_fpn(weights=weights_mobilenet)
    # in_features = model.roi_heads.box_predictor.cls_score.in_features
    # model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

    return model

# 实例化模型
model = get_faster_rcnn_model(MODEL_NUM_CLASSES)
model.to(DEVICE)

# 打印模型信息
print(f"\nModel information:")
print(f"Number of input classes: {MODEL_NUM_CLASSES} (including background class)")
print(f"Used device: {DEVICE}")
print(f"Predictor input features: {model.roi_heads.box_predictor.cls_score.in_features}")
print(f"Predictor output classes: {model.roi_heads.box_predictor.cls_score.out_features}")

# --- 4. 训练设置 ---

# 定义优化器
params = [p for p in model.parameters() if p.requires_grad]
optimizer = torch.optim.SGD(params, lr=LEARNING_RATE, momentum=MOMENTUM, weight_decay=WEIGHT_DECAY)

# 定义学习率调度器 (可选, 但推荐)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=3, gamma=0.1)

# --- 5. 训练循环 ---

print("Starting training...")

# 设置日志配置
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    handlers=[
        logging.FileHandler(f'training_{DATASET_TYPE}_Faster_R-CNN.log'),  # 保存到文件
        logging.StreamHandler()  # 同时输出到控制台
    ]
)

logger = logging.getLogger('faster_rcnn')

# 使用示例
logger.info(f"Selected dataset type: {DATASET_TYPE}")
logger.info(f"Training set size: {len(dataset_train)}")

# --- 定义mAP计算函数 ---
def evaluate_map(model, data_loader, device, dataset_type, category_mapping=None):
    model.eval()  # 设置模型为评估模式
    
    # 用于存储所有预测和目标信息
    all_predictions = []
    all_targets = []
    
    print(f"Caluating mAP...")
    
    with torch.no_grad():
        for images, targets in data_loader:
            images = list(image.to(device) for image in images)
            
            # 执行预测
            predictions = model(images)
            
            # 将预测和目标添加到列表中
            all_predictions.extend(predictions)
            all_targets.extend(targets)
    
    # 将预测和目标转换为COCO格式，以便使用pycocotools计算mAP
    pred_instances = []
    gt_instances = []
    
    image_id = 0
    instance_id = 0
    
    for pred, target in zip(all_predictions, all_targets):
        # 处理预测结果
        boxes = pred['boxes'].cpu().numpy()
        scores = pred['scores'].cpu().numpy()
        labels = pred['labels'].cpu().numpy()
        
        # 只保留置信度较高的预测
        keep = scores > 0.05
        boxes = boxes[keep]
        scores = scores[keep]
        labels = labels[keep]
        
        for box, score, label in zip(boxes, scores, labels):
            pred_instances.append({
                'image_id': image_id,
                'category_id': int(label),
                'bbox': [float(box[0]), float(box[1]), float(box[2] - box[0]), float(box[3] - box[1])],  # COCO格式 [x, y, width, height]
                'score': float(score),
                'id': instance_id
            })
            instance_id += 1
        
        # 处理真实标签
        gt_boxes = target['boxes'].cpu().numpy()
        gt_labels = target['labels'].cpu().numpy()
        
        for gt_box, gt_label in zip(gt_boxes, gt_labels):
            gt_instances.append({
                'image_id': image_id,
                'category_id': int(gt_label),
                'bbox': [float(gt_box[0]), float(gt_box[1]), float(gt_box[2] - gt_box[0]), float(gt_box[3] - gt_box[1])],  # COCO格式
                'area': float((gt_box[2] - gt_box[0]) * (gt_box[3] - gt_box[1])),
                'iscrowd': 0,
                'id': instance_id
            })
            instance_id += 1
        
        image_id += 1
    
    # 创建COCO对象
    coco_gt = COCO()
    coco_dt = COCO()
    
    # 设置类别信息
    categories = []
    if dataset_type == 'coco':
        for id, name in category_mapping.items():
            categories.append({'id': id, 'name': name})
    elif dataset_type == 'pascal':
        for id, name in category_mapping.items():
            categories.append({'id': id, 'name': name})
    elif dataset_type == 'yolo':
        for id, name in category_mapping.items():
            categories.append({'id': id, 'name': name})
    
    # 创建COCO数据结构
    coco_gt.dataset = {
        'images': [{'id': i} for i in range(image_id)],
        'annotations': gt_instances,
        'categories': categories
    }
    coco_gt.createIndex()
    
    coco_dt.dataset = {
        'images': [{'id': i} for i in range(image_id)],
        'annotations': pred_instances,
        'categories': categories
    }
    coco_dt.createIndex()
    
    # 计算mAP
    coco_eval = COCOeval(coco_gt, coco_dt, 'bbox')
    coco_eval.evaluate()
    coco_eval.accumulate()
    coco_eval.summarize()
    
    # 返回[email protected]:0.95 (mAP)
    return coco_eval.stats[0]

for epoch in range(NUM_EPOCHS):
    model.train() # 设置模型为训练模式
    epoch_loss = 0
    for i, (images, targets) in enumerate(data_loader_train):
        images = list(image.to(DEVICE) for image in images)
        # 确保 targets 中的 'boxes' 和 'labels' 是 Tensors
        targets = [{k: v.to(DEVICE) if isinstance(v, torch.Tensor) else v for k, v in t.items()} for t in targets]
        

        
        # 清空之前的梯度
        optimizer.zero_grad()

        # 前向传播并计算损失
        # Faster R-CNN 在训练模式下返回一个损失字典
        loss_dict = model(images, targets)
        losses = sum(loss for loss in loss_dict.values()) # 将所有损失加起来
        
        # 调试信息：打印单个损失组件
        if i < 2 or i % 50 == 0:  # 定期打印损失明细
            print(f"  损失详情: {', '.join([f'{k}: {v.item():.6f}' for k, v in loss_dict.items()])}")

        # 反向传播
        losses.backward()

        # 更新权重
        optimizer.step()

        batch_loss = losses.item()
        epoch_loss += batch_loss

        if (i + 1) % 50 == 0: # 每 50 个 batch 打印一次日志
            print(f'Epoch [{epoch+1}/{NUM_EPOCHS}], Step [{i+1}/{len(data_loader_train)}], Batch Loss: {batch_loss:.6f}')

    # 更新学习率
    lr_scheduler.step()

    avg_epoch_loss = epoch_loss / len(data_loader_train)
    print(f'Epoch [{epoch+1}/{NUM_EPOCHS}] 结束, 平均训练损失: {avg_epoch_loss:.4f}')

    # --- 6. 验证循环 ---
    model.eval() # 设置模型为评估模式
    val_loss = 0
    with torch.no_grad(): # 验证时不需要计算梯度
        for images, targets in data_loader_val:
            images = list(image.to(DEVICE) for image in images)
            # 验证时模型也需要 targets 来计算可能的损失或其他指标 (如果需要)
            # 如果只想做推理，targets 可以是 None，但官方实现通常需要它们
            targets = [{k: v.to(DEVICE) if isinstance(v, torch.Tensor) else v for k, v in t.items()} for t in targets]

            # 注意: 在 eval 模式下，模型默认返回预测结果，而不是损失。
            # 为了获取验证损失（如果需要评估指标），可能需要修改模型或评估逻辑
            # 这里我们暂时跳过精确的验证损失计算，因为配置比较复杂
            # 如果你想评估 mAP，需要使用 cocoapi 或类似工具
            pass # 暂时跳过验证损失计算和 mAP 评估
    
    # 计算验证集上的mAP
    print("Calculating mAP...")
    val_map = evaluate_map(model, data_loader_val, DEVICE, DATASET_TYPE, 
                          category_id_to_name if DATASET_TYPE == 'coco' or DATASET_TYPE == 'yolo' else PASCAL_ID_TO_NAME)
    print(f'Epoch [{epoch+1}/{NUM_EPOCHS}], Validation mAP: {val_map:.4f}')
    logger.info(f"Epoch [{epoch+1}/{NUM_EPOCHS}], Validation mAP: {val_map:.4f}")
    # print(f'Epoch [{epoch+1}/{NUM_EPOCHS}], 验证完成.') # 添加验证指标输出

    # 保存模型 (例如，每个 epoch 或基于验证性能)
    torch.save(model.state_dict(), f'faster_rcnn_{DATASET_TYPE}_epoch_{epoch+1}.pth')
    print(f"Model saved to faster_rcnn_{DATASET_TYPE}_epoch_{epoch+1}.pth")

    # 记录指标
    logger.info(f"Epoch [{epoch+1}/{NUM_EPOCHS}], Average training loss: {avg_epoch_loss:.4f}, Validation mAP: {val_map:.4f}")

print("Training completed!")
# 保存最终模型
torch.save(model.state_dict(), SAVE_MODEL_PATH)
print(f"Final model saved to {SAVE_MODEL_PATH}")


# --- 7. 测试/推理 ---

print("Starting test/inference...")
# 加载训练好的模型权重
model.load_state_dict(torch.load(SAVE_MODEL_PATH, map_location=DEVICE))
model.eval() # 确保模型在评估模式

# 计算训练集mAP
print("计算训练集mAP...")
train_map = evaluate_map(model, data_loader_train, DEVICE, DATASET_TYPE, 
                        category_id_to_name if DATASET_TYPE == 'coco' or DATASET_TYPE == 'yolo' else PASCAL_ID_TO_NAME)
print(f'训练集mAP: {train_map:.4f}')

# 计算验证集mAP
print("计算验证集mAP...")
val_map = evaluate_map(model, data_loader_val, DEVICE, DATASET_TYPE, 
                      category_id_to_name if DATASET_TYPE == 'coco' or DATASET_TYPE == 'yolo' else PASCAL_ID_TO_NAME)
print(f'验证集mAP: {val_map:.4f}')

# 计算测试集mAP
print("计算测试集mAP...")
test_map = evaluate_map(model, data_loader_test, DEVICE, DATASET_TYPE, 
                       category_id_to_name if DATASET_TYPE == 'coco' or DATASET_TYPE == 'yolo' else PASCAL_ID_TO_NAME)
print(f'测试集mAP: {test_map:.4f}')

# 记录最终结果
logger.info(f"Final results - Training mAP: {train_map:.4f}, Validation mAP: {val_map:.4f}, Test mAP: {test_map:.4f}")



# --- About YOLOv3 ---

# --- End ---
print("Code execution completed.")