| """COCO-style evaluation datasets that return image + annotation paths/ids.""" |
|
|
| from __future__ import annotations |
|
|
| import json |
| import os |
| from typing import Any, Dict, List, Optional |
|
|
| import torch |
| from PIL import Image |
| from torch.utils.data import Dataset |
| from torchvision import transforms |
|
|
|
|
| class COCOEvalDataset(Dataset): |
| """COCO val images with annotation ids for metric evaluation. |
| |
| Returns a dict: |
| image: FloatTensor CxHxW in [0,1] |
| image_id: int |
| file_name: str |
| height, width: int |
| path: str |
| """ |
|
|
| def __init__( |
| self, |
| coco_root: str, |
| ann_file: str, |
| image_prefix: str = "val2017", |
| transform=None, |
| ): |
| self.coco_root = coco_root |
| self.image_dir = os.path.join(coco_root, image_prefix) |
| self.ann_file = ann_file if os.path.isabs(ann_file) else os.path.join(coco_root, ann_file) |
| self.transform = transform or transforms.ToTensor() |
|
|
| with open(self.ann_file) as f: |
| coco = json.load(f) |
| self.images: List[Dict[str, Any]] = sorted(coco["images"], key=lambda x: x["id"]) |
| self.categories = coco.get("categories", []) |
|
|
| def __len__(self): |
| return len(self.images) |
|
|
| def __getitem__(self, index: int) -> Dict[str, Any]: |
| info = self.images[index] |
| path = os.path.join(self.image_dir, info["file_name"]) |
| img = Image.open(path).convert("RGB") |
| tensor = self.transform(img) |
| return { |
| "image": tensor, |
| "image_id": int(info["id"]), |
| "file_name": info["file_name"], |
| "height": int(info["height"]), |
| "width": int(info["width"]), |
| "path": path, |
| } |
|
|
|
|
| def coco_eval_collate(batch: List[Dict[str, Any]]) -> Dict[str, Any]: |
| """Collate that keeps variable-size images as a list (batch_size usually 1).""" |
| return { |
| "images": [b["image"] for b in batch], |
| "image_ids": [b["image_id"] for b in batch], |
| "file_names": [b["file_name"] for b in batch], |
| "heights": [b["height"] for b in batch], |
| "widths": [b["width"] for b in batch], |
| "paths": [b["path"] for b in batch], |
| } |
|
|
|
|
| |
| TASK_ANN_FILES = { |
| "detection": "annotations/instances_val2017.json", |
| "instance": "annotations/instances_val2017.json", |
| "semantic": "annotations/panoptic_val2017.json", |
| "panoptic": "annotations/panoptic_val2017.json", |
| "pose": "annotations/coco_wholebody_val_v1.0.json", |
| } |
|
|
