Update app.py

app.py

@@ -2,76 +2,77 @@ import json
 from pathlib import Path
 
 import cv2
-import gradio as gr
-from fastapi import FastAPI
-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
+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 huggingface_hub import hf_hub_download
+from io import BytesIO
+import base64
 
 try:
     from demo.object_detection.inference import YOLOv10
 except (ImportError, ModuleNotFoundError):
     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()
 
-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))
+    # 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)
 
 
-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=get_twilio_turn_credentials() if get_space() else None,
-    concurrency_limit=2 if get_space() else None,
-)
+# Model input format
+class ImagePayload(BaseModel):
+    image: str  # base64 string
+    conf_threshold: float = Field(default=0.3, ge=0, le=1)
 
-app = FastAPI()
 
-stream.mount(app)
+# 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))
 
-@app.get("/")
-async def _():
-    rtc_config = get_twilio_turn_credentials() if get_space() else None
-    html_content = open(cur_dir / "index.html").read()
-    html_content = html_content.replace("__RTC_CONFIGURATION__", json.dumps(rtc_config))
-    return HTMLResponse(content=html_content)
+        # 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"})
 
-class InputData(BaseModel):
-    webrtc_id: str
-    conf_threshold: float = Field(ge=0, le=1)
+    except Exception as e:
+        return JSONResponse(content={"status": "error", "message": str(e)}, status_code=500)
 
 
-@app.post("/input_hook")
-async def _(data: InputData):
-    stream.set_input(data.webrtc_id, data.conf_threshold)
+# Optional: health check
+@app.get("/health")
+async def health():
+    return {"status": "ok"}
 
 
 if __name__ == "__main__":
-    import os
-
-    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)
+    import uvicorn
+    uvicorn.run(app, host="0.0.0.0", port=7860)