You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
281 lines
12 KiB
281 lines
12 KiB
from operator import index
|
|
import sys
|
|
from numpy.core.fromnumeric import size
|
|
|
|
from numpy.lib.nanfunctions import nanstd
|
|
sys.path.insert(0, './yolov5')
|
|
|
|
from yolov5.utils.datasets import LoadImages, LoadStreams
|
|
# from yolov5.utils.general import check_img_size, non_max_suppression, scale_coords
|
|
from yolov5.utils.general import check_img_size, non_max_suppression, scale_coords, xyxy2xywh
|
|
from yolov5.utils.plots import Annotator, colors
|
|
from yolov5.utils.torch_utils import select_device, time_synchronized
|
|
from deep_sort_pytorch.utils.parser import get_config
|
|
from deep_sort_pytorch.deep_sort import DeepSort
|
|
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.utils.Point2GPS import *
|
|
import ec.countunit
|
|
from ec.rtsppushserver import RtstPushServer
|
|
|
|
|
|
palette = (2 ** 11 - 1, 2 ** 15 - 1, 2 ** 20 - 1)
|
|
|
|
|
|
|
|
def detect(opt, save_img=False):
|
|
out, source, weights, view_img, save_txt, imgsz = \
|
|
opt.output, opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
|
|
webcam = source == '0' or source.startswith(
|
|
'rtsp') or source.startswith('rtmp') or source.startswith('http') or source.endswith('.txt')
|
|
|
|
# initialize deepsort
|
|
cfg = get_config()
|
|
cfg.merge_from_file(opt.config_deepsort)
|
|
|
|
deepsort = DeepSort(cfg.DEEPSORT.REID_CKPT,
|
|
max_dist=cfg.DEEPSORT.MAX_DIST, min_confidence=cfg.DEEPSORT.MIN_CONFIDENCE,
|
|
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,
|
|
use_cuda=True)
|
|
# Initialize
|
|
device = select_device(opt.device)
|
|
if os.path.exists(out):
|
|
shutil.rmtree(out) # delete output folder
|
|
os.makedirs(out) # make new output folder
|
|
half = device.type != 'cpu' # half precision only supported on CUDA
|
|
|
|
# Load model
|
|
model = torch.load(weights, map_location=device)[
|
|
'model'].float() # load to FP32
|
|
model.to(device).eval()
|
|
if half:
|
|
model.half() # to FP16
|
|
|
|
# Set Dataloader
|
|
vid_path, vid_writer = None, None
|
|
if webcam:
|
|
view_img = True
|
|
cudnn.benchmark = True # set True to speed up constant image size inference
|
|
dataset = LoadStreams(source, img_size=imgsz)
|
|
else:
|
|
view_img = True
|
|
save_img = True
|
|
dataset = LoadImages(source, img_size=imgsz)
|
|
|
|
# Get names and colors
|
|
names = model.module.names if hasattr(model, 'module') else model.names
|
|
|
|
# Run inference
|
|
t0 = time.time()
|
|
img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img
|
|
# run once
|
|
_ = model(img.half() if half else img) if device.type != 'cpu' else None
|
|
|
|
save_path = str(Path(out))
|
|
txt_path = str(Path(out)) + '/results.txt'
|
|
|
|
rtstPushServer = RtstPushServer(9,opt.view_img_size[0],opt.view_img_size[1],'rtsp://192.168.1.182:8554/video')
|
|
|
|
for frame_idx, (path, img, im0s, vid_cap) in enumerate(dataset):
|
|
draw_box=[]
|
|
# if(framenumber==0):
|
|
# continue
|
|
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)
|
|
|
|
|
|
|
|
# Inference
|
|
t1 = time_synchronized()
|
|
pred = model(img, augment=opt.augment)[0]
|
|
|
|
# Apply NMS
|
|
pred = non_max_suppression(
|
|
pred, opt.conf_thres, opt.iou_thres, classes=opt.classes, agnostic=opt.agnostic_nms)
|
|
t2 = time_synchronized()
|
|
|
|
# Process detections
|
|
for i, det in enumerate(pred): # detections per image
|
|
if webcam: # batch_size >= 1
|
|
p, s, im0 = path[i], '%g: ' % i, im0s[i].copy()
|
|
else:
|
|
p, s, im0 = path, '', im0s
|
|
# 设置显示窗体大小
|
|
im0= cv2.resize(im0,opt.view_img_size)
|
|
|
|
s += '%gx%g ' % img.shape[2:] # print string
|
|
save_path = str(Path(out) / Path(p).name)
|
|
|
|
if det is not None and len(det):
|
|
|
|
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]
|
|
outputs = deepsort.update(xywhs.cpu(), confs.cpu(), clss.cpu(), im0)
|
|
|
|
|
|
if outputs is not None and 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
|
|
# label = f'{names[c]} {conf:.2f}'
|
|
label = f'{id} {names[c]} {conf:.2f}'
|
|
color=colors(c, True)
|
|
draw_box.append([bboxes,label,color])
|
|
|
|
# draw_boxes(im0, bboxes, id,label)
|
|
# ori_im = draw_boxes2(im0, bbox_xyxy, identities, clsTracking, trace, h_inv)
|
|
|
|
|
|
|
|
# Write MOT compliant results to file
|
|
if save_txt and len(outputs) != 0:
|
|
for j, output in enumerate(outputs):
|
|
bbox_left = output[0]
|
|
bbox_top = output[1]
|
|
bbox_w = output[2]
|
|
bbox_h = output[3]
|
|
identity = output[-1]
|
|
with open(txt_path, 'a') as f:
|
|
f.write(('%g ' * 10 + '\n') % (frame_idx, identity, bbox_left,
|
|
bbox_top, bbox_w, bbox_h, -1, -1, -1, -1)) # label format
|
|
ec.countunit.addcount(outputs,im0)
|
|
|
|
else:
|
|
deepsort.increment_ages()
|
|
|
|
# Print time (inference + NMS)
|
|
# print('%sDone. (%.3fs)' % (s, t2 - t1))
|
|
|
|
# Stream results
|
|
if view_img:
|
|
im0= ec.countunit.addtitlerectangle(im0)
|
|
annotator = Annotator(im0, line_width=2, pil=True)
|
|
|
|
for bboxes,label,color in draw_box:
|
|
annotator.box_label(bboxes, label, color)
|
|
|
|
|
|
annotator.drawcount( ec.countunit.countlist, ec.countunit.clsnames)
|
|
im0=annotator.cvimg()
|
|
|
|
im0= ec.countunit.addcolorimg(im0)
|
|
rtstPushServer.write(im0)
|
|
cv2.imshow(p, im0)
|
|
if cv2.waitKey(1) == ord('q'): # q to quit
|
|
raise StopIteration
|
|
if save_img:
|
|
print('saving video!')
|
|
if vid_path != save_path: # new video
|
|
vid_path = save_path
|
|
if isinstance(vid_writer, cv2.VideoWriter):
|
|
vid_writer.release() # release previous video writer
|
|
|
|
fps =20# 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))
|
|
vid_writer = cv2.VideoWriter(
|
|
save_path, cv2.VideoWriter_fourcc(*opt.fourcc), fps,opt.view_img_size)
|
|
vid_writer.write(im0)
|
|
|
|
# Save results (image with detections)
|
|
# if save_img:
|
|
# print('saving img!')
|
|
# if dataset.mode == 'images':
|
|
# cv2.imwrite(save_path, im0)
|
|
# else:
|
|
# print('saving video!')
|
|
# if vid_path != save_path: # new video
|
|
# vid_path = save_path
|
|
# if isinstance(vid_writer, cv2.VideoWriter):
|
|
# vid_writer.release() # release previous video writer
|
|
|
|
# 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))
|
|
# vid_writer = cv2.VideoWriter(
|
|
# save_path, cv2.VideoWriter_fourcc(*opt.fourcc), fps, (w, h))
|
|
# vid_writer.write(im0)
|
|
|
|
# if save_txt or save_img:
|
|
# print('Results saved to %s' % os.getcwd() + os.sep + out)
|
|
# if platform == 'darwin': # MacOS
|
|
# os.system('open ' + save_path)
|
|
|
|
# print('Done. (%.3fs)' % (time.time() - t0))
|
|
|
|
|
|
if __name__ == '__main__':
|
|
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('--weights', type=str,
|
|
default='yolov5/weights/yolov5s.pt', help='model.pt path')
|
|
# file/folder, 0 for webcam
|
|
parser.add_argument('--source', type=str,
|
|
default='rtsp://admin:hk123456@192.168.1.65:554', help='source')
|
|
# default='F:/AIData/video/hqxjy.MP4', help='source')
|
|
# default='rtsp://admin:hk123456@120.240.37.42:554', help='source')
|
|
|
|
parser.add_argument('--output', type=str, default='inference/output',
|
|
help='output folder') # output folder
|
|
parser.add_argument('--img-size', type=int, default=960,
|
|
help='inference size (pixels)')
|
|
# conf 相似度
|
|
parser.add_argument('--conf-thres', type=float,
|
|
default=0.5, 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('--view-img', action='store_true',default=True,
|
|
help='display results')
|
|
#显示窗体大小--- 注意 countunit.py img_size =(1920,1080)也要设置
|
|
parser.add_argument('--view-img-size', type=int, default=(1920,1080),
|
|
help='inference size (pixels)')
|
|
parser.add_argument('--save-txt', action='store_true',
|
|
help='save results to *.txt')
|
|
# class 0 is person 跟踪类别
|
|
parser.add_argument('--classes', nargs='+', type=int,
|
|
default=[0,1,2,3,5,7,16], help='filter by class')
|
|
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("--config_deepsort", type=str,
|
|
default="deep_sort_pytorch/configs/deep_sort.yaml")
|
|
args = parser.parse_args()
|
|
args.img_size = check_img_size(args.img_size)
|
|
print(args)
|
|
|
|
with torch.no_grad():
|
|
print(torch.__version__) #注意是双下划线
|
|
|
|
detect(args)
|
|
|