labelr 0.5.0__py3-none-any.whl → 0.7.0__py3-none-any.whl

labelr/apps/datasets.py

@@ -152,7 +152,10 @@ def export(
     label_studio_url: Optional[str] = LABEL_STUDIO_DEFAULT_URL,
     output_dir: Annotated[
         Optional[Path],
-        typer.Option(help="Path to the output directory", file_okay=False),
+        typer.Option(
+            help="Path to the output directory. Only used if the destintation (`to`) is `ultralytics`",
+            file_okay=False,
+        ),
     ] = None,
     dataset_dir: Annotated[
         Optional[Path],
@@ -200,11 +203,18 @@ def export(
             help="Use the AWS S3 cache for image downloads instead of images.openfoodfacts.org, "
             "it is ignored if the export format is not Ultralytics"
         ),
-    ] = True,
+    ] = False,
     merge_labels: Annotated[
         bool,
         typer.Option(help="Merge multiple labels into a single label"),
     ] = False,
+    revision: Annotated[
+        str,
+        typer.Option(
+            help="Revision (branch, tag or commit) for the Hugging Face Datasets repository. "
+            "Only used when importing from or exporting to Hugging Face Datasets."
+        ),
+    ] = "main",
 ):
     """Export Label Studio annotation, either to Hugging Face Datasets or
     local files (ultralytics format)."""
@@ -261,6 +271,7 @@ def export(
                 project_id=typing.cast(int, project_id),
                 merge_labels=merge_labels,
                 use_aws_cache=use_aws_cache,
+                revision=revision,
             )
         elif to == ExportDestination.ultralytics:
             export_from_ls_to_ultralytics_object_detection(
@@ -286,6 +297,7 @@ def export(
                 download_images=download_images,
                 error_raise=error_raise,
                 use_aws_cache=use_aws_cache,
+                revision=revision,
             )
         else:
             raise typer.BadParameter("Unsupported export format")

labelr/apps/train.py

@@ -6,6 +6,31 @@ from google.cloud import batch_v1
 app = typer.Typer()
 
 
+AVAILABLE_OBJECT_DETECTION_MODELS = [
+    "yolov8n.pt",
+    "yolov8s.pt",
+    "yolov8m.pt",
+    "yolov8l.pt",
+    "yolov8x.pt",
+    "yolov9t.pt",
+    "yolov9s.pt",
+    "yolov9m.pt",
+    "yolov9c.pt",
+    "yolov9e.pt",
+    "yolov10n.pt",
+    "yolov10s.pt",
+    "yolov10m.pt",
+    "yolov10b.pt",
+    "yolov10l.pt",
+    "yolov10x.pt",
+    "yolo11n.pt",
+    "yolo11s.pt",
+    "yolo11m.pt",
+    "yolo11l.pt",
+    "yolo11x.pt",
+]
+
+
 @app.command()
 def train_object_detection(
     wandb_project: str = typer.Option(
@@ -25,9 +50,20 @@ def train_object_detection(
     ),
     epochs: int = typer.Option(100, help="Number of training epochs."),
     imgsz: int = typer.Option(640, help="Size of the image during training."),
-    batch_size: int = typer.Option(64, help="Batch size for training."),
+    batch: int = typer.Option(64, help="Batch size for training."),
+    model_name: str = typer.Option(
+        "yolov8n.pt",
+        help="The YOLO model variant to use for training. "
+        "This should be a valid Ultralytics model name.",
+    ),
 ):
     """Train an object detection model."""
+
+    if model_name not in AVAILABLE_OBJECT_DETECTION_MODELS:
+        raise typer.BadParameter(
+            f"Invalid model name '{model_name}'. Available models are: {', '.join(AVAILABLE_OBJECT_DETECTION_MODELS)}"
+        )
+
     env_variables = {
         "HF_REPO_ID": hf_repo_id,
         "HF_TRAINED_MODEL_REPO_ID": hf_trained_model_repo_id,
@@ -37,8 +73,9 @@ def train_object_detection(
         "WANDB_API_KEY": wandb_api_key,
         "EPOCHS": str(epochs),
         "IMGSZ": str(imgsz),
-        "BATCH_SIZE": str(batch_size),
+        "BATCH_SIZE": str(batch),
         "USE_AWS_IMAGE_CACHE": "False",
+        "YOLO_MODEL_NAME": model_name,
     }
     job_name = "train-yolo-job"
     job_name = job_name + "-" + datetime.datetime.now().strftime("%Y%m%d%H%M%S")
@@ -128,7 +165,7 @@ def launch_job(
     # will run on.
     policy = batch_v1.AllocationPolicy.InstancePolicy()
     # See list of machine types here:
-    # https://docs.cloud.google.com/compute/docs/gpus#l4-gpus
+    # https://docs.cloud.google.com/compute/docs/gpus#t4-gpus
     policy.machine_type = machine_type
 
     accelerator = batch_v1.AllocationPolicy.Accelerator()

labelr/dataset_features.py

@@ -0,0 +1,31 @@
+import datasets
+from datasets import Features
+from datasets import Image as HFImage
+
+# dataset features for predictions.parquet used in evaluation and visualization
+OBJECT_DETECTION_DS_PREDICTION_FEATURES = Features(
+    {
+        "image": HFImage(),
+        "image_with_prediction": HFImage(),
+        "image_id": datasets.Value("string"),
+        "detected": {
+            "bbox": datasets.Sequence(datasets.Sequence(datasets.Value("float32"))),
+            "category_id": datasets.Sequence(datasets.Value("int64")),
+            "category_name": datasets.Sequence(datasets.Value("string")),
+            "confidence": datasets.Sequence(datasets.Value("float32")),
+        },
+        "split": datasets.Value("string"),
+        "width": datasets.Value("int64"),
+        "height": datasets.Value("int64"),
+        "meta": {
+            "barcode": datasets.Value("string"),
+            "off_image_id": datasets.Value("string"),
+            "image_url": datasets.Value("string"),
+        },
+        "objects": {
+            "bbox": datasets.Sequence(datasets.Sequence(datasets.Value("float32"))),
+            "category_id": datasets.Sequence(datasets.Value("int64")),
+            "category_name": datasets.Sequence(datasets.Value("string")),
+        },
+    }
+)

labelr/export.py

@@ -36,11 +36,18 @@ def export_from_ls_to_hf_object_detection(
     project_id: int,
     merge_labels: bool = False,
     use_aws_cache: bool = True,
+    revision: str = "main",
 ):
     if merge_labels:
         label_names = ["object"]
 
-    logger.info("Project ID: %d, label names: %s", project_id, label_names)
+    logger.info(
+        "Project ID: %d, label names: %s, repo_id: %s, revision: %s",
+        project_id,
+        label_names,
+        repo_id,
+        revision,
+    )
 
     for split in ["train", "val"]:
         logger.info("Processing split: %s", split)
@@ -70,7 +77,13 @@ def export_from_ls_to_hf_object_detection(
                 functools.partial(_pickle_sample_generator, tmp_dir),
                 features=HF_DS_OBJECT_DETECTION_FEATURES,
             )
-            hf_ds.push_to_hub(repo_id, split=split)
+            hf_ds.push_to_hub(
+                repo_id,
+                split=split,
+                revision=revision,
+                # Create a PR if not pushing to main branch
+                create_pr=revision != "main",
+            )
 
 
 def export_from_ls_to_ultralytics_object_detection(
@@ -234,19 +247,28 @@ def export_from_hf_to_ultralytics_object_detection(
             is True. Defaults to True.
         revision (str): The dataset revision to load. Defaults to 'main'.
     """
-    logger.info("Repo ID: %s", repo_id)
+    logger.info("Repo ID: %s, revision: %s", repo_id, revision)
     ds = datasets.load_dataset(repo_id, revision=revision)
     data_dir = output_dir / "data"
     data_dir.mkdir(parents=True, exist_ok=True)
     category_id_to_name = {}
 
+    split_map = {
+        "train": "train",
+        "val": "val",
+    }
+    if "val" not in ds and "test" in ds:
+        logger.info("val split not found, using test split instead as val")
+        split_map["val"] = "test"
+
     for split in ["train", "val"]:
+        split_target = split_map[split]
         split_labels_dir = data_dir / "labels" / split
         split_labels_dir.mkdir(parents=True, exist_ok=True)
         split_images_dir = data_dir / "images" / split
         split_images_dir.mkdir(parents=True, exist_ok=True)
 
-        for sample in tqdm.tqdm(ds[split], desc="samples"):
+        for sample in tqdm.tqdm(ds[split_target], desc="samples"):
             image_id = sample["image_id"]
 
             if download_images:

labelr/utils.py

@@ -0,0 +1,13 @@
+def parse_hf_repo_id(hf_repo_id: str) -> tuple[str, str]:
+    """Parse the repo_id and the revision from a hf_repo_id in the format:
+    `org/repo-name@revision`.
+
+    Returns a tuple (repo_id, revision), with revision = 'main' if it
+    was not provided.
+    """
+    if "@" in hf_repo_id:
+        hf_repo_id, revision = hf_repo_id.split("@", 1)
+    else:
+        revision = "main"
+
+    return hf_repo_id, revision

{labelr-0.5.0.dist-info → labelr-0.7.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: labelr
-Version: 0.5.0
+Version: 0.7.0
 Summary: A command-line tool to manage labeling tasks with Label Studio.
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
@@ -12,8 +12,11 @@ Requires-Dist: more-itertools>=10.5.0
 Requires-Dist: openfoodfacts>=2.9.0
 Requires-Dist: typer>=0.15.1
 Requires-Dist: google-cloud-batch==0.18.0
+Requires-Dist: huggingface-hub
 Provides-Extra: ultralytics
-Requires-Dist: ultralytics>=8.3.49; extra == "ultralytics"
+Requires-Dist: ultralytics==8.3.223; extra == "ultralytics"
+Provides-Extra: fiftyone
+Requires-Dist: fiftyone~=1.10.0; extra == "fiftyone"
 Dynamic: license-file
 
 # Labelr

{labelr-0.5.0.dist-info → labelr-0.7.0.dist-info}/RECORD

@@ -3,19 +3,21 @@ labelr/__main__.py,sha256=G4e95-IfhI-lOmkOBP6kQ8wl1x_Fl7dZlLOYr90K83c,66
 labelr/annotate.py,sha256=3fJ9FYbcozcOoKuhNtzPHV8sSnp-45FsNnMc8UeBHGU,3503
 labelr/check.py,sha256=3wK6mE0UsKvoBNm0_lyWhCMq7gxkv5r50pvO70damXY,2476
 labelr/config.py,sha256=3RXF_NdkSuHvfVMGMlYmjlw45fU77zQkLX7gmZq7NxM,64
-labelr/export.py,sha256=h4_n-twXHnrd5FCL9NEeEFjWw6Fzo5b9JayXHcDLAF0,17336
+labelr/dataset_features.py,sha256=ZC9QAUw9oKHqyUPla2h3xQFaRT9sHq8hkPNN4RDDwmo,1257
+labelr/export.py,sha256=gjC2_RJ_yX8zVYXyo1RAgI07iXSgkeqckOTEzSscRXc,17940
 labelr/main.py,sha256=CioMPtaPoGL_5Oxwj8PfalhTyFahMbfp2kd9KdZzm3Y,2258
 labelr/project_config.py,sha256=CIHEcgSOfXb53naHWEBkTDm2V9m3abAu8C54VSzHjAs,1260
 labelr/sample.py,sha256=unu9AQ64FhKPgssuL7gb3qyMd1EQJvMOfqvjdefmWOU,7807
 labelr/types.py,sha256=8CHfLyifF_N94OYDhG-7IcWboOh9o0Z_0LBtQapT8TQ,313
+labelr/utils.py,sha256=e0R15jePWBzRdN8LB6kBSH5Dl_P0MNEtRmeqB9eu5d8,415
 labelr/apps/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-labelr/apps/datasets.py,sha256=twwH8wzbUauDWVZFObOvk0gsohAeYm0usHCigng_ucM,11262
+labelr/apps/datasets.py,sha256=4PMfKS5c7Zw3-NNRBkFbZidMQUI2RBMcXFYBvWHLz3o,11688
 labelr/apps/projects.py,sha256=HpgqIaPrUQzIR7eOLn4EBbEzXRi7hoWStT4jLMQPcBg,15153
-labelr/apps/train.py,sha256=SdbCOPWxBkpnQ7P93flhIcnZIgfEX0Na0ce0RM0M91U,6023
+labelr/apps/train.py,sha256=sI0p3h39LPXhynwl_yMuZnIPlaqlcWSO_81zPC3H3yI,6886
 labelr/apps/users.py,sha256=twQSlpHxE0hrYkgrJpEFbK8lYfWnpJr8vyfLHLtdAUU,909
-labelr-0.5.0.dist-info/licenses/LICENSE,sha256=hIahDEOTzuHCU5J2nd07LWwkLW7Hko4UFO__ffsvB-8,34523
-labelr-0.5.0.dist-info/METADATA,sha256=kK-wEIUdt7LbkLCZ7jrPj88znftFnw6DO_bAR9l8ZDY,6881
-labelr-0.5.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-labelr-0.5.0.dist-info/entry_points.txt,sha256=OACukVeR_2z54i8yQuWqqk_jdEHlyTwmTFOFBmxPp1k,43
-labelr-0.5.0.dist-info/top_level.txt,sha256=bjZo50aGZhXIcZYpYOX4sdAQcamxh8nwfEh7A9RD_Ag,7
-labelr-0.5.0.dist-info/RECORD,,
+labelr-0.7.0.dist-info/licenses/LICENSE,sha256=hIahDEOTzuHCU5J2nd07LWwkLW7Hko4UFO__ffsvB-8,34523
+labelr-0.7.0.dist-info/METADATA,sha256=NghQ_6mNj1Dkets_GlOOOyoAVEQqoPBbbJXhysOKAWI,6991
+labelr-0.7.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+labelr-0.7.0.dist-info/entry_points.txt,sha256=OACukVeR_2z54i8yQuWqqk_jdEHlyTwmTFOFBmxPp1k,43
+labelr-0.7.0.dist-info/top_level.txt,sha256=bjZo50aGZhXIcZYpYOX4sdAQcamxh8nwfEh7A9RD_Ag,7
+labelr-0.7.0.dist-info/RECORD,,