Added visualization for bbox predictions

app.py

@@ -1,19 +1,44 @@
 import gradio as gr
 import mediapipe as mp
-#from mediapipe.tasks import python
-#from mediapipe.tasks.python import vision
+import cv2
+import numpy as np
 
 BaseOptions = mp.tasks.BaseOptions
 ObjectDetector = mp.tasks.vision.ObjectDetector
 ObjectDetectorOptions = mp.tasks.vision.ObjectDetectorOptions
 VisionRunningMode = mp.tasks.vision.RunningMode
 
+MARGIN = 10  # pixels
+ROW_SIZE = 10  # pixels
+FONT_SIZE = 1
+FONT_THICKNESS = 1
+TEXT_COLOR = (255, 0, 0)  # red
+
+def visualize(image, detection_result) -> np.ndarray:
+  for detection in detection_result.detections:
+    # Draw bounding_box
+    bbox = detection.bounding_box
+    start_point = bbox.origin_x, bbox.origin_y
+    end_point = bbox.origin_x + bbox.width, bbox.origin_y + bbox.height
+    cv2.rectangle(image, start_point, end_point, TEXT_COLOR, 3)
+
+    # Draw label and score
+    category = detection.categories[0]
+    category_name = category.category_name
+    probability = round(category.score, 2)
+    result_text = category_name + ' (' + str(probability) + ')'
+    text_location = (MARGIN + bbox.origin_x,
+                     MARGIN + ROW_SIZE + bbox.origin_y)
+    cv2.putText(image, result_text, text_location, cv2.FONT_HERSHEY_PLAIN,
+                FONT_SIZE, TEXT_COLOR, FONT_THICKNESS)
+  return image
+
 def analyze_image(image):
     model_path = "efficientdet_lite0.tflite"
 
     options = ObjectDetectorOptions(
     base_options=BaseOptions(model_asset_path=model_path),
-    max_results=1,
+    max_results=5,
     running_mode=VisionRunningMode.IMAGE)
 
     mp_image = mp.Image(image_format=mp.ImageFormat.SRGB, data=image)
@@ -21,14 +46,14 @@ def analyze_image(image):
     with ObjectDetector.create_from_options(options) as detector:
         detection_result = detector.detect(mp_image)
 
-    results = ""
-    for i,detection in enumerate(detection_result.detections):
-        results += f'Detection {i}: {detection.categories[0].category_name}\n'
+    image_copy = np.copy(image.numpy_view())
+    annotated_image = visualize(image_copy, detection_result)
+    rgb_annotated_image = cv2.cvtColor(annotated_image, cv2.COLOR_BGR2RGB)
 
-    return results
+    return rgb_annotated_image
 
-img = gr.Image()
-txt = gr.Text()
+img_in = gr.Image()
+#img_out = gr.Image()
 
-iface = gr.Interface(fn=analyze_image, inputs=img, outputs=txt)
+iface = gr.Interface(fn=analyze_image, inputs=img_in, outputs="image")
 iface.launch()

requirements.txt

@@ -1,2 +1,3 @@
 numpy
-mediapipe
+mediapipe
+opencv-python