import av
import cv2
import streamlit as st
from streamlit_webrtc import WebRtcMode, webrtc_streamer
import torch

from sample_utils.turn import get_ice_servers

st.markdown(
    "# Object detection with YOLO"
)

# YOLO estandar model
model_yolo = torch.hub.load('Ultralytics/yolov5','yolov5s', pretrained = True, verbose = False)
model_yolo.eval()
labels = model_yolo.names
person_class_index = 0


def callback(frame: av.VideoFrame) -> av.VideoFrame:

    # Get the image
    img = frame.to_ndarray(format="bgr24") 

    # Preprocessing image
    scale = 3 # percent of original size
    width = int(img.shape[1] / scale)
    height = int(img.shape[0] / scale)
    dim = (width, height)

    img2 = cv2.resize(img, dim, interpolation = cv2.INTER_AREA)

    # model detections results
    detections = model_yolo(img2)
    #print(detections)
    
    #   Draw bounding boxes
    for detection in detections.xyxy[0]:
        x1, y1, x2, y2, p, category_id = detection
        x1, y1, x2, y2, category_id = int(x1), int(y1), int(x2), int(y2), int(category_id)
        #if category_id == person_class_index:
        #cv2.rectangle(img, (x1*scale, y1*scale), (x2*scale, y2*scale), (55,139,241), 2)
        cv2.putText(img, 
                        labels[category_id], 
                        (x1*scale, y1*scale - 5), 
                        cv2.FONT_HERSHEY_TRIPLEX, 
                        0.5, (255,255,255), 1)
        xc = int((x1+x2)*scale/2.0)
        yc = int((y1+y2)*scale/2.0)
        cv2.circle(img, (xc, yc), 10, (255,255,255), -1)
        cv2.line(img, (x1*scale, y1*scale), (xc, yc), (255,255,255), 2)


    return av.VideoFrame.from_ndarray(img, format="bgr24")


webrtc_streamer(
    key="opencv-filter",
    mode=WebRtcMode.SENDRECV,
    rtc_configuration={"iceServers": [{"urls": ["stun:stun.l.google.com:19302"]}]},
    
    video_frame_callback=callback,
    media_stream_constraints={"video": True, "audio": False},
    async_processing=True,
)



st.markdown(
    "This demo is based on "
    "https://github.com/whitphx/streamlit-webrtc-example. "  
    "Many thanks to the project."
)