helpers.py 35 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797
  1. # Copyright The PyTorch Lightning team.
  2. #
  3. # Licensed under the Apache License, Version 2.0 (the "License");
  4. # you may not use this file except in compliance with the License.
  5. # You may obtain a copy of the License at
  6. #
  7. # http://www.apache.org/licenses/LICENSE-2.0
  8. #
  9. # Unless required by applicable law or agreed to in writing, software
  10. # distributed under the License is distributed on an "AS IS" BASIS,
  11. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  12. # See the License for the specific language governing permissions and
  13. # limitations under the License.
  14. import contextlib
  15. import io
  16. import json
  17. from collections.abc import Sequence
  18. from importlib.metadata import version
  19. from types import ModuleType
  20. from typing import Any, Dict, List, Literal, Optional, Tuple, Union
  21. import numpy as np
  22. import torch
  23. from lightning_utilities import apply_to_collection
  24. from torch import Tensor
  25. from torchmetrics.utilities import rank_zero_warn
  26. from torchmetrics.utilities.imports import (
  27. _FASTER_COCO_EVAL_AVAILABLE,
  28. _PYCOCOTOOLS_AVAILABLE,
  29. _PYCOCOTOOLS_GREATER_EQUAL_2_0_9,
  30. )
  31. if not (_PYCOCOTOOLS_AVAILABLE or _FASTER_COCO_EVAL_AVAILABLE):
  32. __doctest_skip__ = [
  33. "CocoBackend.tm_to_coco",
  34. "CocoBackend.coco_to_tm",
  35. ]
  36. def _input_validator(
  37. preds: Sequence[dict[str, Tensor]],
  38. targets: Sequence[dict[str, Tensor]],
  39. iou_type: Union[Literal["bbox", "segm"], tuple[Literal["bbox", "segm"], ...]] = "bbox",
  40. ignore_score: bool = False,
  41. ) -> None:
  42. """Ensure the correct input format of `preds` and `targets`."""
  43. if isinstance(iou_type, str):
  44. iou_type = (iou_type,)
  45. name_map = {"bbox": "boxes", "segm": "masks"}
  46. if any(tp not in name_map for tp in iou_type):
  47. raise Exception(f"IOU type {iou_type} is not supported")
  48. item_val_name = [name_map[tp] for tp in iou_type]
  49. if not isinstance(preds, Sequence):
  50. raise ValueError(f"Expected argument `preds` to be of type Sequence, but got {preds}")
  51. if not isinstance(targets, Sequence):
  52. raise ValueError(f"Expected argument `target` to be of type Sequence, but got {targets}")
  53. if len(preds) != len(targets):
  54. raise ValueError(
  55. f"Expected argument `preds` and `target` to have the same length, but got {len(preds)} and {len(targets)}"
  56. )
  57. for k in [*item_val_name, "labels"] + (["scores"] if not ignore_score else []):
  58. if any(k not in p for p in preds):
  59. raise ValueError(f"Expected all dicts in `preds` to contain the `{k}` key")
  60. for k in [*item_val_name, "labels"]:
  61. if any(k not in p for p in targets):
  62. raise ValueError(f"Expected all dicts in `target` to contain the `{k}` key")
  63. for ivn in item_val_name:
  64. if not all(isinstance(pred[ivn], Tensor) for pred in preds):
  65. raise ValueError(f"Expected all {ivn} in `preds` to be of type Tensor")
  66. if not ignore_score and not all(isinstance(pred["scores"], Tensor) for pred in preds):
  67. raise ValueError("Expected all scores in `preds` to be of type Tensor")
  68. if not all(isinstance(pred["labels"], Tensor) for pred in preds):
  69. raise ValueError("Expected all labels in `preds` to be of type Tensor")
  70. for ivn in item_val_name:
  71. if not all(isinstance(target[ivn], Tensor) for target in targets):
  72. raise ValueError(f"Expected all {ivn} in `target` to be of type Tensor")
  73. if not all(isinstance(target["labels"], Tensor) for target in targets):
  74. raise ValueError("Expected all labels in `target` to be of type Tensor")
  75. for i, item in enumerate(targets):
  76. for ivn in item_val_name:
  77. if item[ivn].size(0) != item["labels"].size(0):
  78. raise ValueError(
  79. f"Input '{ivn}' and labels of sample {i} in targets have a"
  80. f" different length (expected {item[ivn].size(0)} labels, got {item['labels'].size(0)})"
  81. )
  82. if ignore_score:
  83. return
  84. for i, item in enumerate(preds):
  85. for ivn in item_val_name:
  86. if not (item[ivn].size(0) == item["labels"].size(0) == item["scores"].size(0)):
  87. raise ValueError(
  88. f"Input '{ivn}', labels and scores of sample {i} in predictions have a"
  89. f" different length (expected {item[ivn].size(0)} labels and scores,"
  90. f" got {item['labels'].size(0)} labels and {item['scores'].size(0)})"
  91. )
  92. def _fix_empty_tensors(boxes: Tensor) -> Tensor:
  93. """Empty tensors can cause problems in DDP mode, this methods corrects them."""
  94. if boxes.numel() == 0 and boxes.ndim == 1:
  95. return boxes.unsqueeze(0)
  96. return boxes
  97. def _validate_iou_type_arg(
  98. iou_type: Union[Literal["bbox", "segm"], tuple[Literal["bbox", "segm"], ...]] = "bbox",
  99. ) -> tuple[Literal["bbox", "segm"], ...]:
  100. """Validate that iou type argument is correct."""
  101. allowed_iou_types = ("segm", "bbox")
  102. if isinstance(iou_type, str):
  103. iou_type = (iou_type,)
  104. if any(tp not in allowed_iou_types for tp in iou_type):
  105. raise ValueError(
  106. f"Expected argument `iou_type` to be one of {allowed_iou_types} or a tuple of, but got {iou_type}"
  107. )
  108. return iou_type
  109. def _load_coco_backend_tools(backend: Literal["pycocotools", "faster_coco_eval"]) -> tuple[object, object, ModuleType]:
  110. """Load the backend tools for the given backend."""
  111. if backend == "pycocotools":
  112. if not _PYCOCOTOOLS_AVAILABLE:
  113. raise ModuleNotFoundError(
  114. "Backend `pycocotools` in metric `MeanAveragePrecision` metric requires that `pycocotools` is"
  115. " installed. Please install with `pip install pycocotools` or `pip install torchmetrics[detection]`"
  116. )
  117. import pycocotools.mask as mask_utils
  118. from pycocotools.coco import COCO
  119. from pycocotools.cocoeval import COCOeval
  120. return COCO, COCOeval, mask_utils
  121. if not _FASTER_COCO_EVAL_AVAILABLE:
  122. raise ModuleNotFoundError(
  123. "Backend `faster_coco_eval` in metric `MeanAveragePrecision` metric requires that `faster-coco-eval` is"
  124. " installed. Please install with `pip install faster-coco-eval`."
  125. )
  126. from faster_coco_eval import COCO
  127. from faster_coco_eval import COCOeval_faster as COCOeval
  128. from faster_coco_eval.core import mask as mask_utils
  129. return COCO, COCOeval, mask_utils
  130. class CocoBackend:
  131. """Backend implementation for COCO-style Mean Average Precision (mAP) calculation.
  132. This class provides the core functionality for evaluating object detection and instance
  133. segmentation predictions using the Common Objects in Context (COCO) evaluation protocol.
  134. It supports both the standard 'pycocotools' and optimized 'faster_coco_eval' backends.
  135. It's used for calculation of mAP in MeanAveragePrecision class. It's a backend that abstracts
  136. away the mAP calculation with coco package
  137. Args:
  138. backend (str): Either 'pycocotools' or 'faster_coco_eval'
  139. """
  140. def __init__(self, backend: Literal["pycocotools", "faster_coco_eval"]) -> None:
  141. if backend not in ("pycocotools", "faster_coco_eval"):
  142. raise ValueError(
  143. f"Expected argument `backend` to be one of ('pycocotools', 'faster_coco_eval') but got {backend}"
  144. )
  145. self.backend = backend
  146. @property
  147. def coco(self) -> object:
  148. """Returns the coco module for the given backend."""
  149. coco, _, _ = _load_coco_backend_tools(self.backend)
  150. return coco
  151. @property
  152. def cocoeval(self) -> object:
  153. """Returns the coco eval module for the given backend."""
  154. _, cocoeval, _ = _load_coco_backend_tools(self.backend)
  155. return cocoeval
  156. @property
  157. def mask_utils(self) -> object:
  158. """Returns the mask utils object for the given backend."""
  159. _, _, mask_utils = _load_coco_backend_tools(self.backend)
  160. return mask_utils
  161. def _get_coco_datasets(
  162. self,
  163. groundtruth_labels: List[Tensor],
  164. groundtruth_box: List[Tensor],
  165. groundtruth_mask: List[Tensor],
  166. groundtruth_crowds: List[Tensor],
  167. groundtruth_area: List[Tensor],
  168. detection_labels: List[Tensor],
  169. detection_box: List[Tensor],
  170. detection_mask: List[Tensor],
  171. detection_scores: List[Tensor],
  172. iou_type: Union[Literal["bbox", "segm"], tuple[Literal["bbox", "segm"], ...]] = ("bbox",),
  173. average: Literal["macro", "micro"] = "micro",
  174. ) -> tuple[object, object]:
  175. """Returns the coco datasets for the target and the predictions."""
  176. if average == "micro":
  177. # for micro averaging we set everything to be the same class
  178. groundtruth_labels = apply_to_collection(groundtruth_labels, Tensor, lambda x: torch.zeros_like(x))
  179. detection_labels = apply_to_collection(detection_labels, Tensor, lambda x: torch.zeros_like(x))
  180. coco_target, coco_preds = self.coco(), self.coco() # type: ignore[operator]
  181. # Equivalent to _get_classes function
  182. all_labels = (
  183. torch.cat(detection_labels + groundtruth_labels).unique().cpu().tolist()
  184. if len(detection_labels) > 0 or len(groundtruth_labels) > 0
  185. else []
  186. )
  187. coco_target.dataset = self._get_coco_format(
  188. labels=groundtruth_labels,
  189. boxes=groundtruth_box if len(groundtruth_box) > 0 else None,
  190. masks=groundtruth_mask if len(groundtruth_mask) > 0 else None,
  191. crowds=groundtruth_crowds,
  192. area=groundtruth_area,
  193. iou_type=iou_type,
  194. all_labels=all_labels,
  195. average=average,
  196. )
  197. coco_preds.dataset = self._get_coco_format(
  198. labels=detection_labels,
  199. boxes=detection_box if len(detection_box) > 0 else None,
  200. masks=detection_mask if len(detection_mask) > 0 else None,
  201. scores=detection_scores,
  202. iou_type=iou_type,
  203. all_labels=all_labels,
  204. average=average,
  205. )
  206. with contextlib.redirect_stdout(io.StringIO()):
  207. coco_target.createIndex()
  208. coco_preds.createIndex()
  209. return coco_preds, coco_target
  210. def _coco_stats_to_tensor_dict(
  211. self, stats: list[float], prefix: str, max_detection_thresholds: list[int]
  212. ) -> dict[str, Tensor]:
  213. """Converts the output of COCOeval.stats to a dict of tensors."""
  214. mdt = max_detection_thresholds
  215. return {
  216. f"{prefix}map": torch.tensor([stats[0]], dtype=torch.float32),
  217. f"{prefix}map_50": torch.tensor([stats[1]], dtype=torch.float32),
  218. f"{prefix}map_75": torch.tensor([stats[2]], dtype=torch.float32),
  219. f"{prefix}map_small": torch.tensor([stats[3]], dtype=torch.float32),
  220. f"{prefix}map_medium": torch.tensor([stats[4]], dtype=torch.float32),
  221. f"{prefix}map_large": torch.tensor([stats[5]], dtype=torch.float32),
  222. f"{prefix}mar_{mdt[0]}": torch.tensor([stats[6]], dtype=torch.float32),
  223. f"{prefix}mar_{mdt[1]}": torch.tensor([stats[7]], dtype=torch.float32),
  224. f"{prefix}mar_{mdt[2]}": torch.tensor([stats[8]], dtype=torch.float32),
  225. f"{prefix}mar_small": torch.tensor([stats[9]], dtype=torch.float32),
  226. f"{prefix}mar_medium": torch.tensor([stats[10]], dtype=torch.float32),
  227. f"{prefix}mar_large": torch.tensor([stats[11]], dtype=torch.float32),
  228. }
  229. @staticmethod
  230. def coco_to_tm(
  231. coco_preds: str,
  232. coco_target: str,
  233. iou_type: Union[Literal["bbox", "segm"], tuple[Literal["bbox", "segm"], ...]] = ("bbox",),
  234. backend: Literal["pycocotools", "faster_coco_eval"] = "pycocotools",
  235. ) -> tuple[list[dict[str, Tensor]], list[dict[str, Tensor]]]:
  236. """Utility function for converting .json coco format files to the input format of the mAP metric.
  237. The function accepts a file for the predictions and a file for the target in coco format and converts them to
  238. a list of dictionaries containing the boxes, labels and scores in the input format of mAP metric.
  239. Args:
  240. coco_preds: Path to the json file containing the predictions in coco format
  241. coco_target: Path to the json file containing the targets in coco format
  242. iou_type: Type of input, either `bbox` for bounding boxes or `segm` for segmentation masks
  243. backend: Backend to use for the conversion. Either `pycocotools` or `faster_coco_eval`.
  244. Returns:
  245. A tuple containing the predictions and targets in the input format of mAP metric. Each element of the
  246. tuple is a list of dictionaries containing the boxes, labels and scores.
  247. Example:
  248. >>> # File formats are defined at https://cocodataset.org/#format-data
  249. >>> # Example files can be found at
  250. >>> # https://github.com/cocodataset/cocoapi/tree/master/results
  251. >>> from torchmetrics.detection import MeanAveragePrecision
  252. >>> preds, target = MeanAveragePrecision().coco_to_tm(
  253. ... "instances_val2014_fakebbox100_results.json",
  254. ... "val2014_fake_eval_res.txt.json"
  255. ... iou_type="bbox"
  256. ... ) # doctest: +SKIP
  257. """
  258. iou_type = _validate_iou_type_arg(iou_type)
  259. coco, _, _ = _load_coco_backend_tools(backend)
  260. with contextlib.redirect_stdout(io.StringIO()):
  261. gt = coco(coco_target) # type: ignore[operator]
  262. dt = gt.loadRes(coco_preds)
  263. gt_dataset = gt.dataset["annotations"]
  264. dt_dataset = dt.dataset["annotations"]
  265. target: dict = {}
  266. for t in gt_dataset:
  267. if t["image_id"] not in target:
  268. target[t["image_id"]] = {
  269. "labels": [],
  270. "iscrowd": [],
  271. "area": [],
  272. }
  273. if "bbox" in iou_type:
  274. target[t["image_id"]]["boxes"] = []
  275. if "segm" in iou_type:
  276. target[t["image_id"]]["masks"] = []
  277. if "bbox" in iou_type:
  278. target[t["image_id"]]["boxes"].append(t["bbox"])
  279. if "segm" in iou_type:
  280. target[t["image_id"]]["masks"].append(gt.annToMask(t))
  281. target[t["image_id"]]["labels"].append(t["category_id"])
  282. target[t["image_id"]]["iscrowd"].append(t["iscrowd"])
  283. target[t["image_id"]]["area"].append(t["area"])
  284. preds: dict = {}
  285. for p in dt_dataset:
  286. if p["image_id"] not in preds:
  287. preds[p["image_id"]] = {"scores": [], "labels": []}
  288. if "bbox" in iou_type:
  289. preds[p["image_id"]]["boxes"] = []
  290. if "segm" in iou_type:
  291. preds[p["image_id"]]["masks"] = []
  292. if "bbox" in iou_type:
  293. preds[p["image_id"]]["boxes"].append(p["bbox"])
  294. if "segm" in iou_type:
  295. preds[p["image_id"]]["masks"].append(gt.annToMask(p))
  296. preds[p["image_id"]]["scores"].append(p["score"])
  297. preds[p["image_id"]]["labels"].append(p["category_id"])
  298. for k in target: # add empty predictions for images without predictions
  299. if k not in preds:
  300. preds[k] = {"scores": [], "labels": []}
  301. if "bbox" in iou_type:
  302. preds[k]["boxes"] = []
  303. if "segm" in iou_type:
  304. preds[k]["masks"] = []
  305. batched_preds, batched_target = [], []
  306. for key in target:
  307. bp = {
  308. "scores": torch.tensor(preds[key]["scores"], dtype=torch.float32),
  309. "labels": torch.tensor(preds[key]["labels"], dtype=torch.int32),
  310. }
  311. if "bbox" in iou_type:
  312. bp["boxes"] = torch.tensor(np.array(preds[key]["boxes"]), dtype=torch.float32)
  313. if "segm" in iou_type:
  314. bp["masks"] = torch.tensor(np.array(preds[key]["masks"]), dtype=torch.uint8)
  315. batched_preds.append(bp)
  316. bt = {
  317. "labels": torch.tensor(target[key]["labels"], dtype=torch.int32),
  318. "iscrowd": torch.tensor(target[key]["iscrowd"], dtype=torch.int32),
  319. "area": torch.tensor(target[key]["area"], dtype=torch.float32),
  320. }
  321. if "bbox" in iou_type:
  322. bt["boxes"] = torch.tensor(target[key]["boxes"], dtype=torch.float32)
  323. if "segm" in iou_type:
  324. bt["masks"] = torch.tensor(np.array(target[key]["masks"]), dtype=torch.uint8)
  325. batched_target.append(bt)
  326. return batched_preds, batched_target
  327. def tm_to_coco(
  328. self,
  329. groundtruth_labels: List[Tensor],
  330. groundtruth_box: List[Tensor],
  331. groundtruth_mask: List[Tensor],
  332. groundtruth_crowds: List[Tensor],
  333. groundtruth_area: List[Tensor],
  334. detection_labels: List[Tensor],
  335. detection_box: List[Tensor],
  336. detection_mask: List[Tensor],
  337. detection_scores: List[Tensor],
  338. name: str = "tm_map_input",
  339. iou_type: Union[Literal["bbox", "segm"], tuple[Literal["bbox", "segm"], ...]] = ("bbox",),
  340. average: Literal["macro", "micro"] = "micro",
  341. ) -> None:
  342. """Utility function for converting the input for mAP metric to coco format and saving it to a json file.
  343. This function should be used after calling `.update(...)` or `.forward(...)` on all data that should be written
  344. to the file, as the input is then internally cached. The function then converts to information to coco format
  345. and writes it to json files.
  346. Args:
  347. groundtruth_labels: List of tensors containing the ground truth labels
  348. groundtruth_box: List of tensors containing the ground truth bounding boxes
  349. groundtruth_mask: List of tensors containing the ground truth segmentation masks
  350. groundtruth_crowds: List of tensors indicating whether ground truth annotations are crowd annotations
  351. groundtruth_area: List of tensors containing the area of ground truth annotations
  352. detection_labels: List of tensors containing the predicted labels
  353. detection_box: List of tensors containing the predicted bounding boxes
  354. detection_mask: List of tensors containing the predicted segmentation masks
  355. detection_scores: List of tensors containing the confidence scores for predictions
  356. name: Name of the output file, which will be appended with "_preds.json" and "_target.json"
  357. iou_type: Type of IoU calculation to use. Can be either "bbox" for bounding box or "segm" for segmentation
  358. average: Type of averaging to use. Can be either "macro" or "micro"
  359. Example:
  360. >>> from torch import tensor
  361. >>> from torchmetrics.detection import MeanAveragePrecision
  362. >>> preds = [
  363. ... dict(
  364. ... boxes=tensor([[258.0, 41.0, 606.0, 285.0]]),
  365. ... scores=tensor([0.536]),
  366. ... labels=tensor([0]),
  367. ... )
  368. ... ]
  369. >>> target = [
  370. ... dict(
  371. ... boxes=tensor([[214.0, 41.0, 562.0, 285.0]]),
  372. ... labels=tensor([0]),
  373. ... )
  374. ... ]
  375. >>> metric = MeanAveragePrecision(iou_type="bbox")
  376. >>> metric.update(preds, target)
  377. >>> metric.tm_to_coco("tm_map_input")
  378. """
  379. all_labels = (
  380. torch.cat(detection_labels + groundtruth_labels).unique().cpu().tolist()
  381. if len(detection_labels) > 0 or len(groundtruth_labels) > 0
  382. else []
  383. )
  384. target_dataset = self._get_coco_format(
  385. labels=groundtruth_labels,
  386. boxes=groundtruth_box if len(groundtruth_box) > 0 else None,
  387. masks=groundtruth_mask if len(groundtruth_mask) > 0 else None,
  388. crowds=groundtruth_crowds,
  389. area=groundtruth_area,
  390. all_labels=all_labels,
  391. iou_type=iou_type,
  392. average=average,
  393. )
  394. preds_dataset = self._get_coco_format(
  395. labels=detection_labels,
  396. boxes=detection_box if len(detection_box) > 0 else None,
  397. masks=detection_mask if len(detection_mask) > 0 else None,
  398. scores=detection_scores,
  399. all_labels=all_labels,
  400. iou_type=iou_type,
  401. average=average,
  402. )
  403. if "segm" in iou_type:
  404. # the rle masks needs to be decoded to be written to a file
  405. preds_dataset["annotations"] = apply_to_collection(
  406. preds_dataset["annotations"], dtype=bytes, function=lambda x: x.decode("utf-8")
  407. )
  408. preds_dataset["annotations"] = apply_to_collection(
  409. preds_dataset["annotations"],
  410. dtype=(np.uint32, np.uint64),
  411. function=lambda x: int(x),
  412. )
  413. target_dataset = apply_to_collection(target_dataset, dtype=bytes, function=lambda x: x.decode("utf-8"))
  414. target_dataset = apply_to_collection(
  415. target_dataset, dtype=(np.uint32, np.uint64), function=lambda x: int(x)
  416. )
  417. preds_json = json.dumps(preds_dataset["annotations"], indent=4)
  418. target_json = json.dumps(target_dataset, indent=4)
  419. with open(f"{name}_preds.json", "w") as f:
  420. f.write(preds_json)
  421. with open(f"{name}_target.json", "w") as f:
  422. f.write(target_json)
  423. def _get_coco_format(
  424. self,
  425. labels: List[Tensor],
  426. all_labels: List[Tensor],
  427. boxes: Optional[List[Tensor]] = None,
  428. masks: Optional[List[Tensor]] = None,
  429. scores: Optional[List[Tensor]] = None,
  430. crowds: Optional[List[Tensor]] = None,
  431. area: Optional[List[Tensor]] = None,
  432. iou_type: Union[Literal["bbox", "segm"], tuple[Literal["bbox", "segm"], ...]] = ("bbox",),
  433. average: Literal["macro", "micro"] = "micro",
  434. ) -> dict:
  435. """Transforms and returns all cached targets or predictions in COCO format.
  436. Format is defined at
  437. https://cocodataset.org/#format-data
  438. """
  439. images = []
  440. annotations = []
  441. annotation_id = 1 # has to start with 1, otherwise COCOEval results are wrong
  442. for image_id, image_labels in enumerate(labels):
  443. if boxes is not None:
  444. image_boxes = boxes[image_id]
  445. image_boxes = image_boxes.cpu().tolist()
  446. if masks is not None:
  447. image_masks = masks[image_id]
  448. if len(image_masks) == 0 and boxes is None:
  449. continue
  450. image_labels = image_labels.cpu().tolist() # type: ignore[assignment]
  451. images.append({"id": image_id})
  452. if "segm" in iou_type and len(image_masks) > 0:
  453. images[-1]["height"], images[-1]["width"] = image_masks[0][0][0], image_masks[0][0][1] # type: ignore[assignment]
  454. for k, image_label in enumerate(image_labels):
  455. if boxes is not None:
  456. image_box = image_boxes[k]
  457. if masks is not None and len(image_masks) > 0:
  458. image_mask = image_masks[k]
  459. image_mask = {"size": image_mask[0], "counts": image_mask[1]}
  460. if "bbox" in iou_type and len(image_box) != 4:
  461. raise ValueError(
  462. f"Invalid input box of sample {image_id}, element {k} (expected 4 values, got {len(image_box)})"
  463. )
  464. if not isinstance(image_label, int):
  465. raise ValueError(
  466. f"Invalid input class of sample {image_id}, element {k}"
  467. f" (expected value of type integer, got type {type(image_label)})"
  468. )
  469. area_stat_box = None
  470. area_stat_mask = None
  471. if area is not None and area[image_id][k].cpu().tolist() > 0: # type: ignore[operator]
  472. area_stat = area[image_id][k].cpu().tolist()
  473. else:
  474. area_stat = self.mask_utils.area(image_mask) if "segm" in iou_type else image_box[2] * image_box[3]
  475. if len(iou_type) > 1:
  476. area_stat_box = image_box[2] * image_box[3]
  477. area_stat_mask = self.mask_utils.area(image_mask)
  478. annotation = {
  479. "id": annotation_id,
  480. "image_id": image_id,
  481. "area": area_stat,
  482. "category_id": image_label,
  483. "iscrowd": crowds[image_id][k].cpu().tolist() if crowds is not None else 0,
  484. }
  485. if area_stat_box is not None:
  486. annotation["area_bbox"] = area_stat_box
  487. annotation["area_segm"] = area_stat_mask
  488. if boxes is not None:
  489. annotation["bbox"] = image_box
  490. if masks is not None:
  491. annotation["segmentation"] = image_mask
  492. if scores is not None:
  493. score = scores[image_id][k].cpu().tolist()
  494. if not isinstance(score, float):
  495. raise ValueError(
  496. f"Invalid input score of sample {image_id}, element {k}"
  497. f" (expected value of type float, got type {type(score)})"
  498. )
  499. annotation["score"] = score
  500. annotations.append(annotation)
  501. annotation_id += 1
  502. classes = [{"id": i, "name": str(i)} for i in all_labels] if average != "micro" else [{"id": 0, "name": "0"}]
  503. result = {
  504. "images": images,
  505. "annotations": annotations,
  506. "categories": classes,
  507. }
  508. if _PYCOCOTOOLS_GREATER_EQUAL_2_0_9:
  509. result["info"] = {
  510. "description": f"Dummy info generated by tm_to_coco to support pycocotools {version('pycocotools')}"
  511. }
  512. return result
  513. def _warning_on_too_many_detections(limit: int) -> None:
  514. rank_zero_warn(
  515. f"Encountered more than {limit} detections in a single image. This means that certain detections with the"
  516. " lowest scores will be ignored, that may have an undesirable impact on performance. Please consider adjusting"
  517. " the `max_detection_threshold` to suit your use case. To disable this warning, set attribute class"
  518. " `warn_on_many_detections=False`, after initializing the metric.",
  519. UserWarning,
  520. )
  521. def _get_safe_item_values(
  522. iou_type: Union[Literal["bbox", "segm"], Tuple[Literal["bbox", "segm"], ...]],
  523. box_format: str,
  524. max_detection_thresholds: List[int],
  525. coco_backend: CocoBackend,
  526. item: dict[str, Any],
  527. warn: bool = False,
  528. ) -> tuple[Optional[Tensor], Optional[tuple]]:
  529. """Convert and return the boxes or masks from the item depending on the iou_type.
  530. Args:
  531. iou_type:
  532. Type of input to process. Supported types are:
  533. - "bbox": Process bounding boxes
  534. - "segm": Process segmentation masks
  535. box_format:
  536. Input format of given boxes. Supported formats are:
  537. - 'xyxy': boxes are represented via corners, x1, y1 being top left and x2, y2 being bottom right.
  538. - 'xywh': boxes are represented via corner, width and height, x1, y2 being top left, w, h being
  539. width and height.
  540. - 'cxcywh': boxes are represented via centre, width and height, cx, cy being center of box, w, h being
  541. width and height.
  542. max_detection_thresholds:
  543. List of thresholds on maximum detections per image. Used to determine if warnings should be raised
  544. when the number of detections exceeds these thresholds.
  545. coco_backend:
  546. The COCO evaluation backend class type to use for processing the items.
  547. item:
  548. Input dictionary containing the boxes or masks to be processed, along with other detection information.
  549. warn:
  550. Whether to warn if the number of boxes or masks exceeds the max_detection_thresholds.
  551. Default is False.
  552. Returns:
  553. A tuple containing processed boxes or masks depending on the iou_type. The first element is the
  554. tensor representation, and the second element contains additional metadata if applicable.
  555. """
  556. from torchvision.ops import box_convert
  557. output = [None, None]
  558. if "bbox" in iou_type:
  559. boxes = _fix_empty_tensors(item["boxes"])
  560. if boxes.numel() > 0:
  561. boxes = box_convert(boxes, in_fmt=box_format, out_fmt="xywh")
  562. output[0] = boxes # type: ignore[call-overload]
  563. if "segm" in iou_type:
  564. masks = []
  565. for i in item["masks"].cpu().numpy():
  566. rle = coco_backend.mask_utils.encode(np.asfortranarray(i))
  567. masks.append((tuple(rle["size"]), rle["counts"]))
  568. output[1] = tuple(masks) # type: ignore[call-overload]
  569. def _valid_output_len(idx: int) -> bool:
  570. val = output[idx]
  571. if val is None:
  572. return False
  573. return len(val) > max_detection_thresholds[-1]
  574. if warn and (_valid_output_len(0) or _valid_output_len(1)):
  575. _warning_on_too_many_detections(max_detection_thresholds[-1])
  576. return output # type: ignore[return-value]
  577. def _get_classes(detection_labels: List[Tensor], groundtruth_labels: List[Tensor]) -> List[int]:
  578. if len(detection_labels) > 0 or len(groundtruth_labels) > 0:
  579. return torch.cat(detection_labels + groundtruth_labels).unique().cpu().tolist()
  580. return []
  581. def _calculate_map_with_coco(
  582. coco_backend: CocoBackend,
  583. groundtruth_labels: List[Tensor],
  584. groundtruth_box: List[Tensor],
  585. groundtruth_mask: List[Tensor],
  586. groundtruth_crowds: List[Tensor],
  587. groundtruth_area: List[Tensor],
  588. detection_labels: List[Tensor],
  589. detection_box: List[Tensor],
  590. detection_mask: List[Tensor],
  591. detection_scores: List[Tensor],
  592. iou_type: Union[Literal["bbox", "segm"], Tuple[Literal["bbox", "segm"], ...]],
  593. average: Literal["macro", "micro"],
  594. iou_thresholds: List[float],
  595. rec_thresholds: List[float],
  596. max_detection_thresholds: List[int],
  597. class_metrics: bool,
  598. extended_summary: bool,
  599. ) -> Dict[str, Tensor]:
  600. coco_preds, coco_target = coco_backend._get_coco_datasets(
  601. groundtruth_labels,
  602. groundtruth_box,
  603. groundtruth_mask,
  604. groundtruth_crowds,
  605. groundtruth_area,
  606. detection_labels,
  607. detection_box,
  608. detection_mask,
  609. detection_scores,
  610. iou_type,
  611. average=average,
  612. )
  613. result_dict = {}
  614. with contextlib.redirect_stdout(io.StringIO()):
  615. for i_type in iou_type:
  616. prefix = "" if len(iou_type) == 1 else f"{i_type}_"
  617. if len(iou_type) > 1:
  618. # the area calculation is different for bbox and segm and therefore to get the small, medium and
  619. # large values correct we need to dynamically change the area attribute of the annotations
  620. for anno in coco_preds.dataset["annotations"]:
  621. anno["area"] = anno[f"area_{i_type}"]
  622. if len(coco_preds.imgs) == 0 or len(coco_target.imgs) == 0:
  623. result_dict.update(
  624. coco_backend._coco_stats_to_tensor_dict(
  625. 12 * [-1.0], prefix=prefix, max_detection_thresholds=max_detection_thresholds
  626. )
  627. )
  628. else:
  629. coco_eval = coco_backend.cocoeval(coco_target, coco_preds, iouType=i_type) # type: ignore[operator]
  630. coco_eval.params.iouThrs = np.array(iou_thresholds, dtype=np.float64)
  631. coco_eval.params.recThrs = np.array(rec_thresholds, dtype=np.float64)
  632. coco_eval.params.maxDets = max_detection_thresholds
  633. coco_eval.evaluate()
  634. coco_eval.accumulate()
  635. coco_eval.summarize()
  636. stats = coco_eval.stats
  637. result_dict.update(
  638. coco_backend._coco_stats_to_tensor_dict(
  639. stats, prefix=prefix, max_detection_thresholds=max_detection_thresholds
  640. )
  641. )
  642. summary = {}
  643. if extended_summary:
  644. summary = {
  645. f"{prefix}ious": apply_to_collection(
  646. coco_eval.ious, np.ndarray, lambda x: torch.tensor(x, dtype=torch.float32)
  647. ),
  648. f"{prefix}precision": torch.tensor(coco_eval.eval["precision"]),
  649. f"{prefix}recall": torch.tensor(coco_eval.eval["recall"]),
  650. f"{prefix}scores": torch.tensor(coco_eval.eval["scores"]),
  651. }
  652. result_dict.update(summary)
  653. # if class mode is enabled, evaluate metrics per class
  654. if class_metrics:
  655. # regardless of average method, reinitialize dataset to get rid of internal state which can
  656. # lead to wrong results when evaluating per class
  657. coco_preds, coco_target = coco_backend._get_coco_datasets(
  658. groundtruth_labels,
  659. groundtruth_box,
  660. groundtruth_mask,
  661. groundtruth_crowds,
  662. groundtruth_area,
  663. detection_labels,
  664. detection_box,
  665. detection_mask,
  666. detection_scores,
  667. iou_type,
  668. average="macro",
  669. )
  670. coco_eval = coco_backend.cocoeval(coco_target, coco_preds, iouType=i_type) # type: ignore[operator]
  671. coco_eval.params.iouThrs = np.array(iou_thresholds, dtype=np.float64)
  672. coco_eval.params.recThrs = np.array(rec_thresholds, dtype=np.float64)
  673. coco_eval.params.maxDets = max_detection_thresholds
  674. map_per_class_list = []
  675. mar_per_class_list = []
  676. for class_id in _get_classes(
  677. detection_labels=detection_labels, groundtruth_labels=groundtruth_labels
  678. ):
  679. coco_eval.params.catIds = [class_id]
  680. with contextlib.redirect_stdout(io.StringIO()):
  681. coco_eval.evaluate()
  682. coco_eval.accumulate()
  683. coco_eval.summarize()
  684. class_stats = coco_eval.stats
  685. map_per_class_list.append(torch.tensor([class_stats[0]]))
  686. mar_per_class_list.append(torch.tensor([class_stats[8]]))
  687. map_per_class_values = torch.tensor(map_per_class_list, dtype=torch.float32)
  688. mar_per_class_values = torch.tensor(mar_per_class_list, dtype=torch.float32)
  689. else:
  690. map_per_class_values = torch.tensor([-1], dtype=torch.float32)
  691. mar_per_class_values = torch.tensor([-1], dtype=torch.float32)
  692. prefix = "" if len(iou_type) == 1 else f"{i_type}_"
  693. result_dict.update(
  694. {
  695. f"{prefix}map_per_class": map_per_class_values,
  696. f"{prefix}mar_{max_detection_thresholds[-1]}_per_class": mar_per_class_values,
  697. },
  698. )
  699. result_dict.update({
  700. "classes": torch.tensor(
  701. _get_classes(detection_labels=detection_labels, groundtruth_labels=groundtruth_labels), dtype=torch.int32
  702. )
  703. })
  704. return result_dict