import cv2
import gradio as gr
import spaces
from fastrtc import Stream, get_twilio_turn_credentials
from gradio.utils import get_space

try:
    from demo.object_detection.inference import YOLOv10
except (ImportError, ModuleNotFoundError):
    from inference import YOLOv10


# Load the model and place it on CUDA at module level. On ZeroGPU this uses the
# PyTorch CUDA emulation; the real GPU is attached inside @spaces.GPU below.
model = YOLOv10("yolov10n.pt").to("cuda")


@spaces.GPU
def detection(image, conf_threshold=0.3):
    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=get_twilio_turn_credentials() if get_space() else None,
    concurrency_limit=2 if get_space() else None,
)

# ZeroGPU only detects @spaces.GPU functions when Gradio is the launched app,
# so we expose and launch the built-in fastrtc UI (a Gradio Blocks) rather than
# mounting the stream on a FastAPI/uvicorn app.
demo = stream.ui

if __name__ == "__main__":
    import os

    if (mode := os.getenv("MODE")) == "PHONE":
        stream.fastphone(host="0.0.0.0", port=7860)
    else:
        demo.launch(server_name="0.0.0.0", server_port=7860)