Repository: voxel51/fiftyone

Issue:
[BUG] Memory leak when importing CocoDetection Dataset with random category_id into FiftyOne
### Describe the problem

When adding a dataset into fiftyone with the following code, the cell run until it crashes the notebook because of an out-of-memory error.
This happens if the dataset contains random category ids (instead of ids ranging from `0` to `num_categories-1`).
Fiftyone should either handle this case (random category ids), or write a clear error message (and definitely not eat all the memory until it gets killed).


### Code to reproduce issue

```python
ds = fo.Dataset.from_dir(
    name=f"name",
    overwrite=True, 
    persistent=True, 
    dataset_type=fo.types.COCODetectionDataset,
    data_path='path_to_images',  
    labels_path='path_to_json', 
    label_types=['segmentations'],
    use_polylines=False, 
    include_id=True,
)
```

Example of a json file that will lead to this issue (you have to have an image with filename `image.jpg`):
```json
{
  "categories": [
    {
      "id": 74673216,
      "name": "Solar Panel",
      "supercategory": "root"
    }
  ],
  "images": [
    {
      "id": 789,
      "file_name": "image.jpg",
      "height": 1024,
      "width": 1024,
      "extra": {}
    }
  ],
  "annotations": [
       {
      "id": 15567,
      "image_id": 789,
      "category_id": 74673216,
      "segmentation": [
        [
          445.6344437096958,
          354.34683544084845,
          452.1216184704017,
          348.23038026732564,
          411.9939701135346,
          303.93243513679033,
          406.06287473360334,
          310.2342501039047
        ]
      ],
      "area": 521.9080653650577,
      "bbox": [
        406.06287473360334,
        303.93243513679033,
        46.05874373679836,
        50.41440030405812
      ],
      "iscrowd": 0,
      "extra": {}
    }
  ]
}
```

### System information

-   **OS Platform and Distribution** (e.g., Linux Ubuntu 22.04): Ubuntu 20.04.6
-   **Python version** (`python --version`): 3.11.6
-   **FiftyOne version** (`fiftyone --version`): 0.23.7 
-   **FiftyOne installed from** (pip or source): pip


### Willingness to contribute

The FiftyOne Community encourages bug fix contributions. Would you or another
member of your organization be willing to contribute a fix for this bug to the
FiftyOne codebase?

-   [ ] Yes. I can contribute a fix for this bug independently
-   [x] Yes. I would be willing to contribute a fix for this bug with guidance
        from the FiftyOne community
-   [ ] No. I cannot contribute a bug fix at this time



Discussion:
Thanks @patrontheo . I'm not finding anywhere that discusses coco format, that the IDs must be sequential, though the examples that are often shown are that way.

It seems to be going bad in this function: https://github.com/voxel51/fiftyone/blob/develop/fiftyone/utils/coco.py#L1408-L1450

I guess the intent is to give all classes a name even if they happened to be missed in the categories dict. However if they are not intended to be sequential then you run into the problem here.

I think a reasonable thing to do is if the max ID is more than x times the number of categories, assume we don't have to fill in 0 through N and skip that whole for-loop. Maybe, x=10 or so?

Are you able to contribute that fix? We always appreciate new (and existing) contributors!
There may be downstream effects since that function returns a list of classes and would now return a dict. Might need some fixing up for this case.
Thanks for the pointer, I'll try to have a look tomorrow, and make a PR.
However I'm not sure why the for loop is useful. Why add classes that are not present in the given categories ?

Let's say we have:
```python
categories = [
    {
      "id": 10,
      "name": "Solar Panel",
      "supercategory": "root"
    }
  ]
```
The function `parse_coco_categories` returns:
`classes`: `['0', '1', '2', '3', '4', '5', '6', '7', '8', '9', 'Solar Panel']`
`supercategory_map`: `{'Solar Panel': {'id': 10, 'name': 'Solar Panel', 'supercategory': 'root'}}`

Do you have any idea why this is wanted ?
I'm not entirely sure to be honest, based on cursory research. Would have to dig/ask around some more.

But because it doesn't really hurt anything (except in this case), I'd rather not change the behavior in general - on the off chance someone is relying on that behavior
@swheaton 
In `_get_object_label_and_attributes` in [coco.py](https://github.com/voxel51/fiftyone/blob/0494cac8d8fb65fb6bd71f9fad1c67827ebd368e/fiftyone/utils/coco.py#L1312), the label is obtained using `label = classes[self.category_id]`. Then, if the category_is very large, `classes[self.category_id]` will return *IndexError: list index out of range* (that's probably why classes is created using the for loop).

I see some options to fix this:
- If the given category ids are not sequential (or are very large), raise an error, explaining that the annotation file should be changed to match this criteria. This could be a bit frustrating for the user that fiftyone does not handle this case.
    - We could provide the user a function that he can call, and that would read the annotation file, remap the ids to be sequential, and save the new annotation file.
    - We could pass a flag `remap_ids`(False by default), that would remap the ids to something sequential (although I would need some directions on how to pass the flag and in which functions).
- If the given category ids are not sequential, we could remap them to `0` to `num_categories-1` in `parse_coco_categories`. Although if the user try to get objects using the category_id there is in the annotation file, it would not work since we changed the ids. Maybe if a remapping is performed, we could print a warning for the user ?
@patrontheo thanks for continuing to look into it.
These 2 things seem to fix it.
1. If we have decided to not fill the classes list from `0..max(ids)` as discussed above, then we can return a `dict` from `id->name` instead. When we go to index via category ID, it will work just the same.
2. In `COCODetectionDatasetImporter.setup()` around line 560, we have to set `info["classes"] = list(classes.values())` if `classes` is a `dict`, so that we set the `default_classes` property on the dataset correctly (we'll get a document validation error otherwise)

feel free to put up a draft PR if you have any code so we can work there instead of in text here.

would also be great to have this scenario as a test if we're going to support it.

Resolving patch:
diff --git a/fiftyone/utils/coco.py b/fiftyone/utils/coco.py
--- a/fiftyone/utils/coco.py
+++ b/fiftyone/utils/coco.py
@@ -129,7 +129,8 @@ def add_coco_labels(
             will be created if necessary
         labels_or_path: a list of COCO annotations or the path to a JSON file
             containing such data on disk
-        classes: the list of class label strings
+        classes: the list of class label strings or a dict mapping class IDs to
+            class labels
         label_type ("detections"): the type of labels to load. Supported values
             are ``("detections", "segmentations", "keypoints")``
         coco_id_field (None): this parameter determines how to map the
@@ -194,6 +195,11 @@ def add_coco_labels(
     view.compute_metadata()
     widths, heights = view.values(["metadata.width", "metadata.height"])
 
+    if isinstance(classes, dict):
+        classes_map = classes
+    else:
+        classes_map = {i: label for i, label in enumerate(classes)}
+
     labels = []
     for _coco_objects, width, height in zip(coco_objects, widths, heights):
         frame_size = (width, height)
@@ -202,7 +208,7 @@ def add_coco_labels(
             _labels = _coco_objects_to_detections(
                 _coco_objects,
                 frame_size,
-                classes,
+                classes_map,
                 None,
                 False,
                 include_annotation_id,
@@ -212,7 +218,7 @@ def add_coco_labels(
                 _labels = _coco_objects_to_polylines(
                     _coco_objects,
                     frame_size,
-                    classes,
+                    classes_map,
                     None,
                     tolerance,
                     include_annotation_id,
@@ -221,7 +227,7 @@ def add_coco_labels(
                 _labels = _coco_objects_to_detections(
                     _coco_objects,
                     frame_size,
-                    classes,
+                    classes_map,
                     None,
                     True,
                     include_annotation_id,
@@ -230,7 +236,7 @@ def add_coco_labels(
             _labels = _coco_objects_to_keypoints(
                 _coco_objects,
                 frame_size,
-                classes,
+                classes_map,
                 None,
                 include_annotation_id,
             )
@@ -408,7 +414,7 @@ def __init__(
 
         self._label_types = _label_types
         self._info = None
-        self._classes = None
+        self._classes_map = None
         self._license_map = None
         self._supercategory_map = None
         self._image_paths_map = None
@@ -451,15 +457,16 @@ def __next__(self):
             frame_size = (width, height)
 
             if self.classes is not None and self.only_matching:
+                all_classes = list(self._classes_map.values())
                 coco_objects = _get_matching_objects(
-                    coco_objects, self.classes, self._classes
+                    coco_objects, self.classes, all_classes
                 )
 
             if "detections" in self._label_types:
                 detections = _coco_objects_to_detections(
                     coco_objects,
                     frame_size,
-                    self._classes,
+                    self._classes_map,
                     self._supercategory_map,
                     False,  # no segmentations
                     self.include_annotation_id,
@@ -472,7 +479,7 @@ def __next__(self):
                     segmentations = _coco_objects_to_polylines(
                         coco_objects,
                         frame_size,
-                        self._classes,
+                        self._classes_map,
                         self._supercategory_map,
                         self.tolerance,
                         self.include_annotation_id,
@@ -481,7 +488,7 @@ def __next__(self):
                     segmentations = _coco_objects_to_detections(
                         coco_objects,
                         frame_size,
-                        self._classes,
+                        self._classes_map,
                         self._supercategory_map,
                         True,  # load segmentations
                         self.include_annotation_id,
@@ -494,7 +501,7 @@ def __next__(self):
                 keypoints = _coco_objects_to_keypoints(
                     coco_objects,
                     frame_size,
-                    self._classes,
+                    self._classes_map,
                     self._supercategory_map,
                     self.include_annotation_id,
                 )
@@ -548,7 +555,7 @@ def setup(self):
         if self.labels_path is not None and os.path.isfile(self.labels_path):
             (
                 info,
-                classes,
+                classes_map,
                 supercategory_map,
                 images,
                 annotations,
@@ -556,6 +563,10 @@ def setup(self):
                 self.labels_path, extra_attrs=self.extra_attrs
             )
 
+            classes = None
+            if classes_map is not None:
+                classes = list(classes_map.values())
+
             if classes is not None:
                 info["classes"] = classes
 
@@ -582,7 +593,7 @@ def setup(self):
             }
         else:
             info = {}
-            classes = None
+            classes_map = None
             supercategory_map = None
             image_dicts_map = {}
             annotations = None
@@ -597,7 +608,7 @@ def setup(self):
             license_map = None
 
         self._info = info
-        self._classes = classes
+        self._classes_map = classes_map
         self._license_map = license_map
         self._supercategory_map = supercategory_map
         self._image_paths_map = image_paths_map
@@ -960,7 +971,7 @@ def __init__(
     def to_polyline(
         self,
         frame_size,
-        classes=None,
+        classes_map=None,
         supercategory_map=None,
         tolerance=None,
         include_id=False,
@@ -970,7 +981,7 @@ def to_polyline(
 
         Args:
             frame_size: the ``(width, height)`` of the image
-            classes (None): the list of classes
+            classes_map (None): a dict mapping class IDs to class labels
             supercategory_map (None): a dict mapping class names to category
                 dicts
             tolerance (None): a tolerance, in pixels, when generating
@@ -987,7 +998,7 @@ def to_polyline(
             return None
 
         label, attributes = self._get_object_label_and_attributes(
-            classes, supercategory_map, include_id
+            classes_map, supercategory_map, include_id
         )
         attributes.update(self.attributes)
 
@@ -1007,7 +1018,7 @@ def to_polyline(
     def to_keypoints(
         self,
         frame_size,
-        classes=None,
+        classes_map=None,
         supercategory_map=None,
         include_id=False,
     ):
@@ -1016,7 +1027,7 @@ def to_keypoints(
 
         Args:
             frame_size: the ``(width, height)`` of the image
-            classes (None): the list of classes
+            classes_map (None): a dict mapping class IDs to class labels
             supercategory_map (None): a dict mapping class names to category
                 dicts
             include_id (False): whether to include the COCO ID of the object as
@@ -1030,7 +1041,7 @@ def to_keypoints(
             return None
 
         label, attributes = self._get_object_label_and_attributes(
-            classes, supercategory_map, include_id
+            classes_map, supercategory_map, include_id
         )
         attributes.update(self.attributes)
 
@@ -1058,7 +1069,7 @@ def to_keypoints(
     def to_detection(
         self,
         frame_size,
-        classes=None,
+        classes_map=None,
         supercategory_map=None,
         load_segmentation=False,
         include_id=False,
@@ -1068,7 +1079,7 @@ def to_detection(
 
         Args:
             frame_size: the ``(width, height)`` of the image
-            classes (None): the list of classes
+            classes_map (None): a dict mapping class IDs to class labels
             supercategory_map (None): a dict mapping class names to category
                 dicts
             load_segmentation (False): whether to load the segmentation mask
@@ -1084,7 +1095,7 @@ def to_detection(
             return None
 
         label, attributes = self._get_object_label_and_attributes(
-            classes, supercategory_map, include_id
+            classes_map, supercategory_map, include_id
         )
         attributes.update(self.attributes)
 
@@ -1310,10 +1321,10 @@ def _get_label(self, classes):
         return str(self.category_id)
 
     def _get_object_label_and_attributes(
-        self, classes, supercategory_map, include_id
+        self, classes_map, supercategory_map, include_id
     ):
-        if classes:
-            label = classes[self.category_id]
+        if classes_map:
+            label = classes_map[self.category_id]
         else:
             label = str(self.category_id)
 
@@ -1354,7 +1365,7 @@ def load_coco_detection_annotations(json_path, extra_attrs=True):
         a tuple of
 
         -   info: a dict of dataset info
-        -   classes: a list of classes
+        -   classes_map: a dict mapping class IDs to labels
         -   supercategory_map: a dict mapping class labels to category dicts
         -   images: a dict mapping image IDs to image dicts
         -   annotations: a dict mapping image IDs to list of
@@ -1381,9 +1392,9 @@ def _parse_coco_detection_annotations(d, extra_attrs=True):
 
     # Load classes
     if categories is not None:
-        classes, supercategory_map = parse_coco_categories(categories)
+        classes_map, supercategory_map = parse_coco_categories(categories)
     else:
-        classes = None
+        classes_map = None
         supercategory_map = None
 
     # Load image metadata
@@ -1402,17 +1413,14 @@ def _parse_coco_detection_annotations(d, extra_attrs=True):
     else:
         annotations = None
 
-    return info, classes, supercategory_map, images, annotations
+    return info, classes_map, supercategory_map, images, annotations
 
 
 def parse_coco_categories(categories):
     """Parses the COCO categories list.
 
-    The returned ``classes`` contains all class IDs from ``[0, max_id]``,
-    inclusive.
-
     Args:
-        categories: a dict of the form::
+        categories: a list of dict of the form::
 
             [
                 ...
@@ -1429,25 +1437,15 @@ def parse_coco_categories(categories):
     Returns:
         a tuple of
 
-        -   classes: a list of classes
+        -   classes_map: a dict mapping class IDs to labels
         -   supercategory_map: a dict mapping class labels to category dicts
     """
-    cat_map = {c["id"]: c for c in categories}
-
-    classes = []
-    supercategory_map = {}
-    for cat_id in range(max(cat_map, default=-1) + 1):
-        category = cat_map.get(cat_id, None)
-        try:
-            name = category["name"]
-        except:
-            name = str(cat_id)
-
-        classes.append(name)
-        if category is not None:
-            supercategory_map[name] = category
+    classes_map = {
+        c["id"]: c["name"] if "name" in c else str(c["id"]) for c in categories
+    }
+    supercategory_map = {c["name"]: c for c in categories}
 
-    return classes, supercategory_map
+    return classes_map, supercategory_map
 
 
 def download_coco_dataset_split(
@@ -1623,12 +1621,15 @@ def download_coco_dataset_split(
         d = etas.load_json(full_anno_path)
         (
             _,
-            all_classes,
+            all_classes_map,
             _,
             images,
             annotations,
         ) = _parse_coco_detection_annotations(d, extra_attrs=True)
 
+        if all_classes_map is not None:
+            all_classes = list(all_classes_map.values())
+
         if image_ids is not None:
             # Start with specific images
             image_ids = _parse_image_ids(image_ids, images, split=split)
@@ -1717,7 +1718,8 @@ def download_coco_dataset_split(
 
             categories = d.get("categories", None)
             if categories is not None:
-                all_classes, _ = parse_coco_categories(categories)
+                all_classes_map, _ = parse_coco_categories(categories)
+                all_classes = list(all_classes_map.values())
             else:
                 all_classes = None
 
@@ -2002,7 +2004,7 @@ def _get_matching_objects(coco_objects, target_classes, all_classes):
 def _coco_objects_to_polylines(
     coco_objects,
     frame_size,
-    classes,
+    classes_map,
     supercategory_map,
     tolerance,
     include_id,
@@ -2011,7 +2013,7 @@ def _coco_objects_to_polylines(
     for coco_obj in coco_objects:
         polyline = coco_obj.to_polyline(
             frame_size,
-            classes=classes,
+            classes_map=classes_map,
             supercategory_map=supercategory_map,
             tolerance=tolerance,
             include_id=include_id,
@@ -2029,7 +2031,7 @@ def _coco_objects_to_polylines(
 def _coco_objects_to_detections(
     coco_objects,
     frame_size,
-    classes,
+    classes_map,
     supercategory_map,
     load_segmentations,
     include_id,
@@ -2038,7 +2040,7 @@ def _coco_objects_to_detections(
     for coco_obj in coco_objects:
         detection = coco_obj.to_detection(
             frame_size,
-            classes=classes,
+            classes_map=classes_map,
             supercategory_map=supercategory_map,
             load_segmentation=load_segmentations,
             include_id=include_id,
@@ -2058,7 +2060,7 @@ def _coco_objects_to_detections(
 def _coco_objects_to_keypoints(
     coco_objects,
     frame_size,
-    classes,
+    classes_map,
     supercategory_map,
     include_id,
 ):
@@ -2066,7 +2068,7 @@ def _coco_objects_to_keypoints(
     for coco_obj in coco_objects:
         keypoint = coco_obj.to_keypoints(
             frame_size,
-            classes=classes,
+            classes_map=classes_map,
             supercategory_map=supercategory_map,
             include_id=include_id,
         )


Test patch: