如何基于ResNet50实现适配txt标注的单类别自定义目标检测

原生PyTorch、TensorFlow的内置数据集接口仅支持固定格式的标注读取，不需要转换你的原始.txt标注，只需要自行实现适配的数据集加载逻辑即可，以下是可直接落地的实现方案：

前提准备 #

先明确你自己的.txt标注的字段规则，常见的有两种：

• VOC风格：每行对应一个目标，字段为左上角x坐标 左上角y坐标 右下角x坐标 右下角y坐标（像素绝对值）
• YOLO风格：每行对应一个目标，字段为中心x归一化值 中心y归一化值 宽度归一化值 高度归一化值
  你只需要根据自己的标注规则调整后续加载逻辑即可，单类别任务不需要在标注中额外加类别字段。

PyTorch框架实现（更灵活适配自定义标注） #

我们直接使用TorchVision内置的fasterrcnn_resnet50_fpn模型，它的主干特征提取网络就是要求的ResNet50架构，完全符合你的需求。

1. 自定义数据集加载类 #

这个类会直接读取你的图片文件夹和.txt标注文件夹，不需要修改原始标注文件：

import os
import torch
from PIL import Image
from torch.utils.data import Dataset
from torchvision.models.detection import fasterrcnn_resnet50_fpn
from torchvision.transforms import functional as F

class TXTLabelDataset(Dataset):
    def __init__(self, img_dir, anno_dir, transforms=None):
        self.img_dir = img_dir
        self.anno_dir = anno_dir
        # 读取所有图片文件名
        self.img_list = [f for f in os.listdir(img_dir) if f.endswith(('.jpg', '.png', '.jpeg'))]
        self.transforms = transforms

    def __len__(self):
        return len(self.img_list)

    def __getitem__(self, idx):
        # 读取图片
        img_filename = self.img_list[idx]
        img_path = os.path.join(self.img_dir, img_filename)
        img = Image.open(img_path).convert("RGB")
        w, h = img.size

        # 读取对应同名的.txt标注
        anno_filename = os.path.splitext(img_filename)[0] + ".txt"
        anno_path = os.path.join(self.anno_dir, anno_filename)
        boxes = []
        with open(anno_path, "r", encoding="utf-8") as f:
            for line in f.readlines():
                line = line.strip().split()
                if not line:
                    continue
                # 以下为VOC风格标注的读取逻辑，如果你是YOLO风格可以替换为对应的转换逻辑
                x1 = float(line[0])
                y1 = float(line[1])
                x2 = float(line[2])
                y2 = float(line[3])
                boxes.append([x1, y1, x2, y2])

                # 如果是YOLO风格标注，用下面的逻辑替换即可
                # cx = float(line[0]) * w
                # cy = float(line[1]) * h
                # bw = float(line[2]) * w
                # bh = float(line[3]) * h
                # x1 = cx - bw/2
                # y1 = cy - bh/2
                # x2 = cx + bw/2
                # y2 = cy + bh/2
                # boxes.append([x1, y1, x2, y2])

        # 转换为模型要求的tensor格式
        boxes = torch.as_tensor(boxes, dtype=torch.float32)
        # 单类别任务，所有目标标签为1（0为默认背景类）
        labels = torch.ones((len(boxes),), dtype=torch.int64)
        target = {
            "boxes": boxes,
            "labels": labels,
            "image_id": torch.tensor([idx]),
            "area": (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0]),
            "iscrowd": torch.zeros((len(boxes),), dtype=torch.int64)
        }

        if self.transforms:
            img, target = self.transforms(img, target)
        else:
            img = F.to_tensor(img)
        return img, target

2. 初始化ResNet50架构的检测模型 #

# 加载预训练模型，单类别任务总类别数为2（背景+自定义目标）
model = fasterrcnn_resnet50_fpn(pretrained=True)
num_classes = 2
# 替换检测头适配单类别
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = fasterrcnn_resnet50_fpn(num_classes=num_classes).roi_heads.box_predictor

3. 训练逻辑 #

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model.to(device)

# 初始化数据集和加载器
dataset = TXTLabelDataset(
    img_dir="你的图片文件夹路径",
    anno_dir="你的.txt标注文件夹路径"
)
# 检测任务需要自定义collate_fn处理不同数量的标注框
data_loader = torch.utils.data.DataLoader(
    dataset,
    batch_size=2,
    shuffle=True,
    collate_fn=lambda x: tuple(zip(*x))
)

# 定义优化器
optimizer = torch.optim.SGD(
    [p for p in model.parameters() if p.requires_grad],
    lr=0.005,
    momentum=0.9,
    weight_decay=0.0005
)

# 训练循环
num_epochs = 10
for epoch in range(num_epochs):
    model.train()
    epoch_loss = 0
    for imgs, targets in data_loader:
        imgs = [img.to(device) for img in imgs]
        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
        # 模型直接返回损失值
        loss_dict = model(imgs, targets)
        total_loss = sum(loss for loss in loss_dict.values())
        epoch_loss += total_loss.item()
        # 反向传播更新参数
        optimizer.zero_grad()
        total_loss.backward()
        optimizer.step()
    print(f"Epoch {epoch+1}/{num_epochs}, 总损失: {epoch_loss:.4f}")

TensorFlow框架实现方案 #

逻辑和PyTorch完全一致，你只需要自定义继承tf.keras.utils.Sequence的数据集类，在__getitem__方法中读取图片和对应.txt标注，将标注转换为你使用的检测模型要求的输入格式即可，全程不需要修改原始.txt标注文件。

注意事项 #

• 确保图片和标注文件的文件名完全一致，仅后缀不同，避免加载时匹配失败
• 若标注为归一化坐标，必须提前乘以图片宽高转换为像素坐标再输入模型
• 单类别训练时不需要额外处理类别逻辑，所有标注框默认归为同一类即可

内容的提问来源于stack exchange，提问作者IronVenom