Hugging Face's logo

Spaces:
Javiai
/
Object-Detection-Yolov5-webrtc
Runtime error

App Files Community
Object-Detection-Yolov5-webrtc
File size: 2,109 Bytes 
5722543
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
c5cab77
 
5722543
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
 
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."
)