Make Style

Author: yeldarby

inference/core/workflows/core_steps/visualizations/icon/v1.py

@@ -12,8 +12,8 @@
 from inference.core.workflows.execution_engine.entities.base import WorkflowImageData
 from inference.core.workflows.execution_engine.entities.types import (
     IMAGE_KIND,
-    INTEGER_KIND,
     INSTANCE_SEGMENTATION_PREDICTION_KIND,
+    INTEGER_KIND,
     KEYPOINT_DETECTION_PREDICTION_KIND,
     OBJECT_DETECTION_PREDICTION_KIND,
     STRING_KIND,
@@ -22,16 +22,15 @@
 from inference.core.workflows.prototypes.block import BlockResult, WorkflowBlockManifest
 
 TYPE: str = "roboflow_core/icon_visualization@v1"
-SHORT_DESCRIPTION = (
-    "Draw icons on an image either at specific static coordinates or dynamically based on detections."
-)
+SHORT_DESCRIPTION = "Draw icons on an image either at specific static coordinates or dynamically based on detections."
 LONG_DESCRIPTION = """
 The `IconVisualization` block draws icons on an image using Supervision's `sv.IconAnnotator`.
 It supports two modes:
 1. **Static Mode**: Position an icon at a fixed location (e.g., for watermarks)
 2. **Dynamic Mode**: Position icons based on detection coordinates
 """
 
+
 class IconManifest(VisualizationManifest):
     type: Literal[f"{TYPE}", "IconVisualization"]
     model_config = ConfigDict(
@@ -173,17 +172,14 @@ class IconManifest(VisualizationManifest):
     def validate_mode_parameters(self) -> "IconManifest":
         if self.mode == "dynamic":
             if self.predictions is None:
-                raise ValueError(
-                    "The 'predictions' field is required for dynamic mode"
-                )
+                raise ValueError("The 'predictions' field is required for dynamic mode")
         return self
 
     @classmethod
     def get_execution_engine_compatibility(cls) -> Optional[str]:
         return ">=1.3.0,<2.0.0"
 
 
-
 class IconVisualizationBlockV1(VisualizationBlock):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
@@ -224,28 +220,39 @@ def run(
     ) -> BlockResult:
         annotated_image = image.numpy_image.copy() if copy_image else image.numpy_image
         icon_np = icon.numpy_image.copy()
-        
+
+        import os
         import tempfile
+
         import cv2
-        import os
-        
+
         # WorkflowImageData loses alpha channels when loading images.
         # Try to recover them from the original source.
         if icon_np.shape[2] == 3:
             # Try reloading from file with IMREAD_UNCHANGED
-            if hasattr(icon, '_image_reference') and icon._image_reference and \
-               not icon._image_reference.startswith('http'):
+            if (
+                hasattr(icon, "_image_reference")
+                and icon._image_reference
+                and not icon._image_reference.startswith("http")
+            ):
                 try:
-                    icon_with_alpha = cv2.imread(icon._image_reference, cv2.IMREAD_UNCHANGED)
+                    icon_with_alpha = cv2.imread(
+                        icon._image_reference, cv2.IMREAD_UNCHANGED
+                    )
                     if icon_with_alpha is not None and icon_with_alpha.shape[2] == 4:
                         icon_np = icon_with_alpha
                 except:
                     pass
-            
+
             # Try decoding base64 with alpha preserved
-            if icon_np.shape[2] == 3 and hasattr(icon, '_base64_image') and icon._base64_image:
+            if (
+                icon_np.shape[2] == 3
+                and hasattr(icon, "_base64_image")
+                and icon._base64_image
+            ):
                 try:
                     import base64
+
                     image_bytes = base64.b64decode(icon._base64_image)
                     nparr = np.frombuffer(image_bytes, np.uint8)
                     decoded = cv2.imdecode(nparr, cv2.IMREAD_UNCHANGED)
@@ -256,30 +263,40 @@ def run(
                             icon_np = decoded
                 except:
                     pass
-        
+
         with tempfile.NamedTemporaryFile(suffix=".png", delete=False) as f:
             # Ensure proper format for IconAnnotator
             if len(icon_np.shape) == 2:
                 icon_np = cv2.cvtColor(icon_np, cv2.COLOR_GRAY2BGR)
-                alpha = np.ones((icon_np.shape[0], icon_np.shape[1], 1), dtype=icon_np.dtype) * 255
+                alpha = (
+                    np.ones(
+                        (icon_np.shape[0], icon_np.shape[1], 1), dtype=icon_np.dtype
+                    )
+                    * 255
+                )
                 icon_np = np.concatenate([icon_np, alpha], axis=2)
             elif icon_np.shape[2] == 3:
-                alpha = np.ones((icon_np.shape[0], icon_np.shape[1], 1), dtype=icon_np.dtype) * 255
+                alpha = (
+                    np.ones(
+                        (icon_np.shape[0], icon_np.shape[1], 1), dtype=icon_np.dtype
+                    )
+                    * 255
+                )
                 icon_np = np.concatenate([icon_np, alpha], axis=2)
-            
+
             cv2.imwrite(f.name, icon_np)
             icon_path = f.name
-        
+
         try:
             if mode == "static":
                 img_height, img_width = annotated_image.shape[:2]
-                
+
                 # Handle negative positioning (from right/bottom edges)
                 if x_position < 0:
                     actual_x = img_width + x_position - icon_width
                 else:
                     actual_x = x_position
-                    
+
                 if y_position < 0:
                     actual_y = img_height + y_position - icon_height
                 else:
@@ -288,48 +305,45 @@ def run(
                 # IconAnnotator expects a detection, so create one at the desired position
                 center_x = actual_x + icon_width // 2
                 center_y = actual_y + icon_height // 2
-                
+
                 static_detections = sv.Detections(
-                    xyxy=np.array([[
-                        center_x - 1, 
-                        center_y - 1, 
-                        center_x + 1, 
-                        center_y + 1
-                    ]], dtype=np.float64),
+                    xyxy=np.array(
+                        [[center_x - 1, center_y - 1, center_x + 1, center_y + 1]],
+                        dtype=np.float64,
+                    ),
                     class_id=np.array([0]),
                     confidence=np.array([1.0]),
                 )
-                
+
                 annotator = sv.IconAnnotator(
                     icon_resolution_wh=(icon_width, icon_height),
                     icon_position=sv.Position.CENTER,
                 )
-                
+
                 annotated_image = annotator.annotate(
                     scene=annotated_image,
                     detections=static_detections,
-                    icon_path=icon_path
+                    icon_path=icon_path,
                 )
-                
+
             elif mode == "dynamic" and predictions is not None and len(predictions) > 0:
                 annotator = self.getAnnotator(
                     icon_width=icon_width,
                     icon_height=icon_height,
                     position=position,
                 )
-                
+
                 if annotator is not None:
                     annotated_image = annotator.annotate(
                         scene=annotated_image,
                         detections=predictions,
-                        icon_path=icon_path
+                        icon_path=icon_path,
                     )
         finally:
             os.unlink(icon_path)
-        
+
         return {
             OUTPUT_IMAGE_KEY: WorkflowImageData.copy_and_replace(
-                origin_image_data=image,
-                numpy_image=annotated_image
+                origin_image_data=image, numpy_image=annotated_image
             )
         }