File size: 4,104 Bytes
6a5bb7e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 | """Faster R-CNN val集评估 —— 输出 mAP@50 / mAP@50-95 / P / R
用法：
python Scripts/eval_fasterrcnn.py
"""
import torch
from torch.utils.data import DataLoader
from torchvision.models.detection import fasterrcnn_resnet50_fpn_v2
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchmetrics.detection.mean_ap import MeanAveragePrecision
from pathlib import Path
from PIL import Image
import torchvision.transforms.functional as F
from torch.utils.data import Dataset
DATA_ROOT = Path("Data/Detection_dataset")
CKPT = Path("Detection_experiments/compare_fasterrcnn/best.pt")
NUM_CLASSES = 2
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
CONF_THRESH = 0.5
class GoatDataset(Dataset):
def __init__(self, split: str):
img_dir = DATA_ROOT / "images" / split
lbl_dir = DATA_ROOT / "labels" / split
self.samples = []
for img_path in sorted(img_dir.glob("*.jpg")):
lbl_path = lbl_dir / (img_path.stem + ".txt")
if lbl_path.exists():
self.samples.append((img_path, lbl_path))
def __len__(self):
return len(self.samples)
def __getitem__(self, idx):
img_path, lbl_path = self.samples[idx]
img = Image.open(img_path).convert("RGB")
w, h = img.size
img_tensor = F.to_tensor(img)
boxes, labels = [], []
for line in lbl_path.read_text().splitlines():
parts = line.strip().split()
if len(parts) != 5:
continue
_, cx, cy, bw, bh = map(float, parts)
x1 = (cx - bw / 2) * w
y1 = (cy - bh / 2) * h
x2 = (cx + bw / 2) * w
y2 = (cy + bh / 2) * h
boxes.append([x1, y1, x2, y2])
labels.append(1)
target = {
"boxes": torch.tensor(boxes, dtype=torch.float32) if boxes else torch.zeros((0, 4), dtype=torch.float32),
"labels": torch.tensor(labels, dtype=torch.int64) if labels else torch.zeros(0, dtype=torch.int64),
}
return img_tensor, target
def collate_fn(batch):
return tuple(zip(*batch))
def build_model():
from torchvision.models import ResNet50_Weights
model = fasterrcnn_resnet50_fpn_v2(weights=None, weights_backbone=ResNet50_Weights.IMAGENET1K_V2)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, NUM_CLASSES)
return model
def main():
model = build_model().to(DEVICE)
model.load_state_dict(torch.load(CKPT, map_location=DEVICE))
model.eval()
val_ds = GoatDataset("val")
val_loader = DataLoader(val_ds, batch_size=1, shuffle=False,
num_workers=2, collate_fn=collate_fn)
metric = MeanAveragePrecision(iou_type="bbox", class_metrics=False)
print(f"Evaluating on {len(val_ds)} val images...")
with torch.no_grad():
for imgs, targets in val_loader:
imgs = [img.to(DEVICE) for img in imgs]
outputs = model(imgs)
preds = []
for out in outputs:
keep = out["scores"] >= CONF_THRESH
preds.append({
"boxes": out["boxes"][keep].cpu(),
"scores": out["scores"][keep].cpu(),
"labels": out["labels"][keep].cpu(),
})
gts = []
for t in targets:
gts.append({
"boxes": t["boxes"],
"labels": t["labels"],
})
metric.update(preds, gts)
result = metric.compute()
map50_95 = result["map"].item()
map50 = result["map_50"].item()
mar = result["mar_100"].item()
print(f"\n{'='*45}")
print(f" Faster R-CNN (ImageNet backbone, 30 epochs)")
print(f"{'='*45}")
print(f" mAP@50 : {map50:.3f}")
print(f" mAP@50-95 : {map50_95:.3f}")
print(f" mAR@100 : {mar:.3f}")
print(f" conf_thresh : {CONF_THRESH}")
print(f"{'='*45}")
if __name__ == "__main__":
main()
|