# limit the number of cpus used by high performance libraries
import os
cpu_num = 1
os.environ["OMP_NUM_THREADS"] = str(cpu_num)
os.environ["OPENBLAS_NUM_THREADS"] = str(cpu_num)
os.environ["MKL_NUM_THREADS"] = str(cpu_num)
os.environ["VECLIB_MAXIMUM_THREADS"] = str(cpu_num)
os.environ["NUMEXPR_NUM_THREADS"] = str(cpu_num)
import sys
# 新添加的目录会优先于其他目录被import检查，程序退出后失效
sys.path.insert(0, './yolov5')
import numpy as np
import argparse
import os
import platform
import shutil
import time
from pathlib import Path
import cv2
import torch
import torch.backends.cudnn as cudnn
from yolov5.models.experimental import attempt_load
from yolov5.utils.downloads import attempt_download
from yolov5.models.common import DetectMultiBackend
from yolov5.utils.datasets import LoadImages, LoadStreams
from yolov5.utils.general import (LOGGER, check_img_size, non_max_suppression, scale_coords,
check_imshow, xyxy2xywh, increment_path)
from yolov5.utils.torch_utils import select_device, time_sync
from yolov5.utils.plots import Annotator, colors
from deep_sort.utils.parser import get_config
from deep_sort.deep_sort import DeepSort
from yolov5.utils.util.saveuitl import sava_dict
from yolov5.utils.util.saveuitl import process_area
from yolov5.utils.util.saveuitl import is_poi_in_poly
from yolov5.utils.util.saveuitl import plot_dots_on_people_and_car
from yolov5.utils.util.saveuitl import distancing
from yolov5.utils.util.saveuitl import call_the_police
from yolov5.utils.util.saveuitl import draw_word_in_img
from yolov5.utils.util.saveuitl import draw_area
FILE = Path(__file__).resolve()
ROOT = FILE.parents[0] # yolov5 deepsort root directory
if str(ROOT) not in sys.path:
sys.path.append(str(ROOT)) # add ROOT to PATH
ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative
def detect(opt, tracks=None):
out, source, yolo_model, deep_sort_model, show_vid, save_vid, save_txt, imgsz, evaluate, half, project, name, exist_ok = \
opt.output, opt.source, opt.yolo_model, opt.deep_sort_model, opt.show_vid, opt.save_vid, \
opt.save_txt, opt.imgsz, opt.evaluate, opt.half, opt.project, opt.name, opt.exist_ok
webcam = source == '0' or source.startswith(
'rtsp') or source.startswith('http') or source.endswith('.txt')
device = select_device(opt.device)
# initialize deepsort
cfg = get_config()
cfg.merge_from_file(opt.config_deepsort)
deepsort = DeepSort(deep_sort_model,
device,
max_dist=cfg.DEEPSORT.MAX_DIST,
max_iou_distance=cfg.DEEPSORT.MAX_IOU_DISTANCE,
max_age=cfg.DEEPSORT.MAX_AGE, n_init=cfg.DEEPSORT.N_INIT, nn_budget=cfg.DEEPSORT.NN_BUDGET,
)
# Initialize
half &= device.type != 'cpu' # half precision only supported on CUDA
# The MOT16 evaluation runs multiple inference streams in parallel, each one writing to
# its own .txt file. Hence, in that case, the output folder is not restored
if not evaluate:
if os.path.exists(out):
pass
shutil.rmtree(out) # delete output folder
os.makedirs(out) # make new output folder
# Directories
if type(yolo_model) is str:
exp_name = yolo_model.split(".")[0]
elif type(yolo_model) is list and len(yolo_model) == 1:
exp_name = yolo_model[0].split(".")[0]
else:
exp_name = "ensemble"
exp_name = exp_name + "_" + deep_sort_model.split('/')[-1].split('.')[0]
save_dir = increment_path(Path(project) / exp_name, exist_ok=exist_ok) # increment run if project name exists
save_dir.mkdir(parents=True, exist_ok=True) # make dir
# Load model
device = select_device(device)
model = DetectMultiBackend(yolo_model, device=device, dnn=opt.dnn)
stride, names, pt, jit, _ = model.stride, model.names, model.pt, model.jit, model.onnx
imgsz = check_img_size(imgsz, s=stride) # check image size
# Half
half &= pt and device.type != 'cpu' # half precision only supported by PyTorch on CUDA
if pt:
model.model.half() if half else model.model.float()
# Set Dataloader
vid_path, vid_writer = None, None
# Check if environment supports image displays
if show_vid:
show_vid = check_imshow()
# Dataloader
if webcam:
show_vid = check_imshow()
cudnn.benchmark = True # set True to speed up constant image size inference
dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt and not jit)
bs = len(dataset) # batch_size
else:
dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt and not jit)
bs = 1 # batch_size
vid_path, vid_writer = [None] * bs, [None] * bs
# Get names and colors
names = model.module.names if hasattr(model, 'module') else model.names
# extract what is in between the last '/' and last '.'
txt_file_name = source.split('/')[-1].split('.')[0]
txt_path = str(Path(save_dir)) + '/' + txt_file_name + '.txt'
if pt and device.type != 'cpu':
model(torch.zeros(1, 3, *imgsz).to(device).type_as(next(model.model.parameters()))) # warmup
dt, seen = [0.0, 0.0, 0.0, 0.0], 0
# 定义保存数据的结构
mydict = {}
# 被检测车数量
checkNum = 0
# 未被检测车数量
unCheckNum = 0
# 未被检测车字典信息
errorDict = {}
# 已报过警的人id
already_alarm_people_ids = []
# # 车指定区域
area_poly = [[315, 326], [12, 812], [1009, 1069], [1179, 545]]
# 人指定区域
pepple_area_poly = [[1052, 405], [817, 968], [1349, 1071], [1423, 493]]
# 红色区域
red_area = [[88, 421], [1371, 625], [1353, 740], [45, 504]]
# 蓝色区域
blue_area = [[12, 538], [1335, 794], [1343, 922], [27, 655]]
# 上行人数
up_persons = 0
# 下行人数
down_persons = 0
#已经过红区的人id
red_area_ids = []
#已经过蓝区的人id
blue_area_ids = []
for frame_idx, (path, img, im0s, vid_cap, s) in enumerate(dataset):
t1 = time_sync()
img = torch.from_numpy(img).to(device)
img = img.half() if half else img.float() # uint8 to fp16/32
img /= 255.0 # 0 - 255 to 0.0 - 1.0
if img.ndimension() == 3:
img = img.unsqueeze(0)
t2 = time_sync()
dt[0] += t2 - t1
# Inference
visualize = increment_path(save_dir / Path(path[0]).stem, mkdir=True) if opt.visualize else False
pred = model(img, augment=opt.augment, visualize=visualize)
t3 = time_sync()
dt[1] += t3 - t2
# Apply NMS
pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, opt.classes, opt.agnostic_nms,
max_det=opt.max_det)
dt[2] += time_sync() - t3
# Process detections 每一帧里面开始遍历
for i, det in enumerate(pred): # detections per image
seen += 1
if webcam: # batch_size >= 1
p, im0, _ = path[i], im0s[i].copy(), dataset.count
s += f'{i}: '
else:
p, im0, _ = path, im0s.copy(), getattr(dataset, 'frame', 0)
p = Path(p) # to Path
save_path = str(save_dir / p.name) # im.jpg, vid.mp4, ...
s += '%gx%g ' % img.shape[2:] # print string
annotator = Annotator(im0, line_width=2, pil=not ascii)
if det is not None and len(det):
# Rescale boxes from img_size to im0 size
det[:, :4] = scale_coords(
img.shape[2:], det[:, :4], im0.shape).round()
# Print results
for c in det[:, -1].unique():
n = (det[:, -1] == c).sum() # detections per class
s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string
xywhs = xyxy2xywh(det[:, 0:4])
confs = det[:, 4]
clss = det[:, 5]
# pass detections to deepsort(xywhs是一帧中的多个舞台)
t4 = time_sync()
outputs = deepsort.update(xywhs.cpu(), confs.cpu(), clss.cpu(), im0)
t5 = time_sync()
dt[3] += t5 - t4
# 描绘指定区域
# start_draw_time = time.time()
process_area(im0, area_poly,color=(192,192,192))
# process_area(im0, pepple_area_poly,color=(0,97,255))
process_area(im0, red_area,color=(0,0,255))
process_area(im0, blue_area,color=(255,0,0))
# if frame_idx == 500:
# end_time = time.time()
# print("每帧都画框,共500帧,画区域和检测共用时", (end_time - start_time) * 1000, "ms")
# end_draw_time = time.time()
# print("在一帧中画需要所需要的时间为:" , (end_draw_time-start_draw_time)*1000, "ms")
# 每一帧后清空车和人坐标
people_coords = []
car_coords = []
# 每一帧需要报警的人数
people_num_of_frame = 0
# 每一帧需要报警的id信息
people_ids_of_frame = []
# 每帧中的车数
cars_of_farme = 0
# draw boxes for visualization
if len(outputs) > 0:
for j, (output, conf) in enumerate(zip(outputs, confs)):
bboxes = output[0:4]
id = output[4]
cls = output[5]
c = int(cls) # integer class
idbboxes = output[0:5]
label = f'{id} {names[c]} {conf:.2f}'
right_x = (bboxes[2] + bboxes[0]) // 2
# 第一个参数 右下角坐标 第二个 帧数 第三个 是否被检查的状态1为被检查 第四个参数是人和车距离在阈值内的帧数
dict = {id : [[right_x,bboxes[3]],1,0,0]}
# 统计上行人和下行人
if is_poi_in_poly([right_x,bboxes[3]],blue_area):
if red_area_ids.count(id) == 0:
if blue_area_ids.count(id) == 0:
blue_area_ids.append(id)
else:
if blue_area_ids.count(id) == 0:
blue_area_ids.append(id)
down_persons += 1
elif is_poi_in_poly([right_x,bboxes[3]],red_area):
if blue_area_ids.count(id) == 0:
if red_area_ids.count(id) == 0:
red_area_ids.append(id)
else:
if red_area_ids.count(id) == 0:
red_area_ids.append(id)
up_persons += 1
# 检测人指定区域的数量
if is_poi_in_poly([right_x,bboxes[3]],pepple_area_poly) and already_alarm_people_ids.count(id) == 0:
if label is not None:
if ((label.split())[1] == 'person'):
people_num_of_frame += 1
already_alarm_people_ids.append(id)
people_ids_of_frame.append(id)
# 检测指定区域内车是否被检测
if is_poi_in_poly([right_x,bboxes[3]],area_poly):
# 人车坐标分别添加到对应列表中
if label is not None:
if ((label.split())[1] == 'person'):
people_coords.append(idbboxes)
elif ((label.split())[1] == 'car'):
car_coords.append(idbboxes)
# 给字典添加数据
mydict, checkNum = sava_dict(mydict, id, dict, checkNum)
# 给人和车描绘中心点
plot_dots_on_people_and_car(idbboxes, im0)
else:
mydict, errorDict, unCheckNum = call_the_police(id,mydict,errorDict,area_poly,unCheckNum)
annotator.box_label(bboxes, label, color=colors(c, True))
bboxes[2] = right_x
# 描绘车右下中心点
annotator.repoint(label, bboxes, color=(0, 255, 255))
#统计一帧中区域内的车数
if is_poi_in_poly([right_x,bboxes[3]],area_poly) and (label.split())[1] == 'car':
cars_of_farme += 1
if save_txt:
# to MOT format
bbox_left = output[0]
bbox_top = output[1]
bbox_w = output[2] - output[0]
bbox_h = output[3] - output[1]
# Write MOT compliant results to file
with open(txt_path, 'a') as f:
f.write(('%g ' * 10 + '\n') % (frame_idx + 1, id, bbox_left, # MOT format
bbox_top, bbox_w, bbox_h, -1, -1, -1, -1))
# 检测人和车之间的距离
distancing(people_coords, car_coords,im0,dist_thres_lim=(0,200),mydict=mydict)
LOGGER.info(f'{s}Done. YOLO:({t3 - t2:.3f}s), DeepSort:({t5 - t4:.3f}s), mydict:({mydict})')
# 如果每帧中有未报警的人数,则调用报警接口
if people_num_of_frame > 0:
LOGGER.info(f'people_num_of_frame:({people_num_of_frame}), people_ids_of_frame:({people_ids_of_frame})')
# 一帧中的区域内没有任何物体,则清空车字典
if cars_of_farme == 0:
mydict.clear()
else:
deepsort.increment_ages()
LOGGER.info('No detections')
# Stream results
im0 = annotator.result()
text_checkDraw = 'CHECK: ' + str(checkNum)
test_unCheckDraw = 'NOCKECK: ' + str(unCheckNum)
text_upperson = 'UP_Per: ' + str(up_persons)
test_downperson = 'DOWN_Per: ' + str(down_persons)
# 给图片添加文字
im0 = draw_word_in_img(im0,text_checkDraw,org=(50,50),color=(0, 255, 0))
im0 = draw_word_in_img(im0, test_unCheckDraw, org=(50, 100), color=(0, 0, 255))
im0 = draw_word_in_img(im0, text_upperson, org=(50, 150), color=(0, 255, 0))
im0 = draw_word_in_img(im0, test_downperson, org=(50, 200), color=(0, 0, 255))
if show_vid:
cv2.imshow(str(p), im0)
if cv2.waitKey(1) == ord('q'): # q to quit
raise StopIteration
# Save results (image with detections)
if save_vid:
if vid_path != save_path: # new video
vid_path = save_path
if isinstance(vid_writer, cv2.VideoWriter):
vid_writer.release() # release previous video writer
if vid_cap: # video
fps = vid_cap.get(cv2.CAP_PROP_FPS)
w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))
h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
else: # stream
fps, w, h = 30, im0.shape[1], im0.shape[0]
vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h))
vid_writer.write(im0)
# Print results
t = tuple(x / seen * 1E3 for x in dt) # speeds per image
LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS, %.1fms deep sort update \
per image at shape {(1, 3, *imgsz)}' % t)
if save_txt or save_vid:
print('Results saved to %s' % save_path)
if platform == 'darwin': # MacOS
os.system('open ' + save_path)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--yolo_model', nargs='+', type=str, default='yolov5m.pt', help='model.pt path(s)')
parser.add_argument('--deep_sort_model', type=str, default='osnet_ibn_x1_0_MSMT17.pth')
parser.add_argument('--source', type=str, default='che.mp4', help='source') # file/folder, 0 for webcam --source http://admin:admin@192.168.0.102:8081
parser.add_argument('--output', type=str, default='inference/output', help='output folder') # output folder
parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w')
parser.add_argument('--conf-thres', type=float, default=0.2, help='object confidence threshold')
parser.add_argument('--iou-thres', type=float, default=0.5, help='IOU threshold for NMS')
parser.add_argument('--fourcc', type=str, default='mp4v', help='output video codec (verify ffmpeg support)')
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--show-vid', action='store_true', help='display tracking video results')
parser.add_argument('--save-vid', action='store_true', help='save video tracking results')
parser.add_argument('--save-txt', action='store_true', help='save MOT compliant results to *.txt')
# class 0 is person, 1 is bycicle, 2 is car... 79 is oven
parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 16 17')
parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')
parser.add_argument('--augment', action='store_true', help='augmented inference')
parser.add_argument('--evaluate', action='store_true', help='augmented inference')
parser.add_argument("--config_deepsort", type=str, default="deep_sort/configs/deep_sort.yaml")
parser.add_argument("--half", action="store_true", help="use FP16 half-precision inference")
parser.add_argument('--visualize', action='store_true', help='visualize features')
parser.add_argument('--max-det', type=int, default=1000, help='maximum detection per image')
parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')
parser.add_argument('--project', default=ROOT / 'runs/track', help='save results to project/name')
parser.add_argument('--name', default='exp', help='save results to project/name')
parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')
opt = parser.parse_args()
opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1 # expand
# 可以不被追踪到：一是减少内存，二是可以把这个operator从computation graph中detach出来，这样就不会在BP过程中计算到
with torch.no_grad():
detect(opt)
02-03
2528
2528
被折叠的 条评论
为什么被折叠?
到【灌水乐园】发言
