File size: 6,365 Bytes
b64c69e dc52a28 b64c69e 3ac1bb0 |
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 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 |
import cv2
import torch
import numpy as np
import math
from super_gradients.training import models
from super_gradients.training.processing import (DetectionCenterPadding,StandardizeImage,
ImagePermute, ComposeProcessing,
DetectionLongestMaxSizeRescale)
from deep_sort_pytorch.utils.parser import get_config
from deep_sort_pytorch.deep_sort import DeepSort
import streamlit as st
file_path = 'coco-labels-paper.txt'
palette = (2 ** 11 - 1, 2 ** 15 - 1, 2 ** 20 - 1)
names = []
with open(file_path, 'r') as file:
for line in file:
names.append(line.strip())
st.header(":hand: Welcome To YoLo Nas Object Detection and Tracking : ")
st.info("""
This app uses the cutting-edge YOLO Nas algorithm to detect objects in real-time video streams.
But that's not all it also employs the powerful DeepSort algorithm to track these objects,
providing you with seamless tracking capabilities. Easily upload a video feed,
and watch as our app identifies and tracks objects with precision.
It's simple, efficient, and ready to help you monitor and analyze moving objects effortlessly!
""")
with st.sidebar :
device_name =st.selectbox("Device : " , ["cpu" , "cuda"])
if device_name == 'cuda' :
device = torch.device("cuda:0")
else :
device = torch.device("cpu")
source_name = st.selectbox("select you source feed : " , ["URL"])
conf = st.slider("Select threshold confidence value : " , min_value=0.1 , max_value=1.0 , value=0.25)
iou = st.slider("Select Intersection over union (iou) value : " , min_value=0.1 , max_value=1.0 , value=0.5)
#model=models.get('yolo_nas_s',num_classes=len(names) ,
# checkpoint_path="yolo_nas_s_coco.pth").to(device)
if source_name == "URL" :
source = st.text_input("Input your Url Camera feed and press Entre ex : http://IP:8080/video")
cap = cv2.VideoCapture(source)
model=models.get('yolo_nas_s', pretrained_weights="coco").to(device)
model.set_dataset_processing_params(
class_names=names,
image_processor=ComposeProcessing(
[DetectionLongestMaxSizeRescale(output_shape=(636, 636)),
DetectionCenterPadding(output_shape=(640, 640),
pad_value=114),
StandardizeImage(max_value=255.0),
ImagePermute(permutation=(2, 0, 1)),]),
iou=iou ,conf=conf)
cfg_deep = get_config()
cfg_deep.merge_from_file("deep_sort_pytorch/configs/deep_sort.yaml")
deepsort = DeepSort(cfg_deep.DEEPSORT.REID_CKPT,
max_dist=cfg_deep.DEEPSORT.MAX_DIST, min_confidence=cfg_deep.DEEPSORT.MIN_CONFIDENCE,
nms_max_overlap=cfg_deep.DEEPSORT.NMS_MAX_OVERLAP,
max_iou_distance=cfg_deep.DEEPSORT.MAX_IOU_DISTANCE,
max_age=cfg_deep.DEEPSORT.MAX_AGE, n_init=cfg_deep.DEEPSORT.N_INIT,
nn_budget=cfg_deep.DEEPSORT.NN_BUDGET,
use_cuda=False)
def compute_color_for_labels(label):
"""
Simple function that adds fixed color depending on the class
"""
if label == 0: #person
color = (85,45,255)
elif label == 2: # Car
color = (222,82,175)
elif label == 3: # Motobike
color = (0, 204, 255)
elif label == 5: # Bus
color = (0, 149, 255)
else:
color = [int((p * (label ** 2 - label + 1)) % 255) for p in palette]
return tuple(color)
def draw_boxes(img, bbox, identities=None, categories=None, names=None, offset=(0,0)):
for i, box in enumerate(bbox):
x1, y1, x2, y2 = [int(i) for i in box]
x1 += offset[0]
x2 += offset[0]
y1 += offset[0]
y2 += offset[0]
cat = int(categories[i]) if categories is not None else 0
id = int(identities[i]) if identities is not None else 0
cv2.rectangle(img, (x1, y1), (x2, y2), color= compute_color_for_labels(cat),thickness=2, lineType=cv2.LINE_AA)
label = str(id) + ":" + names[cat]
(w,h), _ = cv2.getTextSize(str(label), cv2.FONT_HERSHEY_SIMPLEX, fontScale=1/2, thickness=1)
t_size=cv2.getTextSize(str(label), cv2.FONT_HERSHEY_SIMPLEX, fontScale=1/2, thickness=1)[0]
c2=x1+t_size[0], y1-t_size[1]-3
cv2.rectangle(img, (x1, y1), c2, color=compute_color_for_labels(cat), thickness=-1, lineType=cv2.LINE_AA)
cv2.putText(img, str(label), (x1, y1-2), 0, 1/2, [255, 255, 255], thickness=1, lineType=cv2.LINE_AA)
return img
if st.button("Start detection and Tracking") :
frame_window = st.image( [] )
while True:
xywh_bboxs = []
confs = []
oids = []
ret, frame = cap.read()
if ret:
result = list(model.predict(frame))[0]
bbox_xyxys = result.prediction.bboxes_xyxy.tolist()
confidences = result.prediction.confidence
labels = result.prediction.labels.tolist()
for (bbox_xyxy, confidence, cls) in zip(bbox_xyxys, confidences, labels):
bbox = np.array(bbox_xyxy)
x1, y1, x2, y2 = bbox[0], bbox[1], bbox[2], bbox[3]
x1, y1, x2, y2 = int(x1), int(y1), int(x2), int(y2)
conf = math.ceil((confidence*100))/100
cx, cy = int((x1+x2)/2), int((y1+y2)/2)
bbox_width = abs(x1-x2)
bbox_height = abs(y1-y2)
xcycwh = [cx, cy, bbox_width, bbox_height]
xywh_bboxs.append(xcycwh)
confs.append(conf)
oids.append(int(cls))
xywhs = torch.tensor(xywh_bboxs)
confss= torch.tensor(confs)
outputs = deepsort.update(xywhs, confss, oids, frame)
if len(outputs)>0:
bbox_xyxy = outputs[:,:4]
identities = outputs[:, -2]
object_id = outputs[:, -1]
draw_boxes(frame, bbox_xyxy, identities, object_id , names=names)
#output.write(frame)
#cv2.imshow('Video', frame)
#if cv2.waitKey(25) & 0xFF == ord('q'):
# break
frame = cv2.cvtColor( frame , cv2.COLOR_BGR2RGB)
frame_window.image(frame)
else:
break |