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 class VideoPeopleDetection(): time_reference = datetime.datetime.now() counter_frame = 0 processed_fps = 0 def __init__(self,video_path=None, model=None): self.model = model self.classes = self.model.names self.frame = [None] if video_path is not None: self.video_name = video_path else: self.video_name = 'vid2.mp4' # A default video file self.dataset = LoadImages(self.video_name) def use_webcam(self, source): # self.dataset.release() # Release any existing video capture # self.cap = cv2.VideoCapture(0) # Open default webcam # print('use_webcam') source = source self.imgsz = 640 cudnn.benchmark = True self.dataset = LoadStreams(source, img_size=self.imgsz) def class_to_label(self, x): return self.classes[int(x)] # def transpic(self, img): # return cv2.cvtColor(img, cv2.COLOR_BGR2RGB) def get_frame(self): for im0s in self.dataset: # print(self.dataset.mode) # print(self.dataset) if self.dataset.mode == 'stream': img = im0s[0].copy() else: img = im0s.copy() results = self.model(img, size=640) # Loop through each detected object and count the people accuracy = 0 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) resText=f'检测到{num_people}人' # print(num_people) return jpeg.tobytes(), resText def time_synchronized(): # pytorch-accurate time if torch.cuda.is_available(): torch.cuda.synchronize() return time.time()