labelr 0.1.0__py3-none-any.whl

labelr/sample.py

@@ -0,0 +1,186 @@
+import logging
+import random
+import string
+
+import datasets
+from openfoodfacts.images import download_image
+
+logger = logging.getLogger(__name__)
+
+
+def format_annotation_results_from_hf(
+    objects: dict, image_width: int, image_height: int
+):
+    """Format annotation results from a HF object detection dataset into Label
+    Studio format."""
+    annotation_results = []
+    for i in range(len(objects["bbox"])):
+        bbox = objects["bbox"][i]
+        # category_id = objects["category_id"][i]
+        category_name = objects["category_name"][i]
+        # These are relative coordinates (between 0.0 and 1.0)
+        y_min, x_min, y_max, x_max = bbox
+        # Make sure the coordinates are within the image boundaries,
+        # and convert them to percentages
+        y_min = min(max(0, y_min), 1.0) * 100
+        x_min = min(max(0, x_min), 1.0) * 100
+        y_max = min(max(0, y_max), 1.0) * 100
+        x_max = min(max(0, x_max), 1.0) * 100
+        x = x_min
+        y = y_min
+        width = x_max - x_min
+        height = y_max - y_min
+
+        id_ = "".join(random.choices(string.ascii_letters + string.digits, k=10))
+        annotation_results.append(
+            {
+                "id": id_,
+                "type": "rectanglelabels",
+                "from_name": "label",
+                "to_name": "image",
+                "original_width": image_width,
+                "original_height": image_height,
+                "image_rotation": 0,
+                "value": {
+                    "rotation": 0,
+                    "x": x,
+                    "y": y,
+                    "width": width,
+                    "height": height,
+                    "rectanglelabels": [category_name],
+                },
+            },
+        )
+    return annotation_results
+
+
+def format_object_detection_sample_from_hf(hf_sample: dict, split: str) -> dict:
+    hf_meta = hf_sample["meta"]
+    objects = hf_sample["objects"]
+    image_width = hf_sample["width"]
+    image_height = hf_sample["height"]
+    annotation_results = format_annotation_results_from_hf(
+        objects, image_width, image_height
+    )
+    return {
+        "data": {
+            "image_id": hf_sample["image_id"],
+            "image_url": hf_meta["image_url"],
+            "batch": "null",
+            "split": split,
+            "meta": {
+                "width": image_width,
+                "height": image_height,
+                "barcode": hf_meta["barcode"],
+                "off_image_id": hf_meta["off_image_id"],
+            },
+        },
+        "predictions": [{"result": annotation_results}],
+    }
+
+
+def format_object_detection_sample_to_ls(
+    image_id: str,
+    image_url: str,
+    width: int,
+    height: int,
+    extra_meta: dict | None = None,
+) -> dict:
+    """Format an object detection sample in Label Studio format.
+
+    Args:
+        image_id: The image ID.
+        image_url: The URL of the image.
+        width: The width of the image.
+        height: The height of the image.
+        extra_meta: Extra metadata to include in the sample.
+    """
+    extra_meta = extra_meta or {}
+    return {
+        "data": {
+            "image_id": image_id,
+            "image_url": image_url,
+            "batch": "null",
+            "meta": {
+                "width": width,
+                "height": height,
+                **extra_meta,
+            },
+        },
+    }
+
+
+def format_object_detection_sample_to_hf(
+    task_data: dict, annotations: list[dict], category_names: list[str]
+) -> dict | None:
+    if len(annotations) > 1:
+        logger.info("More than one annotation found, skipping")
+        return None
+    elif len(annotations) == 0:
+        logger.info("No annotation found, skipping")
+        return None
+
+    annotation = annotations[0]
+    bboxes = []
+    bbox_category_ids = []
+    bbox_category_names = []
+
+    for annotation_result in annotation["result"]:
+        if annotation_result["type"] != "rectanglelabels":
+            raise ValueError("Invalid annotation type: %s" % annotation_result["type"])
+
+        value = annotation_result["value"]
+        x_min = value["x"] / 100
+        y_min = value["y"] / 100
+        width = value["width"] / 100
+        height = value["height"] / 100
+        x_max = x_min + width
+        y_max = y_min + height
+        bboxes.append([y_min, x_min, y_max, x_max])
+        category_name = value["rectanglelabels"][0]
+        bbox_category_names.append(category_name)
+        bbox_category_ids.append(category_names.index(category_name))
+
+    image_url = task_data["image_url"]
+    image = download_image(image_url, error_raise=False)
+    if image is None:
+        logger.error("Failed to download image: %s", image_url)
+        return None
+
+    return {
+        "image_id": task_data["image_id"],
+        "image": image,
+        "width": task_data["meta"]["width"],
+        "height": task_data["meta"]["height"],
+        "meta": {
+            "barcode": task_data["meta"]["barcode"],
+            "off_image_id": task_data["meta"]["off_image_id"],
+            "image_url": image_url,
+        },
+        "objects": {
+            "bbox": bboxes,
+            "category_id": bbox_category_ids,
+            "category_name": bbox_category_names,
+        },
+    }
+
+
+# The HuggingFace Dataset features
+HF_DS_FEATURES = datasets.Features(
+    {
+        "image_id": datasets.Value("string"),
+        "image": datasets.features.Image(),
+        "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/triton/object_detection.py

@@ -0,0 +1,241 @@
+import dataclasses
+import functools
+import logging
+import time
+from typing import Any, Optional
+
+import numpy as np
+from PIL import Image
+
+try:
+    import grpc
+    from tritonclient.grpc import service_pb2, service_pb2_grpc
+    from tritonclient.grpc.service_pb2_grpc import GRPCInferenceServiceStub
+except ImportError:
+    pass
+
+logger = logging.getLogger(__name__)
+
+
+JSONType = dict[str, Any]
+
+OBJECT_DETECTION_MODEL_VERSION = {
+    "nutriscore": "tf-nutriscore-1.0",
+    "nutrition_table": "tf-nutrition-table-1.0",
+    "universal_logo_detector": "tf-universal-logo-detector-1.0",
+}
+
+LABELS = {
+    "nutriscore": [
+        "NULL",
+        "nutriscore-a",
+        "nutriscore-b",
+        "nutriscore-c",
+        "nutriscore-d",
+        "nutriscore-e",
+    ],
+}
+
+OBJECT_DETECTION_IMAGE_MAX_SIZE = (1024, 1024)
+
+
+@functools.cache
+def get_triton_inference_stub(
+    triton_uri: str,
+) -> "GRPCInferenceServiceStub":
+    """Return a gRPC stub for Triton Inference Server.
+
+    :param triton_uri: URI of the Triton Inference Server
+    :return: gRPC stub for Triton Inference Server
+    """
+    triton_uri = triton_uri
+    channel = grpc.insecure_channel(triton_uri)
+    return service_pb2_grpc.GRPCInferenceServiceStub(channel)
+
+
+def convert_image_to_array(image: Image.Image) -> np.ndarray:
+    """Convert a PIL Image into a numpy array.
+
+    The image is converted to RGB if needed before generating the array.
+
+    :param image: the input image
+    :return: the generated numpy array of shape (width, height, 3)
+    """
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+
+    (im_width, im_height) = image.size
+
+    return np.array(image.getdata()).reshape((im_height, im_width, 3)).astype(np.uint8)
+
+
+@dataclasses.dataclass
+class ObjectDetectionResult:
+    bounding_box: tuple[int, int, int, int]
+    score: float
+    label: str
+
+
+@dataclasses.dataclass
+class ObjectDetectionRawResult:
+    num_detections: int
+    detection_boxes: np.ndarray
+    detection_scores: np.ndarray
+    detection_classes: np.ndarray
+    label_names: list[str]
+    detection_masks: Optional[np.ndarray] = None
+    boxed_image: Optional[Image.Image] = None
+
+    def select(self, threshold: Optional[float] = None) -> list[ObjectDetectionResult]:
+        if threshold is None:
+            threshold = 0.5
+
+        box_masks = self.detection_scores > threshold
+        selected_boxes = self.detection_boxes[box_masks]
+        selected_scores = self.detection_scores[box_masks]
+        selected_classes = self.detection_classes[box_masks]
+
+        results = []
+        for bounding_box, score, label in zip(
+            selected_boxes, selected_scores, selected_classes
+        ):
+            label_int = int(label)
+            label_str = self.label_names[label_int]
+            if label_str is not None:
+                result = ObjectDetectionResult(
+                    bounding_box=tuple(bounding_box.tolist()),  # type: ignore
+                    score=float(score),
+                    label=label_str,
+                )
+                results.append(result)
+
+        return results
+
+    def to_json(self, threshold: Optional[float] = None) -> list[JSONType]:
+        return [dataclasses.asdict(r) for r in self.select(threshold)]
+
+
+def resize_image(image: Image.Image, max_size: tuple[int, int]) -> Image.Image:
+    width, height = image.size
+    max_width, max_height = max_size
+
+    if width > max_width or height > max_height:
+        new_image = image.copy()
+        new_image.thumbnail((max_width, max_height))
+        return new_image
+
+    return image
+
+
+class RemoteModel:
+    def __init__(self, name: str, label_names: list[str]):
+        self.name: str = name
+        self.label_names = label_names
+
+    def detect_from_image(
+        self,
+        image: Image.Image,
+        triton_uri: str,
+    ) -> ObjectDetectionRawResult:
+        """Run object detection model on an image.
+
+        :param image: the input Pillow image
+        :param triton_uri: URI of the Triton Inference Server.
+        :return: the detection result
+        """
+        resized_image = resize_image(image, OBJECT_DETECTION_IMAGE_MAX_SIZE)
+        image_array = convert_image_to_array(resized_image)
+        grpc_stub = get_triton_inference_stub(triton_uri)
+        request = service_pb2.ModelInferRequest()
+        request.model_name = self.name
+
+        image_input = service_pb2.ModelInferRequest().InferInputTensor()
+        image_input.name = "inputs"
+        image_input.datatype = "UINT8"
+        image_input.shape.extend([1, image_array.shape[0], image_array.shape[1], 3])
+        request.inputs.extend([image_input])
+
+        for output_name in (
+            "num_detections",
+            "detection_classes",
+            "detection_scores",
+            "detection_boxes",
+        ):
+            output = service_pb2.ModelInferRequest().InferRequestedOutputTensor()
+            output.name = output_name
+            request.outputs.extend([output])
+
+        request.raw_input_contents.extend([image_array.tobytes()])
+        start_time = time.monotonic()
+        response = grpc_stub.ModelInfer(request)
+        logger.debug(
+            "Inference time for %s: %s", self.name, time.monotonic() - start_time
+        )
+
+        if len(response.outputs) != 4:
+            raise Exception(f"expected 4 output, got {len(response.outputs)}")
+
+        if len(response.raw_output_contents) != 4:
+            raise Exception(
+                f"expected 4 raw output content, got {len(response.raw_output_contents)}"
+            )
+
+        output_index = {output.name: i for i, output in enumerate(response.outputs)}
+        num_detections = (
+            np.frombuffer(
+                response.raw_output_contents[output_index["num_detections"]],
+                dtype=np.float32,
+            )
+            .reshape((1, 1))
+            .astype(int)[0][0]  # type: ignore
+        )
+        detection_scores = np.frombuffer(
+            response.raw_output_contents[output_index["detection_scores"]],
+            dtype=np.float32,
+        ).reshape((1, -1))[0]
+        detection_classes = (
+            np.frombuffer(
+                response.raw_output_contents[output_index["detection_classes"]],
+                dtype=np.float32,
+            )
+            .reshape((1, -1))
+            .astype(int)  # type: ignore
+        )[0]
+        detection_boxes = np.frombuffer(
+            response.raw_output_contents[output_index["detection_boxes"]],
+            dtype=np.float32,
+        ).reshape((1, -1, 4))[0]
+
+        result = ObjectDetectionRawResult(
+            num_detections=num_detections,
+            detection_classes=detection_classes,
+            detection_boxes=detection_boxes,
+            detection_scores=detection_scores,
+            detection_masks=None,
+            label_names=self.label_names,
+        )
+
+        return result
+
+
+class ObjectDetectionModelRegistry:
+    models: dict[str, RemoteModel] = {}
+    _loaded = False
+
+    @classmethod
+    def get_available_models(cls) -> list[str]:
+        cls.load_all()
+        return list(cls.models.keys())
+
+    @classmethod
+    def load(cls, name: str) -> RemoteModel:
+        label_names = LABELS[name]
+        model = RemoteModel(name, label_names)
+        cls.models[name] = model
+        return model
+
+    @classmethod
+    def get(cls, name: str) -> RemoteModel:
+        if name not in cls.models:
+            cls.load(name)
+        return cls.models[name]

labelr/types.py

@@ -0,0 +1,16 @@
+import enum
+
+
+class ExportSource(str, enum.Enum):
+    hf = "hf"
+    ls = "ls"
+
+
+class ExportDestination(str, enum.Enum):
+    hf = "hf"
+    ultralytics = "ultralytics"
+
+
+class TaskType(str, enum.Enum):
+    object_detection = "object_detection"
+    classification = "classification"