Update app.py

app.py

@@ -178,23 +178,58 @@ def draw_classification_boxes(image_with_insects, detections):
             f"{label}: {score:.2f}",
             (box.xmin, box.ymin - baseline),
             cv2.FONT_HERSHEY_SIMPLEX,
-            0.5,
+             0.5,
             (255, 255, 255),
              2
         )
     return image_with_insects
 
+def plot_detections_plotly(image: np.ndarray, detections: List[DetectionResult]) -> str:
+    from plotly import graph_objects as go
+    import plotly.express as px
+    fig = px.imshow(image)
+    class_colors = {i: f'rgb({random.randint(0, 255)}, {random.randint(0, 255)}, {random.randint(0, 255)})' for i in range(len(detections))}
+    for idx, detection in enumerate(detections):
+        label = detection.label
+        box = detection.box
+        score = detection.score
+        mask = detection.mask
+        polygon = mask_to_polygon(mask)
+        fig.add_trace(go.Scatter(
+            x=[point[0] for point in polygon] + [polygon[0][0]],
+            y=[point[1] for point in polygon] + [polygon[0][1]],
+            mode='lines',
+            line=dict(color=class_colors[idx], width=2),
+            fill='toself',
+            name=f"{label}: {score:.2f}"
+        ))
+        xmin, ymin, xmax, ymax = box.xyxy
+        fig.add_shape(
+            type="rect",
+            x0=xmin, y0=ymin, x1=xmax, y1=ymax,
+            line=dict(color=class_colors[idx])
+        )
+        fig.add_annotation(
+            x=(xmin + xmax) // 2, y=(ymin + ymax) // 2,
+            text=f"{label}: {score:.2f}",
+        )
+    fig.update_layout(xaxis=dict(visible=False), yaxis=dict(visible=False))
+    file_path = "/tmp/plotly_image.html"
+    fig.write_html(file_path)
+    return file_path
+
 def process_image(image):
     labels = ["insect"]
     original_image, detections = grounded_segmentation(image, labels, threshold=0.3, polygon_refinement=True)
     annotated_image = plot_detections(original_image, detections)
     yellow_background_with_insects = create_yellow_background_with_insects(np.array(original_image), detections)
     yellow_background_with_boxes = draw_classification_boxes(yellow_background_with_insects.copy(), detections)
-    return annotated_image, yellow_background_with_boxes
+    plotly_image_path = plot_detections_plotly(original_image, detections)
+    return annotated_image, yellow_background_with_boxes, plotly_image_path
 
 gr.Interface(
     fn=process_image,
     inputs=gr.Image(type="pil"),
-    outputs=[gr.Image(type="numpy"), gr.Image(type="numpy")],
+    outputs=[gr.Image(type="numpy"), gr.Image(type="numpy"), gr.HTML()],
     title="🐞 InsectSAM + GroundingDINO Inference"
 ).launch()