Make COCO exports remember original img resolutions (#1142)

* Make COCO exports remember original img resolutions * black formatting
airctic · Sep 29, 2022 · b6b70ad · b6b70ad
1 parent fc19096
commit b6b70ad
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/icevision/data/convert_records_to_coco_style.py b/icevision/data/convert_records_to_coco_style.py
@@ -132,6 +132,14 @@ def export_batch_inferences_as_coco_annotations(
         )
 
     coco_style_preds = convert_preds_to_coco_style(preds)
+    imgs_array = [PIL.Image.open(Path(fname)) for fname in img_files]
+
+    sizes = [{"x": img._size[0], "y": img._size[1]} for img in imgs_array]
+
+    for idx, image in enumerate(coco_style_preds["images"]):
+        coco_style_preds["images"][idx]["width"] = sizes[idx]["x"]
+        coco_style_preds["images"][idx]["height"] = sizes[idx]["y"]
+
     finalized_pseudo_labels = {**addl_info, **coco_style_preds}
 
     # Serialize