74 lines
2.1 KiB
Python
74 lines
2.1 KiB
Python
|
import datetime
|
||
|
|
import os
|
||
|
|
import time
|
||
|
|
import ffmpeg
|
||
|
|
import torch
|
||
|
|
import cv2
|
||
|
|
import numpy as np
|
||
|
|
from multiprocessing import Process, Manager
|
||
|
|
from threading import Thread
|
||
|
|
from read_data import LoadImages, LoadStreams
|
||
|
|
import torch.backends.cudnn as cudnn
|
||
|
|
|
||
|
|
|
||
|
|
def use_webcam(source, model):
|
||
|
|
|
||
|
|
source = source
|
||
|
|
imgsz = 640
|
||
|
|
cudnn.benchmark = True
|
||
|
|
dataset = LoadStreams(source, img_size=imgsz)
|
||
|
|
|
||
|
|
|
||
|
|
for im0s in dataset:
|
||
|
|
# print(self.dataset.mode)
|
||
|
|
# print(self.dataset)
|
||
|
|
if dataset.mode == 'stream':
|
||
|
|
img = im0s[0].copy()
|
||
|
|
else:
|
||
|
|
img = im0s.copy()
|
||
|
|
|
||
|
|
results = model(img, size=640)
|
||
|
|
|
||
|
|
|
||
|
|
# Loop through each detected object and count the people
|
||
|
|
num_people = 0
|
||
|
|
bgr = (0, 255, 0)
|
||
|
|
|
||
|
|
|
||
|
|
for obj in results.xyxy[0]:
|
||
|
|
# xmin, ymin, xmax, ymax = map(int, obj[:4])
|
||
|
|
# accuracy = obj[4]
|
||
|
|
# if (accuracy > 0.5):
|
||
|
|
|
||
|
|
# cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (0, 0, 255), 2)
|
||
|
|
# cv2.putText(img, f" {round(float(accuracy), 2), self.classes[obj[-1].item()]}", (xmin, ymin),
|
||
|
|
# cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 0, 255), 2)
|
||
|
|
|
||
|
|
if obj[-1] == 0: # 0 is the class ID for 'person'
|
||
|
|
|
||
|
|
# Draw bounding boxes around people
|
||
|
|
xmin, ymin, xmax, ymax = map(int, obj[:4])
|
||
|
|
accuracy = obj[4]
|
||
|
|
if (accuracy > 0.5):
|
||
|
|
num_people += 1
|
||
|
|
cv2.rectangle(img, (xmin, ymin), (xmax, ymax), (0, 0, 255), 2)
|
||
|
|
cv2.putText(img, f" {round(float(accuracy), 2)}", (xmin, ymin),
|
||
|
|
cv2.FONT_HERSHEY_SIMPLEX, 0.9, (0, 0, 255), 2)
|
||
|
|
|
||
|
|
|
||
|
|
# Draw the number of people on the frame and display it
|
||
|
|
|
||
|
|
|
||
|
|
ret, jpeg = cv2.imencode(".jpg", img)
|
||
|
|
|
||
|
|
|
||
|
|
return jpeg.tobytes()
|
||
|
|
|
||
|
|
|
||
|
|
def time_synchronized():
|
||
|
|
# pytorch-accurate time
|
||
|
|
if torch.cuda.is_available():
|
||
|
|
torch.cuda.synchronize()
|
||
|
|
return time.time()
|
||
|
|
|