Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Add support for exports using existing COCO IDs #4530

Merged
merged 1 commit into from
Jun 25, 2024
Merged

Add support for exports using existing COCO IDs #4530

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
30 changes: 26 additions & 4 deletions fiftyone/utils/coco.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -706,6 +706,8 @@ class COCODetectionDatasetExporter(
- ``True``: export all extra attributes found
- ``False``: do not export extra attributes
- a name or list of names of specific attributes to export
coco_id (None): the name of a sample field containing the COCO IDs of
each image
annotation_id (None): the name of a label field containing the COCO
annotation ID of each label
iscrowd ("iscrowd"): the name of a detection attribute that indicates
Expand All @@ -731,6 +733,7 @@ def __init__(
categories=None,
info=None,
extra_attrs=True,
coco_id=None,
annotation_id=None,
iscrowd="iscrowd",
num_decimals=None,
Expand Down Expand Up @@ -761,12 +764,14 @@ def __init__(
self.categories = categories
self.info = info
self.extra_attrs = extra_attrs
self.coco_id = coco_id
self.annotation_id = annotation_id
self.iscrowd = iscrowd
self.num_decimals = num_decimals
self.tolerance = tolerance

self._image_id = None
self._image_id_map = None
self._anno_id = None
self._images = None
self._annotations = None
Expand Down Expand Up @@ -805,6 +810,11 @@ def log_collection(self, sample_collection):
if self.info is None:
self.info = sample_collection.info

if self.coco_id is not None:
self._image_id_map = dict(
zip(*sample_collection.values(["filepath", self.coco_id]))
)

def export_sample(self, image_or_path, label, metadata=None):
out_image_path, uuid = self._media_exporter.export(image_or_path)

Expand All @@ -816,12 +826,24 @@ def export_sample(self, image_or_path, label, metadata=None):
else:
file_name = uuid

# @todo would be nice to support using existing COCO ID here
self._image_id += 1
if self._image_id_map is not None:
image_id = self._image_id_map.get(image_or_path, None)
if image_id is None:
msg = (
"Ignoring sample with filepath '%s' that has no image ID"
% image_or_path
)
warnings.warn(msg)
return

image_id = int(image_id)
else:
self._image_id += 1
image_id = self._image_id

self._images.append(
{
"id": self._image_id,
"id": image_id,
"file_name": file_name,
"height": metadata.height,
"width": metadata.width,
Expand Down Expand Up @@ -869,7 +891,7 @@ def export_sample(self, image_or_path, label, metadata=None):
obj = COCOObject.from_label(
label,
metadata,
image_id=self._image_id,
image_id=image_id,
category_id=category_id,
extra_attrs=self.extra_attrs,
id_attr=self.annotation_id,
Expand Down
Loading