k20hcmus/FishEye8K
3
1import argparse2 3import os4# limit the number of cpus used by high performance libraries5# os.environ["OMP_NUM_THREADS"] = "8"6# os.environ["OPENBLAS_NUM_THREADS"] = "8"7# os.environ["MKL_NUM_THREADS"] = "8"8# os.environ["VECLIB_MAXIMUM_THREADS"] = "8"9# os.environ["NUMEXPR_NUM_THREADS"] = "8"10import platform11import sys12import numpy as np13from pathlib import Path14import torch15import torch.backends.cudnn as cudnn16from numpy import random17from time import time18import pandas as pd19 20 21FILE = Path(__file__).resolve()22ROOT = FILE.parents[0] # yolov5 strongsort root directory23WEIGHTS = ROOT / 'weights'24if str(ROOT) not in sys.path:25 sys.path.append(str(ROOT)) # add ROOT to PATH26if str(ROOT / 'yolov9') not in sys.path:27 sys.path.append(str(ROOT / 'yolov9')) # add yolov5 ROOT to PATH28if str(ROOT / 'strong_sort') not in sys.path:29 sys.path.append(str(ROOT / 'strong_sort')) # add strong_sort ROOT to PATH30ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative31from models.experimental import attempt_load32from models.common import DetectMultiBackend33from utils.dataloaders import LoadImages, LoadStreams, LoadScreenshots34from utils.general import (LOGGER, Profile, check_file, check_img_size, check_imshow, check_requirements, colorstr, cv2,35 increment_path, non_max_suppression, print_args, scale_boxes, strip_optimizer, xyxy2xywh)36from utils.torch_utils import select_device, time_sync, smart_inference_mode37from utils.plots import Annotator, colors, save_one_box38from strong_sort.utils.parser import get_config39from strong_sort.strong_sort import StrongSORT40 41 42VID_FORMATS = 'asf', 'avi', 'gif', 'm4v', 'mkv', 'mov', 'mp4', 'mpeg', 'mpg', 'ts', 'wmv' # include video suffixes43 44 45def plot_one_box(x, img, color=None, label=None, line_thickness=3):46 # Plots one bounding box on image img47 tl = line_thickness or round(0.002 * (img.shape[0] + img.shape[1]) / 2) + 1 # line/font thickness48 c1, c2 = (int(x[0]), int(x[1])), (int(x[2]), int(x[3]))49 cv2.rectangle(img, c1, c2, color, thickness=tl, lineType=cv2.LINE_AA)50 if label:51 tf = max(tl - 1, 1) # font thickness52 t_size = cv2.getTextSize(label, 0, fontScale=tl / 3, thickness=tf)[0]53 c2 = c1[0] + t_size[0], c1[1] - t_size[1] - 354 cv2.rectangle(img, c1, c2, color, -1, cv2.LINE_AA) # filled55 cv2.putText(img, label, (c1[0], c1[1] - 2), 0, tl / 3, [225, 255, 255], thickness=tf, lineType=cv2.LINE_AA)56 57 58 59def convert_to_int(tensor):60 return tensor.type(torch.int16).item()61 62@smart_inference_mode()63def run_strongsort(64 source='0',65 data = ROOT / 'data/coco.yaml', # data.yaml path66 yolo_weights=WEIGHTS / 'yolo.pt', # model.pt path(s),67 strong_sort_weights=WEIGHTS / 'osnet_x0_25_msmt17.pt', # model.pt path,68 config_strongsort=ROOT / 'strong_sort/configs/strong_sort.yaml',69 imgsz=(640, 640), # inference size (height, width)70 conf_thres=0.25, # confidence threshold71 iou_thres=0.45, # NMS IOU threshold72 max_det=1000, # maximum detections per image73 device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu74 view_img=False, # show results75 save_txt=False, # save results to *.txt76 save_conf=False, # save confidences in --save-txt labels77 save_crop=False, # save cropped prediction boxes78 nosave=False, # do not save images/videos79 classes=None, # filter by class: --class 0, or --class 0 2 380 agnostic_nms=False, # class-agnostic NMS81 augment=False, # augmented inference82 visualize=False, # visualize features83 update=False, # update all models84 project=ROOT / 'runs/track', # save results to project/name85 name='exp', # save results to project/name86 exist_ok=False, # existing project/name ok, do not increment87 line_thickness=3, # bounding box thickness (pixels)88 hide_labels=False, # hide labels89 hide_conf=False, # hide confidences90 half=False, # use FP16 half-precision inference91 dnn=False, # use OpenCV DNN for ONNX inference92 vid_stride=1, # video frame-rate stride93):94 95 source = str(source)96 save_img = not nosave and not source.endswith('.txt') # save inference images97 is_file = Path(source).suffix[1:] in (VID_FORMATS)98 is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://'))99 webcam = source.isnumeric() or source.endswith('.txt') or (is_url and not is_file)100 screenshot = source.lower().startswith('screen')101 102 if is_url and is_file:103 source = check_file(source) # download104 105 # Directories106 if not isinstance(yolo_weights, list): # single yolo model107 exp_name = Path(yolo_weights).stem108 elif type(yolo_weights) is list and len(yolo_weights) == 1: # single models after --yolo_weights109 exp_name = Path(yolo_weights[0]).stem110 yolo_weights = Path(yolo_weights[0])111 else: # multiple models after --yolo_weights112 exp_name = 'ensemble'113 exp_name = name if name else exp_name + "_" + Path(strong_sort_weights).stem114 save_dir = increment_path(Path(project) / exp_name, exist_ok=exist_ok) # increment run115 save_dir = Path(save_dir)116 (save_dir / 'tracks' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir117 118 # Load model119 device = select_device(device)120 model = DetectMultiBackend(yolo_weights, device=device, dnn=dnn, data=data, fp16=half)121 stride, names, pt = model.stride, model.names, model.pt122 imgsz = check_img_size(imgsz, s=stride) # check image size123 124 # Dataloader125 126 # Dataloader127 bs = 1 # batch_size128 if webcam:129 view_img = check_imshow(warn=True)130 dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)131 bs = len(dataset)132 elif screenshot:133 dataset = LoadScreenshots(source, img_size=imgsz, stride=stride, auto=pt)134 else:135 dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride)136 vid_path, vid_writer,txt_path = [None] * bs, [None] * bs, [None] * bs137 138 139 # initialize StrongSORT140 cfg = get_config()141 cfg.merge_from_file(config_strongsort)142 143 # Create as many strong sort instances as there are video sources144 strongsort_list = []145 for i in range(bs):146 strongsort_list.append(147 StrongSORT(148 strong_sort_weights,149 device,150 half,151 max_dist=cfg.STRONGSORT.MAX_DIST,152 max_iou_distance=cfg.STRONGSORT.MAX_IOU_DISTANCE,153 max_age=cfg.STRONGSORT.MAX_AGE,154 n_init=cfg.STRONGSORT.N_INIT,155 nn_budget=cfg.STRONGSORT.NN_BUDGET,156 mc_lambda=cfg.STRONGSORT.MC_LAMBDA,157 ema_alpha=cfg.STRONGSORT.EMA_ALPHA,158 159 )160 )161 strongsort_list[i].model.warmup()162 outputs = [None] * bs163 164 colors = [[0, 0, 255], [255, 148, 0], [0, 255, 10], [0, 247, 250], [235,0,255]]165 #250, 247, 0166 # Run tracking167 model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *imgsz)) # warmup168 seen, windows, dt,sdt = 0, [], (Profile(), Profile(), Profile(), Profile()),[0.0, 0.0, 0.0, 0.0]169 curr_frames, prev_frames = [None] * bs, [None] * bs170 frame_counts = []171 172 for frame_idx, (path, im, im0s, vid_cap, s) in enumerate(dataset):173 # s = ''174 t1 = time_sync()175 with dt[0]:176 im = torch.from_numpy(im).to(model.device)177 im = im.half() if model.fp16 else im.float() # uint8 to fp16/32178 im /= 255 # 0 - 255 to 0.0 - 1.0179 if len(im.shape) == 3:180 im = im[None] # expand for batch dim181 t2 = time_sync()182 sdt[0] += t2 - t1183 184 # Inference185 with dt[1]:186 visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False187 pred = model(im, augment=augment, visualize=visualize)188 # pred = pred[0][1]189 t3 = time_sync()190 sdt[1] += t3 - t2191 192 # Apply NMS193 with dt[2]:194 pred = pred[0][1] if isinstance(pred[0], list) else pred[0] # single model or ensemble195 pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det)196 sdt[2] += time_sync() - t3197 198 # Second-stage classifier (optional)199 # pred = utils.general.apply_classifier(pred, classifier_model, im, im0s)200 201 counts = {}202 # Process detections203 for i, det in enumerate(pred): # detections per image204 seen += 1205 if webcam: # bs >= 1206 p, im0, _ = path[i], im0s[i].copy(), dataset.count207 p = Path(p) # to Path208 s += f'{i}: '209 # txt_file_name = p.name210 txt_file_name = p.stem + f'_{i}' # Unique text file name211 # save_path = str(save_dir / p.name) + str(i) # im.jpg, vid.mp4, ...212 save_path = str(save_dir / p.stem) + f'_{i}' # Unique video file name213 214 else:215 p, im0, _ = path, im0s.copy(), getattr(dataset, 'frame', 0)216 217 218 p = Path(p) # to Path219 # video file220 if source.endswith(VID_FORMATS):221 txt_file_name = p.stem222 save_path = str(save_dir / p.name) # im.jpg, vid.mp4, ...223 # folder with imgs224 else:225 txt_file_name = p.parent.name # get folder name containing current img226 save_path = str(save_dir / p.parent.name) # im.jpg, vid.mp4, ...227 228 curr_frames[i] = im0229 230 txt_path = str(save_dir / 'tracks' / txt_file_name) # im.txt231 s += '%gx%g ' % im.shape[2:] # print string232 gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh233 imc = im0.copy() if save_crop else im0 # for save_crop234 annotator = Annotator(im0, line_width=line_thickness, example=str(names))235 236 237 if cfg.STRONGSORT.ECC: # camera motion compensation238 strongsort_list[i].tracker.camera_update(prev_frames[i], curr_frames[i])239 240 if det is not None and len(det):241 # Rescale boxes from img_size to im0 size242 det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round()243 244 # Print results245 for c in det[:, -1].unique():246 n = (det[:, -1] == c).sum() # detections per class247 s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string248 counts[names[int(c)]] = n249 xywhs = xyxy2xywh(det[:, 0:4])250 confs = det[:, 4]251 clss = det[:, 5]252 253 # pass detections to strongsort254 t4 = time_sync()255 outputs[i] = strongsort_list[i].update(xywhs.cpu(), confs.cpu(), clss.cpu(), im0)256 t5 = time_sync()257 sdt[3] += t5 - t4258 259 # Write results260 for j, (output, conf) in enumerate(zip(outputs[i], confs)):261 xyxy = output[0:4]262 id = output[4]263 cls = output[5]264 label = names[int(cls)]265 # for *xyxy, conf, cls in reversed(det):266 if save_txt: # Write to file267 xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh268 # line = (id , cls, *xywh, conf) if save_conf else (cls, *xywh) # label format269 line = ( int(p.stem), frame_idx, id , cls, *xywh, conf) if save_conf else ( p.stem, frame_idx, cls, *xywh) # label format270 with open(txt_path + '.txt', 'a') as file:271 file.write(('%g ' * len(line) + '\n') % line)272 273 if save_img or save_crop or view_img: # Add bbox to image274 c = int(cls) # integer class275 label = None if hide_labels else ( str(id) + ' ' + names[c] if hide_conf else f' { id } {names[c]} {conf:.2f}')276 plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=2)277 if save_crop:278 save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True)279 280 frame_counts.append({'frame': frame_idx, 'counts': counts.copy()})281 # # draw boxes for visualization282 # if len(outputs[i]) > 0:283 # for j, (output, conf) in enumerate(zip(outputs[i], confs)):284 285 # bboxes = output[0:4]286 # id = output[4]287 # cls = output[5]288 289 # if save_txt:290 # # to MOT format291 # bbox_left = output[0]292 # bbox_top = output[1]293 # bbox_w = output[2] - output[0]294 # bbox_h = output[3] - output[1]295 # # format video_name frame id xmin ymin width height score class 296 # with open(txt_path + '.txt', 'a') as file:297 # file.write(f'{p.stem} {frame_idx} {id} {bbox_left} {bbox_top} {bbox_w} {bbox_h} {conf:.2f} {cls}\n')298 299 # if save_img or save_crop or view_img: # Add bbox to image300 # c = int(cls) # integer class301 # id = int(id) # integer id302 # label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}')303 # plot_one_box(bboxes, im0, label=label, color=colors[int(cls)], line_thickness=2)304 # if save_crop:305 # txt_file_name = txt_file_name if (isinstance(path, list) and len(path) > 1) else ''306 # save_one_box(bboxes, imc, file=save_dir / 'crops' / txt_file_name / names[c] / f'{id}' / f'{p.stem}.jpg', BGR=True)307 308 print(f'{s}Done. YOLO:({t3 - t2:.3f}s), StrongSORT:({t5 - t4:.3f}s)')309 310 else:311 strongsort_list[i].increment_ages()312 print('No detections')313 314 # Stream results315 im0 = annotator.result()316 317 318 if view_img:319 if platform.system() == 'Linux' and p not in windows:320 windows.append(p)321 cv2.namedWindow(str(p), cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO) # allow window resize (Linux)322 cv2.resizeWindow(str(p), im0.shape[1], im0.shape[0])323 cv2.imshow(str(p), im0)324 cv2.waitKey(1) # 1 millisecond325 326 # Save results (image with detections)327 if save_img:328 if dataset.mode == 'image':329 cv2.imwrite(save_path, im0)330 else: # 'video' or 'stream'331 if vid_path[i] != save_path: # new video332 vid_path[i] = save_path333 if isinstance(vid_writer[i], cv2.VideoWriter):334 vid_writer[i].release() # release previous video writer335 if vid_cap: # video336 fps = vid_cap.get(cv2.CAP_PROP_FPS)337 w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH))338 h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT))339 else: # stream340 fps, w, h = 30, im0.shape[1], im0.shape[0]341 save_path = str(Path(save_path).with_suffix('.mp4')) # force *.mp4 suffix on results videos342 vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc('m','p','4','v'), fps, (w, h))343 vid_writer[i].write(im0)344 345 prev_frames[i] = curr_frames[i]346 347 348 # Print time (inference-only)349 LOGGER.info(f"{s}{'' if len(det) else '(no detections), '}{dt[1].dt * 1E3:.1f}ms")350 351 flattened_counts = [352 {'frame': entry['frame'], 'label': label, 'count': count}353 for entry in frame_counts for label, count in entry['counts'].items()354 ]355 frame_counts_df = pd.DataFrame(flattened_counts)356 frame_counts_df['count'] = frame_counts_df['count'].apply(convert_to_int)357 counts_df = None358 # Print results359 LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape, %.1fms StrongSORT' % tuple(1E3 * x / seen for x in sdt))360 if save_txt or save_img:361 s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else ''362 LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}")363 if update:364 strip_optimizer(yolo_weights[0]) # update model (to fix SourceChangeWarning)365 return save_path, counts_df, frame_counts_df366def parse_opt():367 parser = argparse.ArgumentParser()368 parser.add_argument('--yolo-weights', nargs='+', type=str, default=WEIGHTS / 'yolov9.pt', help='model.pt path(s)')369 parser.add_argument('--strong-sort-weights', type=str, default=WEIGHTS / 'osnet_x0_25_msmt17.pt')370 parser.add_argument('--config-strongsort', type=str, default='strong_sort/configs/strong_sort.yaml')371 parser.add_argument('--source', type=str, default='0', help='file/dir/URL/glob, 0 for webcam') 372 parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='(optional) dataset.yaml path')373 parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w')374 parser.add_argument('--conf-thres', type=float, default=0.5, help='confidence threshold')375 parser.add_argument('--iou-thres', type=float, default=0.5, help='NMS IoU threshold')376 parser.add_argument('--max-det', type=int, default=1000, help='maximum detections per image')377 parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')378 parser.add_argument('--view-img', action='store_true', help='show results')379 parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')380 parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels')381 parser.add_argument('--save-crop', action='store_true', help='save cropped prediction boxes')382 parser.add_argument('--nosave', action='store_true', help='do not save images/videos')383 # class 0 is person, 1 is bycicle, 2 is car... 79 is oven384 parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --classes 0, or --classes 0 2 3')385 parser.add_argument('--agnostic-nms', action='store_true', help='class-agnostic NMS')386 parser.add_argument('--augment', action='store_true', help='augmented inference')387 parser.add_argument('--visualize', action='store_true', help='visualize features')388 parser.add_argument('--update', action='store_true', help='update all models')389 parser.add_argument('--project', default=ROOT / 'runs/track', help='save results to project/name')390 parser.add_argument('--name', default='exp', help='save results to project/name')391 parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment')392 parser.add_argument('--line-thickness', default=3, type=int, help='bounding box thickness (pixels)')393 parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels')394 parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences')395 parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference')396 parser.add_argument('--vid-stride', type=int, default=1, help='video frame-rate stride')397 parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference')398 opt = parser.parse_args()399 opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1 # expand400 401 return opt402 403 404def main(opt):405 # check_requirements(requirements=ROOT / 'requirements.txt', exclude=('tensorboard', 'thop'))406 run_strongsort(**vars(opt))407 408 409if __name__ == "__main__":410 opt = parse_opt()411 main(opt)