feat: ✨ new annotators added from supervision

Signed-off-by: Onuralp SEZER <[email protected]>

.gitignore

@@ -158,3 +158,12 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+# Mac
+.DS_Store
+.AppleDouble
+.LSOverride
+
+# YoloV8 files
+yolov8s-seg.pt
+yolov8s.pt

app.py

@@ -1,35 +1,57 @@
+import os  # added for cache_examples
+
 import gradio as gr
+import numpy as np
 import supervision as sv
+from PIL import Image
+from torch import cuda, device
 from ultralytics import YOLO
-import os #added for cache_examples
-from PIL import Image, ImageColor
-import numpy as np
 
-def load_model(img):
-  # Load model, get results and return detections/labels
-  model = YOLO('yolov8s-seg.pt')
-  result = model(img, verbose=False, imgsz=1280)[0]
-  detections = sv.Detections.from_ultralytics(result)
-  labels = [
-            f"{model.model.names[class_id]} {confidence:.2f}"
-            for class_id, confidence in zip(detections.class_id, detections.confidence)
-           ]
-
-  return detections, labels
-
-def calculate_crop_dim(a,b):
-  #Calculates the crop dimensions of the image resultant
-  if a>b:
-    width= a
-    height = a
-  else:
-    width = b
-    height = b
-
-  return width, height
-
-def annotator(img,annotators,colorbb,colormask,colorellipse,colorbc,colorcir,colorlabel):
+# Use GPU if available
+if cuda.is_available():
+    device = device("cuda")
+else:
+    device = device("cpu")
 
+
+def load_model(img):
+    # Load model, get results and return detections/labels
+    model = YOLO("yolov8s-seg.pt")
+    result = model(img, verbose=False, imgsz=1280)[0]
+    detections = sv.Detections.from_ultralytics(result)
+    labels = [
+        f"{model.model.names[class_id]} {confidence:.2f}"
+        for class_id, confidence in zip(detections.class_id, detections.confidence)
+    ]
+
+    print(labels)
+    return detections, labels
+
+
+def calculate_crop_dim(a, b):
+    # Calculates the crop dimensions of the image resultant
+    if a > b:
+        width = a
+        height = a
+    else:
+        width = b
+        height = b
+
+    return width, height
+
+
+def annotator(
+    img,
+    annotators,
+    colorbb,
+    colormask,
+    colorellipse,
+    colorbc,
+    colorcir,
+    colorlabel,
+    colorhalo,
+    colortri
+):
     """
     Function that changes the color of annotators
     Args:
@@ -40,14 +62,10 @@ def annotator(img,annotators,colorbb,colormask,colorellipse,colorbc,colorcir,col
         annotators: annotated image
     """
 
-
-
-    img = img[...,::-1].copy() # BGR to RGB using numpy
-
+    img = img[..., ::-1].copy()  # BGR to RGB using numpy
 
     detections, labels = load_model(img)
 
-
     if "Blur" in annotators:
         # Apply Blur
         blur_annotator = sv.BlurAnnotator()
@@ -73,8 +91,7 @@ def annotator(img,annotators,colorbb,colormask,colorellipse,colorbc,colorcir,col
         corner_annotator = sv.BoxCornerAnnotator(sv.Color.from_hex(str(colorbc)))
         img = corner_annotator.annotate(img, detections=detections)
 
-
-    if "Circle" in annotators: 
+    if "Circle" in annotators:
         # Draw circle
         circle_annotator = sv.CircleAnnotator(sv.Color.from_hex(str(colorcir)))
         img = circle_annotator.annotate(img, detections=detections)
@@ -85,71 +102,133 @@ def annotator(img,annotators,colorbb,colormask,colorellipse,colorbc,colorcir,col
         label_annotator = sv.LabelAnnotator(sv.Color.from_hex(str(colorlabel)))
         img = label_annotator.annotate(img, detections=detections, labels=labels)
 
+    if "Pixelate" in annotators:
+        # Draw PixelateAnnotator
+        pixelate_annotator = sv.PixelateAnnotator()
+        img = pixelate_annotator.annotate(img, detections=detections)
 
+    if "Halo" in annotators:
+        # Draw HaloAnnotator
+        halo_annotator = sv.HaloAnnotator(sv.Color.from_hex(str(colorhalo)))
+        img = halo_annotator.annotate(img, detections=detections)
 
-    #crop image for the largest possible square
-    res_img = Image.fromarray(img)
-    print(type(res_img))
-
-    x=0
-    y=0
+    if "HeatMap" in annotators:
+        # Draw HeatMapAnnotator
+        heatmap_annotator = sv.HeatMapAnnotator()
+        img = heatmap_annotator.annotate(img, detections=detections)
+    
+    if "Dot" in annotators:
+        # Draw DotAnnotator
+        dot_annotator = sv.DotAnnotator()
+        img = dot_annotator.annotate(img, detections=detections)
+    
+    if "Triangle" in annotators:
+        # Draw TriangleAnnotator
+        tri_annotator = sv.TriangleAnnotator(sv.Color.from_hex(str(colortri)))
+        img = tri_annotator.annotate(img, detections=detections)
 
+    # crop image for the largest possible square
+    res_img = Image.fromarray(img)
+    # print(type(res_img))
+    x = 0
+    y = 0
 
-    print("size of the pil im=", res_img.size)
-    (v1,v2) = res_img.size
+    # print("size of the pil im=", res_img.size)
+    (v1, v2) = res_img.size
     width, height = calculate_crop_dim(v1, v2)
-    print(width, height)
+    # print(width, height)
     my_img = np.array(res_img)
 
-    crop_img = my_img[y:y+height, x:x+width]
-    print(type(crop_img))
+    crop_img = my_img[y : y + height, x : x + width]
+    # print(type(crop_img))
 
-    return crop_img[...,::-1].copy() # BGR to RGB using numpy
+    return crop_img[..., ::-1].copy()  # BGR to RGB using numpy
 
 
-with gr.Blocks(theme=gr.themes.Soft(primary_hue=gr.themes.colors.purple)
-    .set(
-      button_primary_background_fill="*primary_600",
-      button_primary_background_fill_hover="*primary_700",
-      checkbox_label_background_fill_selected="*primary_600",
-      checkbox_background_color_selected="*primary_400",
-    )) as demo:
+purple_theme = theme = gr.themes.Soft(primary_hue=gr.themes.colors.purple).set(
+    button_primary_background_fill="*primary_600",
+    button_primary_background_fill_hover="*primary_700",
+    checkbox_label_background_fill_selected="*primary_600",
+    checkbox_background_color_selected="*primary_400",
+)
+
+with gr.Blocks(theme=purple_theme) as app:
     gr.Markdown("""# Supervision Annotators""")
-    annotators = gr.CheckboxGroup(choices=["BoundingBox", "Mask", "Ellipse", "BoxCorner", "Circle", "Label", "Blur"], value=["BoundingBox", "Mask"], label="Select Annotators:")
+    annotators = gr.CheckboxGroup(
+        choices=[
+            "BoundingBox",
+            "Mask",
+            "Halo",
+            "Ellipse",
+            "BoxCorner",
+            "Circle",
+            "Label",
+            "Blur",
+            "Pixelate",
+            "HeatMap",
+            "Dot",
+            "Triangle"
+        ],
+        value=["BoundingBox", "Mask"],
+        label="Select Annotators:",
+    )
 
-    with gr.Accordion("**Color Picker**"):
-      with gr.Row():
+    gr.Markdown("🎨 **Color Picker**")
+    with gr.Row(variant="compact"):
         with gr.Column():
-          colorbb = gr.ColorPicker(value="#A351FB",label="BoundingBox")
+            colorbb = gr.ColorPicker(value="#A351FB", label="BoundingBox")
         with gr.Column():
-          colormask = gr.ColorPicker(value="#A351FB",label="Mask")
+            colormask = gr.ColorPicker(value="#A351FB", label="Mask")
         with gr.Column():
-          colorellipse = gr.ColorPicker(value="#A351FB",label="Ellipse")
+            colorellipse = gr.ColorPicker(value="#A351FB", label="Ellipse")
         with gr.Column():
-          colorbc = gr.ColorPicker(value="#A351FB",label="BoxCorner")
+            colorbc = gr.ColorPicker(value="#A351FB", label="BoxCorner")
         with gr.Column():
-          colorcir = gr.ColorPicker(value="#A351FB",label="Circle")
+            colorcir = gr.ColorPicker(value="#A351FB", label="Circle")
         with gr.Column():
-          colorlabel = gr.ColorPicker(value="#A351FB",label="Label")
-    
+            colorlabel = gr.ColorPicker(value="#A351FB", label="Label")
+        with gr.Column():
+            colorhalo = gr.ColorPicker(value="#A351FB", label="Halo")
+        with gr.Column():
+            colordot = gr.ColorPicker(value="#A351FB", label="Dot")
+        with gr.Column():
+            colortri = gr.ColorPicker(value="#A351FB", label="Triangle")
+
     with gr.Row():
-      with gr.Column():
-        with gr.Tab("Input image"):
-          image_input = gr.Image(type="numpy", show_label=False)
-      with gr.Column():
-        with gr.Tab("Result image"):
-          image_output = gr.Image(type="numpy", show_label=False)
+        with gr.Column():
+            with gr.Tab("Input image"):
+                image_input = gr.Image(type="numpy", show_label=False)
+        with gr.Column():
+            with gr.Tab("Result image"):
+                image_output = gr.Image(type="numpy", show_label=False)
     image_button = gr.Button(value="Annotate it!", variant="primary")
 
-    image_button.click(annotator, inputs=[image_input,annotators,colorbb,colormask,colorellipse,colorbc,colorcir,colorlabel], outputs=image_output)
+    image_button.click(
+        annotator,
+        inputs=[
+            image_input,
+            annotators,
+            colorbb,
+            colormask,
+            colorellipse,
+            colorbc,
+            colorcir,
+            colorlabel,
+            colorhalo,
+            colortri,
+        ],
+        outputs=image_output,
+    )
 
     gr.Markdown("## Image Examples")
     gr.Examples(
-        examples=[os.path.join(os.path.abspath(''), "city.jpg"), 
-                  os.path.join(os.path.abspath(''), "household.jpg"), 
-                  os.path.join(os.path.abspath(''), "industry.jpg"),
-                  os.path.join(os.path.abspath(''), "retail.jpg"),
-                  os.path.join(os.path.abspath(''), "aerodefence.jpg")],
+        examples=[
+            os.path.join(os.path.abspath(""), "city.jpg"),
+            os.path.join(os.path.abspath(""), "household.jpg"),
+            os.path.join(os.path.abspath(""), "industry.jpg"),
+            os.path.join(os.path.abspath(""), "retail.jpg"),
+            os.path.join(os.path.abspath(""), "aerodefence.jpg"),
+        ],
         inputs=image_input,
         outputs=image_output,
         fn=annotator,
@@ -157,4 +236,7 @@ with gr.Blocks(theme=gr.themes.Soft(primary_hue=gr.themes.colors.purple)
     )
 
 
-demo.launch(debug=False)
+if __name__ == "__main__":
+    print("Starting app...")
+    print("Dark theme is available at: http://localhost:7860/?__theme=dark")
+    app.launch(debug=False)