import argparse import time from pathlib import Path import sys import cv2 import torch import torch.backends.cudnn as cudnn from numpy import random from yolov7.models.experimental import attempt_load from yolov7.utils.datasets import LoadStreams, LoadImages from yolov7.utils.general import check_img_size, check_requirements, check_imshow, non_max_suppression, \ apply_classifier, \ scale_coords, xyxy2xywh, strip_optimizer, set_logging, increment_path from yolov7.utils.plots import plot_one_box from yolov7.utils.torch_utils import select_device, load_classifier, time_synchronized, TracedModel from tracker.mc_bot_sort import BoTSORT from tracker.tracking_utils.timer import Timer sys.path.insert(0, './yolov7') sys.path.append('.') # Global trackerTimer = Timer() timer = Timer() def write_results(filename, results): save_format = '{frame},{id},{x1},{y1},{w},{h},{s},-1,-1,-1\n' with open(filename, 'w') as f: for frame_id, tlwhs, track_ids, scores in results: for tlwh, track_id, score in zip(tlwhs, track_ids, scores): if track_id < 0: continue x1, y1, w, h = tlwh line = save_format.format(frame=frame_id, id=track_id, x1=round(x1, 1), y1=round(y1, 1), w=round(w, 1), h=round(h, 1), s=round(score, 2)) f.write(line) print('save results to {}'.format(filename)) def detect(): source, weights, view_img, imgsz, trace = opt.source, opt.weights, opt.view_img, opt.img_size, opt.trace save_img = opt.save_frames and not source.endswith('.txt') # save inference images webcam = source.isnumeric() or source.endswith('.txt') or source.lower().startswith( ('rtsp://', 'rtmp://', 'http://', 'https://')) # Directories save_dir = Path(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)) # increment run if save_img: save_dir.mkdir(parents=True, exist_ok=True) # make dir # Initialize set_logging() device = select_device(opt.device) half = device.type != 'cpu' # half precision only supported on CUDA # Load model model = attempt_load(weights, map_location=device) # load FP32 model stride = int(model.stride.max()) # model stride imgsz = check_img_size(imgsz, s=stride) # check img_size if trace: model = TracedModel(model, device, opt.img_size) if half: model.half() # to FP16 # Second-stage classifier classify = False if classify: modelc = load_classifier(name='resnet101', n=2) # initialize modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']).to(device).eval() # Set Dataloader vid_path, vid_writer = None, None if webcam: view_img = check_imshow() cudnn.benchmark = True # set True to speed up constant image size inference dataset = LoadStreams(source, img_size=imgsz, stride=stride) else: dataset = LoadImages(source, img_size=imgsz, stride=stride) if opt.ablation: dataset.files = dataset.files[len(dataset.files) // 2 + 1:] dataset.nf = len(dataset.files) # Get names and colors names = model.module.names if hasattr(model, 'module') else model.names colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(100)] # Create tracker tracker = BoTSORT(opt, frame_rate=30.0) # Run inference if device.type != 'cpu': model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters()))) # run once t0 = time.time() results = [] fn = 0 for path, img, im0s, vid_cap in dataset: fn += 1 timer.tic() 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() # Apply Classifier if classify: pred = apply_classifier(pred, modelc, img, im0s) # Process detections for i, det in enumerate(pred): # detections per image if webcam: # batch_size >= 1 p, s, im0, frame = path[i], '%g: ' % i, im0s[i].copy(), dataset.count else: p, s, im0, frame = path, '', im0s, getattr(dataset, 'frame', 0) # Run tracker detections = [] if len(det): boxes = scale_coords(img.shape[2:], det[:, :4], im0.shape) boxes = boxes.cpu().numpy() detections = det.cpu().numpy() detections[:, :4] = boxes trackerTimer.tic() online_targets = tracker.update(detections, im0) trackerTimer.toc() timer.toc() online_tlwhs = [] online_ids = [] online_scores = [] online_cls = [] for t in online_targets: tlwh = t.tlwh tlbr = t.tlbr tid = t.track_id tcls = t.cls vertical = tlwh[2] / tlwh[3] > opt.aspect_ratio_thresh if tlwh[2] * tlwh[3] > opt.min_box_area and not vertical: online_tlwhs.append(tlwh) online_ids.append(tid) online_scores.append(t.score) online_cls.append(t.cls) # save results results.append( f"{fn},{tid},{tlwh[0]:.2f},{tlwh[1]:.2f},{tlwh[2]:.2f},{tlwh[3]:.2f},{t.score:.2f},-1,-1,-1\n" ) if save_img or view_img: # Add bbox to image if opt.hide_labels_name: label = f'{tid}, {int(tcls)}' else: label = f'{tid}, {names[int(tcls)]}' plot_one_box(tlbr, im0, label=label, color=colors[int(tid) % len(colors)], line_thickness=2) p = Path(p) # to Path save_path = str(save_dir / p.name) # img.jpg # Stream results if view_img: cv2.imshow('BoT-SORT', im0) cv2.waitKey(1) # 1 millisecond # Save results (image with detections) if save_img: if dataset.mode == 'image': cv2.imwrite(save_path, im0) else: # 'video' or 'stream' 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] save_path += '.mp4' vid_writer = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h)) vid_writer.write(im0) res_file = opt.project + '/' + opt.name + ".txt" with open(res_file, 'w') as f: f.writelines(results) print(f"save results to {res_file}") print(f'Done. ({time.time() - t0:.3f}s)') if __name__ == '__main__': parser = argparse.ArgumentParser() parser.add_argument('source', type=str, default='inference/images', help='source') # file/folder, 0 for webcam parser.add_argument('--weights', nargs='+', type=str, default='yolov7.pt', help='model.pt path(s)') parser.add_argument("--benchmark", dest="benchmark", type=str, default='MOT17', help="benchmark to evaluate: MOT17 | MOT20") parser.add_argument("--eval", dest="split_to_eval", type=str, default='test', help="split to evaluate: train | val | test") parser.add_argument('--img-size', type=int, default=1280, help='inference size (pixels)') parser.add_argument('--conf-thres', type=float, default=0.09, help='object confidence threshold') parser.add_argument('--iou-thres', type=float, default=0.7, help='IOU threshold for NMS') 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', help='display results') parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --class 0, or --class 0 2 3') 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("--fp16", dest="fp16", default=False, action="store_true", help="Adopting mix precision evaluating.") parser.add_argument("--fuse", dest="fuse", default=False, action="store_true", help="Fuse conv and bn for testing.") parser.add_argument('--project', default='runs/track', help='save results to project/name') parser.add_argument('--name', default='exp', help='save results to project/name') parser.add_argument('--trace', action='store_true', help='trace model') parser.add_argument('--hide-labels-name', default=False, action='store_true', help='hide labels') parser.add_argument("--default-parameters", dest="default_parameters", default=False, action="store_true", help="use the default parameters as in the paper") parser.add_argument("--save-frames", dest="save_frames", default=False, action="store_true", help="save sequences with tracks.") # tracking args parser.add_argument("--track_high_thresh", type=float, default=0.5, help="tracking confidence threshold") parser.add_argument("--track_low_thresh", default=0.1, type=float, help="lowest detection threshold") parser.add_argument("--new_track_thresh", default=0.6, type=float, help="new track thresh") parser.add_argument("--track_buffer", type=int, default=30, help="the frames for keep lost tracks") parser.add_argument("--match_thresh", type=float, default=0.8, help="matching threshold for tracking") parser.add_argument("--aspect_ratio_thresh", type=float, default=1.6, help="threshold for filtering out boxes of which aspect ratio are above the given value.") parser.add_argument('--min_box_area', type=float, default=10, help='filter out tiny boxes') parser.add_argument("--fuse-score", dest="mot20", default=False, action="store_true", help="fuse score and iou for association") # CMC parser.add_argument("--cmc-method", default="file", type=str, help="cmc method: files (Vidstab GMC) | sparseOptFlow |orb | ecc") parser.add_argument("--ablation", dest="ablation", default=False, action="store_true", help="ablation ") # ReID parser.add_argument("--with-reid", dest="with_reid", default=False, action="store_true", help="with ReID module.") parser.add_argument("--fast-reid-config", dest="fast_reid_config", default=r"fast_reid/configs/MOT17/sbs_S50.yml", type=str, help="reid config file path") parser.add_argument("--fast-reid-weights", dest="fast_reid_weights", default=r"pretrained/mot17_sbs_S50.pth", type=str, help="reid config file path") parser.add_argument('--proximity_thresh', type=float, default=0.5, help='threshold for rejecting low overlap reid matches') parser.add_argument('--appearance_thresh', type=float, default=0.25, help='threshold for rejecting low appearance similarity reid matches') opt = parser.parse_args() opt.jde = False print(opt) opt.exist_ok = True data_path = opt.source device = opt.device if opt.benchmark == 'MOT20': train_seqs = [1, 2, 3, 5] test_seqs = [4, 6, 7, 8] seqs_ext = [''] MOT = 20 elif opt.benchmark == 'MOT17': train_seqs = [2, 4, 5, 9, 10, 11, 13] test_seqs = [1, 3, 6, 7, 8, 12, 14] seqs_ext = ['FRCNN', 'DPM', 'SDP'] MOT = 17 else: raise ValueError("Error: Unsupported benchmark:" + opt.benchmark) ablation = False if opt.split_to_eval == 'train': seqs = train_seqs elif opt.split_to_eval == 'val': seqs = train_seqs ablation = True elif opt.split_to_eval == 'test': seqs = test_seqs else: raise ValueError("Error: Unsupported split to evaluate:" + opt.split_to_eval) mainTimer = Timer() mainTimer.tic() for ext in seqs_ext: for i in seqs: if i < 10: seq = 'MOT' + str(MOT) + '-0' + str(i) else: seq = 'MOT' + str(MOT) + '-' + str(i) if ext != '': seq += '-' + ext opt.name = seq opt.ablation = ablation opt.mot20 = MOT == 20 opt.fps = 30 opt.device = device opt.batch_size = 1 opt.trt = False split = 'train' if i in train_seqs else 'test' opt.source = data_path + '/' + split + '/' + seq + '/' + 'img1' if opt.default_parameters: # TODO: Set YOLOv7 models if MOT == 20: # MOT20 opt.weights = r'./pretrained/yolov7-d6-mot20.pt' opt.match_thresh = 0.7 else: # MOT17 if ablation: opt.weights = r'./pretrained/yolov7-d6-mot17-ablation.pt' else: opt.weights = r'./pretrained/yolov7-d6-mot17.pt' # TODO: Set tracking parameters for YOLOv7 opt.track_high_thresh = 0.3 opt.track_low_thresh = 0.1 opt.track_buffer = 30 opt.agnostic_nms = True opt.new_track_thresh = opt.track_high_thresh + 0.1 opt.test_conf = max(0.001, opt.track_low_thresh - 0.01) with torch.no_grad(): detect() mainTimer.toc() print("TOTAL TIME END-to-END (with loading networks and images): ", mainTimer.total_time) print("TOTAL TIME (Detector + Tracker): " + str(timer.total_time) + ", FPS: " + str(1.0 /timer.average_time)) print("TOTAL TIME (Tracker only): " + str(trackerTimer.total_time) + ", FPS: " + str(1.0 / trackerTimer.average_time))