Update app.py

app.py

@@ -3,10 +3,20 @@ import numpy as np
 from PIL import Image, ImageDraw, ImageFont
 from collections import Counter
 import time
-import tempfile
+import os
 from ultralytics import YOLO
 import cv2
-import os
+from gradio_client.documentation import document, DocumentedType
+
+# Import WebRTC components
+from gradio_webrtc import (
+    RTCConfiguration,
+    WebRtcStreamerContext,
+    WebRtcMode,
+    WebRtcStreamer,
+    VideoTransformerBase,
+    VideoTransformerContext,
+)
 
 # Constants
 COIN_CLASS_ID = 11  # 10sen coin
@@ -114,7 +124,7 @@ def non_max_suppression(detections, iou_threshold):
 
     return [detections[i] for i in keep_indices]
 
-class VideoProcessor:
+class ScrewDetectionProcessor:
     def __init__(self):
         self.px_to_mm_ratio = None
         self.detected_objects = []
@@ -147,12 +157,12 @@ class VideoProcessor:
         if isinstance(frame, np.ndarray):
             frame_np = frame
         else:
-            # This handles the case if frame comes from Gradio's webcam which is already numpy
+            # This handles the case if frame comes from other sources
             frame_np = np.array(frame)
             
         results = model(frame_np, conf=self.confidence_threshold)
         
-        if not results:
+        if not results or len(results) == 0:
             return frame_np, []
         
         result = results[0]
@@ -226,6 +236,30 @@ class VideoProcessor:
         # Convert back to BGR for OpenCV operations
         return cv2.cvtColor(processed_img, cv2.COLOR_RGB2BGR), frame_detected_objects
 
+# WebRTC Video Transformer
+class ScrewDetectionTransformer(VideoTransformerBase):
+    def __init__(self):
+        self.processor = ScrewDetectionProcessor()
+        self.summary_text = "No detections yet."
+        
+    def update_settings(self, iou_threshold, confidence_threshold, show_detections, show_summary):
+        self.processor.update_settings(
+            iou_threshold=iou_threshold,
+            confidence_threshold=confidence_threshold,
+            show_detections=show_detections,
+            show_summary=show_summary
+        )
+        
+    def get_summary(self):
+        return self.processor.get_summary()
+        
+    def transform(self, frame):
+        # Process frame will be called on each video frame
+        img = frame.to_ndarray(format="bgr24")
+        processed_frame, _ = self.processor.process_frame(img)
+        self.summary_text = self.processor.get_summary()
+        return processed_frame
+
 def process_image(input_image, iou_threshold, confidence_threshold, show_detections, show_summary):
     if input_image is None:
         return None, "Please upload an image first."
@@ -237,7 +271,7 @@ def process_image(input_image, iou_threshold, confidence_threshold, show_detecti
         frame = input_image
     
     # Create a temporary processor for image processing
-    processor = VideoProcessor()
+    processor = ScrewDetectionProcessor()
     processor.update_settings(iou_threshold, confidence_threshold, show_detections, show_summary)
     processed_frame, _ = processor.process_frame(frame)
     
@@ -254,7 +288,7 @@ def process_video(video_path, iou_threshold, confidence_threshold, show_detectio
         
     try:
         # Create a processor for video processing
-        processor = VideoProcessor()
+        processor = ScrewDetectionProcessor()
         processor.update_settings(iou_threshold, confidence_threshold, show_detections, show_summary)
         
         cap = cv2.VideoCapture(video_path)
@@ -288,23 +322,20 @@ def process_video(video_path, iou_threshold, confidence_threshold, show_detectio
     except Exception as e:
         return [], f"Error processing video: {str(e)}"
 
-def process_webcam(frame, iou_threshold, confidence_threshold, show_detections, show_summary):
-    if frame is None:
-        return None
-        
-    # Convert from RGB to BGR for processing
-    frame_bgr = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)
-    
-    # Create a temporary processor for webcam processing
-    processor = VideoProcessor()
-    processor.update_settings(iou_threshold, confidence_threshold, show_detections, show_summary)
-    
-    # Process the frame
-    processed_frame, _ = processor.process_frame(frame_bgr)
-    
-    # Convert back to RGB for Gradio
-    processed_frame_rgb = cv2.cvtColor(processed_frame, cv2.COLOR_BGR2RGB)
-    return processed_frame_rgb
+def update_webrtc_settings(iou_threshold, confidence_threshold, show_detections, show_summary, webrtc_ctx):
+    if webrtc_ctx and webrtc_ctx.video_transformer:
+        webrtc_ctx.video_transformer.update_settings(
+            iou_threshold=iou_threshold,
+            confidence_threshold=confidence_threshold,
+            show_detections=show_detections,
+            show_summary=show_summary
+        )
+    return "Settings updated"
+
+def get_webrtc_summary(webrtc_ctx):
+    if webrtc_ctx and webrtc_ctx.video_transformer:
+        return webrtc_ctx.video_transformer.get_summary()
+    return "WebRTC not active"
     
 # Gradio Interface
 with gr.Blocks(title="Screw Detection and Measurement") as demo:
@@ -348,28 +379,57 @@ with gr.Blocks(title="Screw Detection and Measurement") as demo:
             outputs=[video_output, video_summary]
         )
     
-    with gr.Tab("Webcam"):
+    with gr.Tab("WebRTC Webcam"):
         with gr.Row():
-            with gr.Column():
+            with gr.Column(scale=1):
                 webcam_iou = gr.Slider(label="IoU Threshold (NMS)", minimum=0.0, maximum=1.0, value=0.7, step=0.05)
                 webcam_conf = gr.Slider(label="Confidence Threshold", minimum=0.0, maximum=1.0, value=0.5, step=0.05)
                 webcam_show_det = gr.Checkbox(label="Show Detections", value=True)
                 webcam_show_sum = gr.Checkbox(label="Show Summary", value=True)
-            with gr.Column():
-                # Use the compatible webcam syntax - this is compatible with older versions of gradio
-                webcam_input = gr.Image(label="Live Camera")
-                webcam_output = gr.Image(label="Processed Output")
-                webcam_button = gr.Button("Process Webcam Image")
+                
+                # Create a settings update button
+                update_settings = gr.Button("Update Settings")
+                
+                # Summary textbox
+                webcam_summary = gr.Textbox(label="Detection Summary", interactive=False)
+                
+                # Button to get summary
+                get_summary = gr.Button("Get Detection Summary")
+                
+            with gr.Column(scale=2):
+                # Configure WebRTC with STUN servers
+                rtc_config = RTCConfiguration(
+                    {"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]}
+                )
+                
+                # Create the WebRTC component with our transformer
+                webrtc_ctx = gr.State(None)
+                
+                # Use WebRtcStreamer with our transformer
+                webrtc = WebRtcStreamer(
+                    key="screw-detection",
+                    mode=WebRtcMode.SENDRECV,
+                    rtc_configuration=rtc_config,
+                    video_transformer_factory=ScrewDetectionTransformer,
+                    async_transform=True,
+                )
     
-        # Use click instead of change for webcam processing
-        webcam_button.click(
-            fn=process_webcam,
-            inputs=[webcam_input, webcam_iou, webcam_conf, webcam_show_det, webcam_show_sum],
-            outputs=webcam_output
+        # Connect the update settings button
+        update_settings.click(
+            update_webrtc_settings,
+            inputs=[webcam_iou, webcam_conf, webcam_show_det, webcam_show_sum, webrtc_ctx],
+            outputs=gr.Textbox(value="Settings updated", visible=False)
+        )
+        
+        # Connect the get summary button
+        get_summary.click(
+            get_webrtc_summary,
+            inputs=[webrtc_ctx],
+            outputs=webcam_summary
         )
 
 # Add warning about model loading
 if model is None:
     gr.Warning("Model could not be loaded. Please ensure 'yolo11-obb12classes.pt' is available.")
 
-demo.launch()
+demo.launch()