diff --git a/plugins/kitti_export_object_detection/pyproject.toml b/plugins/kitti_export_object_detection/pyproject.toml index 67cf4e8..ba2caaa 100644 --- a/plugins/kitti_export_object_detection/pyproject.toml +++ b/plugins/kitti_export_object_detection/pyproject.toml @@ -4,7 +4,7 @@ version = "0.1.0" description = "KITTI object detection export plugin for Lightly Studio" requires-python = ">=3.9" dependencies = [ - "lightly_studio>=1.0.0", + "lightly_studio>=1.0.4", "labelformat", "sqlmodel", ] diff --git a/plugins/kitti_export_object_detection/src/lightly_plugins_kitti_export_object_detection/operator.py b/plugins/kitti_export_object_detection/src/lightly_plugins_kitti_export_object_detection/operator.py index d01c2ea..77deb11 100644 --- a/plugins/kitti_export_object_detection/src/lightly_plugins_kitti_export_object_detection/operator.py +++ b/plugins/kitti_export_object_detection/src/lightly_plugins_kitti_export_object_detection/operator.py @@ -15,6 +15,7 @@ from sqlmodel import Session from lightly_studio.core.image.image_sample import ImageSample +from lightly_studio.export.image_dataset_export import image_sample_to_image from lightly_studio.export.lightly_studio_label_input import ( LightlyStudioObjectDetectionInput, ) @@ -51,7 +52,13 @@ def __init__( images_root: Common root path used to preserve nested image folders. """ self._images_root = images_root - super().__init__(session=session, dataset_id=dataset_id, samples=samples) + super().__init__( + session=session, + dataset_id=dataset_id, + samples=samples, + annotation_collection_id=None, + sample_to_image=image_sample_to_image, + ) def get_images(self) -> list[Image]: """Return images with filenames relative to the KITTI output folder.""" diff --git a/plugins/yolo_object_detection/src/lightly_plugins_yolo_object_detection/operator.py b/plugins/yolo_object_detection/src/lightly_plugins_yolo_object_detection/operator.py index e2a532f..8c8bb91 100644 --- a/plugins/yolo_object_detection/src/lightly_plugins_yolo_object_detection/operator.py +++ b/plugins/yolo_object_detection/src/lightly_plugins_yolo_object_detection/operator.py @@ -9,6 +9,7 @@ from sqlmodel import Session from ultralytics import YOLO # type: ignore[attr-defined] +from ultralytics.engine.results import Results from lightly_studio.models.annotation.annotation_base import ( AnnotationCreate, @@ -140,8 +141,8 @@ def execute( total_annotations_created = 0 for i, image_entry in enumerate(samples, start=1): try: - results = model( - image_entry.file_path_abs, conf=confidence, verbose=False + result = list( + model(image_entry.file_path_abs, conf=confidence, verbose=False) )[0] except Exception as e: logger.error( @@ -153,12 +154,22 @@ def execute( success=False, message=f"Failed to run inference on '{image_entry.file_path_abs}': {e}", ) - for box in results.boxes: - category_id = int(box.cls) + # A single image always yields one `Results`; `embed=` is never passed. + assert isinstance(result, Results) + boxes = result.boxes + if boxes is None: + logger.warning( + "No boxes returned for '%s'; is '%s' a detection model?", + image_entry.file_path_abs, + model_path, + ) + continue + for box_index in range(len(boxes)): + category_id = int(boxes.cls[box_index]) label_id = label_map.get(category_id) if label_id is None: continue - x_center, y_center, w, h = box.xywh[0].tolist() + x_center, y_center, w, h = boxes.xywh[box_index].tolist() annotations_to_create.append( AnnotationCreate( annotation_label_id=label_id, @@ -168,7 +179,7 @@ def execute( y=round(y_center - h / 2), width=max(1, round(w)), height=max(1, round(h)), - confidence=float(box.conf), + confidence=float(boxes.conf[box_index]), ) )