File size: 9,283 Bytes
4409fdb | 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 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 | """Shared locked-protocol dataset, model, and COCO evaluator utilities."""
from __future__ import annotations
import contextlib
import io
import json
import xml.etree.ElementTree as ET
from pathlib import Path
import torch
from PIL import Image
from torch.utils.data import Dataset
CLASS_NAMES = ("car", "motorbike", "bicycle", "chair", "diningtable", "bottle", "tvmonitor", "bus")
EXPECTED = {"train": 1647, "val": 183, "test": 400}
def resolve_image(dataset_root: Path, domain: str, basename: str) -> Path:
folder = "RGB_normal" if domain == "bright" else "RGB_Dark"
hits = list(dataset_root.rglob(f"{folder}/{basename}"))
if len(hits) != 1:
hits = [p for p in dataset_root.rglob(basename) if folder.lower() in str(p).lower()]
if len(hits) != 1:
raise RuntimeError(f"expected exactly one {domain} image for {basename}, found {hits}")
return hits[0]
def read_yolo(path: Path, width: int, height: int):
boxes, labels = [], []
for line in path.read_text().splitlines():
if not line.strip():
continue
cls, cx, cy, bw, bh = map(float, line.split())
if not 0 <= int(cls) < len(CLASS_NAMES):
raise RuntimeError(f"invalid class id {cls} in {path}")
boxes.append([(cx - bw / 2) * width, (cy - bh / 2) * height,
(cx + bw / 2) * width, (cy + bh / 2) * height])
labels.append(int(cls))
return (torch.tensor(boxes, dtype=torch.float32).reshape(-1, 4),
torch.tensor(labels, dtype=torch.int64))
def read_voc(path: Path):
boxes, labels = [], []
for obj in ET.parse(path).getroot().findall("object"):
name = obj.findtext("name")
if name not in CLASS_NAMES:
raise RuntimeError(f"unknown VOC class {name!r} in {path}")
box = obj.find("bndbox")
boxes.append([float(box.findtext("xmin")), float(box.findtext("ymin")),
float(box.findtext("xmax")), float(box.findtext("ymax"))])
labels.append(CLASS_NAMES.index(name))
return (torch.tensor(boxes, dtype=torch.float32).reshape(-1, 4),
torch.tensor(labels, dtype=torch.int64))
class LockedLODDataset(Dataset):
def __init__(self, manifest: Path, dataset_root: Path, labels_root: Path, limit: int | None = None):
payload = json.loads(Path(manifest).read_text())
self.domain = payload["domain"]
self.split = payload["split"]
self.items = payload["items"][:limit]
if len(payload["items"]) != EXPECTED[self.split]:
raise RuntimeError(f"{manifest}: expected {EXPECTED[self.split]} locked rows, got {len(payload['items'])}")
self.dataset_root, self.labels_root = Path(dataset_root), Path(labels_root)
def __len__(self):
return len(self.items)
def __getitem__(self, index):
item = self.items[index]
image = Image.open(resolve_image(self.dataset_root, self.domain, item["image_basename"])).convert("RGB")
# The public LOD Kaggle dataset contains the source VOC annotations.
# Prefer them, retaining copied YOLO labels only as a local fallback.
group = "RGB-normal-Annotations" if self.domain == "bright" else "RGB-dark-Annotations"
voc = self.dataset_root / group / group / (Path(item["image_basename"]).stem + ".xml")
if voc.is_file():
boxes, labels = read_voc(voc)
else:
label_domain = "normal" if self.domain == "bright" else "dark"
label = self.labels_root / label_domain / self.split / "labels" / item["label_basename"]
if not label.is_file():
raise FileNotFoundError(f"missing VOC and fallback label for {item['pair_id']}")
boxes, labels = read_yolo(label, *image.size)
return image, {"boxes": boxes, "labels": labels,
"image_id": torch.tensor(index),
"orig_size": torch.tensor([image.height, image.width])}
def collate(batch):
return tuple(zip(*batch))
def build_model(cfg):
detector = cfg["detector"]
if detector == "rtdetr":
from transformers import RTDetrV2ForObjectDetection, RTDetrImageProcessor
processor = RTDetrImageProcessor.from_pretrained(cfg["pretrained"])
model = RTDetrV2ForObjectDetection.from_pretrained(
cfg["pretrained"], num_labels=len(CLASS_NAMES),
id2label=dict(enumerate(CLASS_NAMES)), label2id={n: i for i, n in enumerate(CLASS_NAMES)},
ignore_mismatched_sizes=True,
)
return model, processor, "hf"
if detector == "fasterrcnn_r50_fpn":
from torchvision.models.detection import fasterrcnn_resnet50_fpn_v2, FasterRCNN_ResNet50_FPN_V2_Weights
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
model = fasterrcnn_resnet50_fpn_v2(weights=FasterRCNN_ResNet50_FPN_V2_Weights.COCO_V1)
model.roi_heads.box_predictor = FastRCNNPredictor(model.roi_heads.box_predictor.cls_score.in_features, len(CLASS_NAMES) + 1)
return model, None, "torchvision"
if detector == "ssdlite320_mobilenet_v3_large":
from torchvision.models.detection import ssdlite320_mobilenet_v3_large, SSDLite320_MobileNet_V3_Large_Weights
from torchvision.models.detection.ssdlite import SSDLiteClassificationHead
model = ssdlite320_mobilenet_v3_large(weights=SSDLite320_MobileNet_V3_Large_Weights.COCO_V1)
# TorchVision exposes each SSDLite predictor as a Sequential block;
# its depthwise Conv2d is nested under block[0][0]. Derive the
# feature widths from the loaded COCO head rather than hard-coding a
# version-specific list.
in_channels = [module[0][0].in_channels for module in model.head.classification_head.module_list]
num_anchors = model.anchor_generator.num_anchors_per_location()
model.head.classification_head = SSDLiteClassificationHead(
in_channels, num_anchors, len(CLASS_NAMES) + 1, torch.nn.BatchNorm2d
)
return model, None, "torchvision"
raise ValueError(f"unsupported detector: {detector}")
def image_tensors(images, device):
from torchvision.transforms.functional import pil_to_tensor
return [pil_to_tensor(image).float().div(255).to(device) for image in images]
def hf_batch(processor, images, targets, device):
annotations = []
for index, target in enumerate(targets):
boxes = target["boxes"]
xywh = torch.stack((boxes[:, 0], boxes[:, 1], boxes[:, 2] - boxes[:, 0], boxes[:, 3] - boxes[:, 1]), 1)
annotations.append({"image_id": index, "annotations": [
{"bbox": box.tolist(), "category_id": int(label), "area": float(box[2] * box[3]), "iscrowd": 0}
for box, label in zip(xywh, target["labels"])
]})
encoded = processor(images=list(images), annotations=annotations, return_tensors="pt")
return {key: (value.to(device) if hasattr(value, "to") else [{k: v.to(device) for k, v in x.items()} for x in value])
for key, value in encoded.items()}
class CanonicalMAP:
def __init__(self):
self.images, self.annotations, self.predictions, self.annotation_id = {}, [], [], 1
@staticmethod
def _xywh(box):
x1, y1, x2, y2 = map(float, box.tolist())
return [x1, y1, max(0.0, x2 - x1), max(0.0, y2 - y1)]
def update(self, predictions, targets):
for pred, target in zip(predictions, targets):
image_id = int(target["image_id"])
h, w = map(int, target["orig_size"].tolist())
self.images[image_id] = {"id": image_id, "height": h, "width": w}
for box, label in zip(target["boxes"], target["labels"]):
bbox = self._xywh(box)
self.annotations.append({"id": self.annotation_id, "image_id": image_id, "category_id": int(label) + 1,
"bbox": bbox, "area": bbox[2] * bbox[3], "iscrowd": 0})
self.annotation_id += 1
for box, score, label in zip(pred["boxes"], pred["scores"], pred["labels"]):
self.predictions.append({"image_id": image_id, "category_id": int(label) + 1,
"bbox": self._xywh(box), "score": float(score)})
def compute(self):
from pycocotools.coco import COCO
from pycocotools.cocoeval import COCOeval
gt = COCO(); gt.dataset = {"images": list(self.images.values()), "annotations": self.annotations,
"categories": [{"id": i + 1, "name": n} for i, n in enumerate(CLASS_NAMES)]}; gt.createIndex()
dt = gt.loadRes(self.predictions) if self.predictions else COCO()
if not self.predictions:
dt.dataset = {"images": list(self.images.values()), "annotations": [], "categories": gt.dataset["categories"]}; dt.createIndex()
evaluator = COCOeval(gt, dt, "bbox"); evaluator.params.imgIds = sorted(self.images); evaluator.params.catIds = list(range(1, 9)); evaluator.params.maxDets = [1, 10, 100]
with contextlib.redirect_stdout(io.StringIO()):
evaluator.evaluate(); evaluator.accumulate(); evaluator.summarize()
return {"map50_95": float(evaluator.stats[0]), "map50": float(evaluator.stats[1])}
|