Spaces:

fastrtc
/

object-detection

Running on T4

App Files Files Community

freddyaboulton HF Staff commited on Feb 20

Commit

534fffb

verified ·

1 Parent(s): 023470c

Upload folder using huggingface_hub

Browse files

Files changed (2) hide show

app.py +28 -10
index.html +45 -1

app.py CHANGED Viewed

@@ -3,8 +3,9 @@ from pathlib import Path
 import cv2
 import gradio as gr
 from fastapi.responses import HTMLResponse
-from fastrtc import Stream, get_twilio_turn_credentials
 from gradio.utils import get_space
 from huggingface_hub import hf_hub_download
 from pydantic import BaseModel, Field
@@ -12,7 +13,7 @@ from pydantic import BaseModel, Field
 try:
     from demo.object_detection.inference import YOLOv10
 except (ImportError, ModuleNotFoundError):
-    from inference import YOLOv10
 cur_dir = Path(__file__).parent
@@ -25,10 +26,16 @@ model = YOLOv10(model_file)
 def detection(image, conf_threshold=0.3):
-    image = cv2.resize(image, (model.input_width, model.input_height))
-    print("conf_threshold", conf_threshold)
-    new_image = model.detect_objects(image, conf_threshold)
-    return cv2.resize(new_image, (500, 500))
 stream = Stream(
@@ -40,8 +47,12 @@ stream = Stream(
     concurrency_limit=20 if get_space() else None,
 )
-@stream.get("/")
 async def _():
     rtc_config = get_twilio_turn_credentials() if get_space() else None
     html_content = open(cur_dir / "index.html").read()
@@ -54,12 +65,19 @@ class InputData(BaseModel):
     conf_threshold: float = Field(ge=0, le=1)
-@stream.post("/input_hook")
 async def _(data: InputData):
     stream.set_input(data.webrtc_id, data.conf_threshold)
 if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(stream, host="0.0.0.0", port=7860)

 import cv2
 import gradio as gr
+from fastapi import FastAPI
 from fastapi.responses import HTMLResponse
+from fastrtc import Stream, WebRTCError, get_twilio_turn_credentials
 from gradio.utils import get_space
 from huggingface_hub import hf_hub_download
 from pydantic import BaseModel, Field
 try:
     from demo.object_detection.inference import YOLOv10
 except (ImportError, ModuleNotFoundError):
+    from .inference import YOLOv10
 cur_dir = Path(__file__).parent
 def detection(image, conf_threshold=0.3):
+    try:
+        image = cv2.resize(image, (model.input_width, model.input_height))
+        print("conf_threshold", conf_threshold)
+        new_image = model.detect_objects(image, conf_threshold)
+        return cv2.resize(new_image, (500, 500))
+    except Exception as e:
+        import traceback
+        traceback.print_exc()
+        raise WebRTCError(str(e))
 stream = Stream(
     concurrency_limit=20 if get_space() else None,
 )
+app = FastAPI()
+stream.mount(app)
+@app.get("/")
 async def _():
     rtc_config = get_twilio_turn_credentials() if get_space() else None
     html_content = open(cur_dir / "index.html").read()
     conf_threshold: float = Field(ge=0, le=1)
+@app.post("/input_hook")
 async def _(data: InputData):
     stream.set_input(data.webrtc_id, data.conf_threshold)
 if __name__ == "__main__":
+    import os
+    if (mode := os.getenv("MODE")) == "UI":
+        stream.ui.launch(server_port=7860, server_name="0.0.0.0")
+    elif mode == "PHONE":
+        stream.fastphone(host="0.0.0.0", port=7860)
+    else:
+        import uvicorn
+        uvicorn.run(app, host="0.0.0.0", port=7860)

index.html CHANGED Viewed

@@ -112,10 +112,28 @@
             border-radius: 50%;
             cursor: pointer;
         }
     </style>
 </head>
 <body>
     <div class="container">
         <h1>Real-time Object Detection</h1>
         <p>Using YOLOv10 to detect objects in your webcam feed</p>
@@ -160,6 +178,17 @@
             });
         }
         async function setupWebRTC() {
             const config = __RTC_CONFIGURATION__;
             peerConnection = new RTCPeerConnection(config);
@@ -182,7 +211,9 @@
                 const dataChannel = peerConnection.createDataChannel('text');
                 dataChannel.onmessage = (event) => {
                     const eventJson = JSON.parse(event.data);
-                    if (eventJson.type === "send_input") {
                         updateConfThreshold(confThreshold.value);
                     }
                 };
@@ -217,6 +248,16 @@
                 });
                 const serverResponse = await response.json();
                 await peerConnection.setRemoteDescription(serverResponse);
                 // Send initial confidence threshold
@@ -224,6 +265,9 @@
             } catch (err) {
                 console.error('Error setting up WebRTC:', err);
             }
         }

             border-radius: 50%;
             cursor: pointer;
         }
+        /* Add styles for toast notifications */
+        .toast {
+            position: fixed;
+            top: 20px;
+            left: 50%;
+            transform: translateX(-50%);
+            background-color: #f44336;
+            color: white;
+            padding: 16px 24px;
+            border-radius: 4px;
+            font-size: 14px;
+            z-index: 1000;
+            display: none;
+            box-shadow: 0 2px 5px rgba(0, 0, 0, 0.2);
+        }
     </style>
 </head>
 <body>
+    <!-- Add toast element after body opening tag -->
+    <div id="error-toast" class="toast"></div>
     <div class="container">
         <h1>Real-time Object Detection</h1>
         <p>Using YOLOv10 to detect objects in your webcam feed</p>
             });
         }
+        function showError(message) {
+            const toast = document.getElementById('error-toast');
+            toast.textContent = message;
+            toast.style.display = 'block';
+            // Hide toast after 5 seconds
+            setTimeout(() => {
+                toast.style.display = 'none';
+            }, 5000);
+        }
         async function setupWebRTC() {
             const config = __RTC_CONFIGURATION__;
             peerConnection = new RTCPeerConnection(config);
                 const dataChannel = peerConnection.createDataChannel('text');
                 dataChannel.onmessage = (event) => {
                     const eventJson = JSON.parse(event.data);
+                    if (eventJson.type === "error") {
+                        showError(eventJson.message);
+                    } else if (eventJson.type === "send_input") {
                         updateConfThreshold(confThreshold.value);
                     }
                 };
                 });
                 const serverResponse = await response.json();
+                if (serverResponse.status === 'failed') {
+                    showError(serverResponse.meta.error === 'concurrency_limit_reached'
+                        ? `Too many connections. Maximum limit is ${serverResponse.meta.limit}`
+                        : serverResponse.meta.error);
+                    stop();
+                    startButton.textContent = 'Start';
+                    return;
+                }
                 await peerConnection.setRemoteDescription(serverResponse);
                 // Send initial confidence threshold
             } catch (err) {
                 console.error('Error setting up WebRTC:', err);
+                showError('Failed to establish connection. Please try again.');
+                stop();
+                startButton.textContent = 'Start';
             }
         }