Update app.py

app.py

@@ -1,78 +1,58 @@
 import json
 from pathlib import Path
-
-import cv2
-import numpy as np
-from PIL import Image
-from fastapi import FastAPI, Request
-from fastapi.responses import HTMLResponse, JSONResponse
-from pydantic import BaseModel, Field
+from fastapi import FastAPI
+from fastapi.responses import HTMLResponse
+from fastrtc import Stream
 from huggingface_hub import hf_hub_download
-from io import BytesIO
-import base64
+from pydantic import BaseModel, Field
+import os
+import cv2
 
-try:
-    from demo.object_detection.inference import YOLOv10
-except (ImportError, ModuleNotFoundError):
-    from inference import YOLOv10
+from inference import YOLOv10
 
-# Define app and paths
-app = FastAPI()
 cur_dir = Path(__file__).parent
 
-# Load YOLOv10 ONNX model
 model_file = hf_hub_download(
     repo_id="onnx-community/yolov10n", filename="onnx/model.onnx"
 )
-model = YOLOv10(model_file)
-
-# Serve the index.html file
-@app.get("/", response_class=HTMLResponse)
-async def serve_frontend():
-    html_path = cur_dir / "index.html"
-    with open(html_path, "r", encoding="utf-8") as f:
-        html_content = f.read()
-
-    # Replace placeholder with empty RTC config or other configs if needed
-    html_content = html_content.replace("__RTC_CONFIGURATION__", json.dumps({}))
-    return HTMLResponse(content=html_content)
 
+model = YOLOv10(model_file)
 
-# Model input format
-class ImagePayload(BaseModel):
-    image: str  # base64 string
-    conf_threshold: float = Field(default=0.3, ge=0, le=1)
-
-
-# Inference route
-@app.post("/detect")
-async def detect_objects(payload: ImagePayload):
-    try:
-        # Decode base64 image
-        header, encoded = payload.image.split(",", 1)
-        img_bytes = base64.b64decode(encoded)
-        img = Image.open(BytesIO(img_bytes)).convert("RGB")
-        img_np = np.array(img)
-
-        # Resize for model input
-        img_resized = cv2.resize(img_np, (model.input_width, model.input_height))
-
-        # Run detection
-        output_image = model.detect_objects(img_resized, payload.conf_threshold)
-
-        # Return detections (if you want to send image back, convert to base64)
-        return JSONResponse(content={"status": "success"})
+def detection(image, conf_threshold=0.3):
+    image = cv2.resize(image, (model.input_width, model.input_height))
+    new_image = model.detect_objects(image, conf_threshold)
+    return cv2.resize(new_image, (500, 500))
+
+stream = Stream(
+    handler=detection,
+    modality="video",
+    mode="send-receive",
+    additional_inputs=[gr.Slider(minimum=0, maximum=1, step=0.01, value=0.3)],
+    rtc_configuration=None,
+    concurrency_limit=None,
+)
 
-    except Exception as e:
-        return JSONResponse(content={"status": "error", "message": str(e)}, status_code=500)
+app = FastAPI()
+stream.mount(app)
 
+@app.get("/")
+async def serve_index():
+    html_content = open(cur_dir / "index.html").read()
+    return HTMLResponse(content=html_content)
 
-# Optional: health check
-@app.get("/health")
-async def health():
-    return {"status": "ok"}
+class InputData(BaseModel):
+    webrtc_id: str
+    conf_threshold: float = Field(ge=0, le=1)
 
+@app.post("/input_hook")
+async def update_input(data: InputData):
+    stream.set_input(data.webrtc_id, data.conf_threshold)
 
 if __name__ == "__main__":
-    import uvicorn
-    uvicorn.run(app, host="0.0.0.0", port=7860)
+    if (mode := os.getenv("MODE")) == "UI":
+        stream.ui.launch(server_port=7860)
+    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)