def create(name, pretrained, channels, classes, autoshape): """Creates a specified YOLOv3 model Arguments: name (str): name of model, i.e. 'yolov3' pretrained (bool): load pretrained weights into the model channels (int): number of input channels classes (int): number of model classes Returns: pytorch model """ try: cfg = list((Path(__file__).parent / 'models').rglob(f'{name}.yaml'))[0] # model.yaml path model = Model(cfg, channels, classes) if pretrained: fname = f'{name}.pt' # checkpoint filename attempt_download(fname) # download if not found locally ckpt = torch.load(fname, map_location=torch.device('cpu')) # load msd = model.state_dict() # model state_dict csd = ckpt['model'].float().state_dict() # checkpoint state_dict as FP32 csd = {k: v for k, v in csd.items() if msd[k].shape == v.shape} # filter model.load_state_dict(csd, strict=False) # load if len(ckpt['model'].names) == classes: model.names = ckpt['model'].names # set class names attribute if autoshape: model = model.autoshape() # for file/URI/PIL/cv2/np inputs and NMS device = select_device('0' if torch.cuda.is_available() else 'cpu') # default to GPU if available return model.to(device) except Exception as e: help_url = 'https://github.com/ultralytics/yolov5/issues/36' s = 'Cache maybe be out of date, try force_reload=True. See %s for help.' % help_url raise Exception(s) from e
def custom(path_or_model='path/to/model.pt', autoshape=True, verbose=True): """YOLOv5-custom model https://github.com/ultralytics/yolov5 Arguments (3 options): path_or_model (str): 'path/to/model.pt' path_or_model (dict): torch.load('path/to/model.pt') path_or_model (nn.Module): torch.load('path/to/model.pt')['model'] Returns: pytorch model """ set_logging(verbose=verbose) model = torch.load(path_or_model) if isinstance( path_or_model, str) else path_or_model # load checkpoint if isinstance(model, dict): model = model['ema' if model.get('ema') else 'model'] # load model hub_model = Model(model.yaml).to(next(model.parameters()).device) # create hub_model.load_state_dict(model.float().state_dict()) # load state_dict hub_model.names = model.names # class names if autoshape: hub_model = hub_model.autoshape( ) # for file/URI/PIL/cv2/np inputs and NMS device = select_device('0' if torch.cuda.is_available() else 'cpu') # default to GPU if available return hub_model.to(device)
def test_2(): from models.yolo import Model device = 'cuda:0' model_path = '/home/lintao/jobs/logo/yolov5-4.0/yolov5s.pt' # model = torch.load(model_path)['model'].float() model = Model( '/home/lintao/jobs/logo/yolov5-4.0/models/yolov5s.yaml').eval() model.to(device) model.model[-1].export = True x = torch.ones((1, 3, 640, 640)).float().cuda() y = model(x) save_path = model_path.replace('.pt', '.onnx') # filename torch_2_onnx(model, x, input_names=['images'], output_names=['classes'], save_path=save_path)
def create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True): """Creates a specified YOLOv5 model Arguments: name (str): name of model, i.e. 'yolov5s' pretrained (bool): load pretrained weights into the model channels (int): number of input channels classes (int): number of model classes autoshape (bool): apply YOLOv5 .autoshape() wrapper to model verbose (bool): print all information to screen Returns: YOLOv5 pytorch model """ set_logging(verbose=verbose) fname = Path(name).with_suffix('.pt') # checkpoint filename try: if pretrained and channels == 3 and classes == 80: model = attempt_load( fname, map_location=torch.device('cpu')) # download/load FP32 model else: cfg = list((Path(__file__).parent / 'models').rglob(f'{name}.yaml'))[0] # model.yaml path model = Model(cfg, channels, classes) # create model if pretrained: attempt_download(fname) # download if not found locally ckpt = torch.load(fname, map_location=torch.device('cpu')) # load msd = model.state_dict() # model state_dict csd = ckpt['model'].float().state_dict( ) # checkpoint state_dict as FP32 csd = { k: v for k, v in csd.items() if msd[k].shape == v.shape } # filter model.load_state_dict(csd, strict=False) # load if len(ckpt['model'].names) == classes: model.names = ckpt[ 'model'].names # set class names attribute if autoshape: model = model.autoshape() # for file/URI/PIL/cv2/np inputs and NMS device = select_device('0' if torch.cuda.is_available() else 'cpu') # default to GPU if available return model.to(device) except Exception as e: help_url = 'https://github.com/ultralytics/yolov5/issues/36' s = 'Cache may be out of date, try `force_reload=True`. See %s for help.' % help_url raise Exception(s) from e
def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): """Creates a specified YOLOv5 model Arguments: name (str): name of model, i.e. 'yolov5s' pretrained (bool): load pretrained weights into the model channels (int): number of input channels classes (int): number of model classes autoshape (bool): apply YOLOv5 .autoshape() wrapper to model verbose (bool): print all information to screen device (str, torch.device, None): device to use for model parameters Returns: YOLOv5 pytorch model """ from pathlib import Path from models.common import AutoShape from models.experimental import attempt_load from models.yolo import Model from utils.downloads import attempt_download from utils.general import check_requirements, intersect_dicts, set_logging from utils.torch_utils import select_device file = Path(__file__).resolve() check_requirements(exclude=('tensorboard', 'thop', 'opencv-python')) set_logging(verbose=verbose) save_dir = Path('') if str(name).endswith('.pt') else file.parent path = (save_dir / name).with_suffix('.pt') # checkpoint path try: device = select_device(('0' if torch.cuda.is_available() else 'cpu') if device is None else device) if pretrained and channels == 3 and classes == 80: model = attempt_load(path, map_location=device) # download/load FP32 model else: cfg = list((Path(__file__).parent / 'models').rglob(f'{name}.yaml'))[0] # model.yaml path model = Model(cfg, channels, classes) # create model if pretrained: ckpt = torch.load(attempt_download(path), map_location=device) # load csd = ckpt['model'].float().state_dict() # checkpoint state_dict as FP32 csd = intersect_dicts(csd, model.state_dict(), exclude=['anchors']) # intersect model.load_state_dict(csd, strict=False) # load if len(ckpt['model'].names) == classes: model.names = ckpt['model'].names # set class names attribute if autoshape: model = AutoShape(model) # for file/URI/PIL/cv2/np inputs and NMS return model.to(device) except Exception as e: help_url = 'https://github.com/ultralytics/yolov5/issues/36' s = 'Cache may be out of date, try `force_reload=True`. See %s for help.' % help_url raise Exception(s) from e
def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): """Creates a specified YOLOv3 model Arguments: name (str): name of model, i.e. 'yolov3' pretrained (bool): load pretrained weights into the model channels (int): number of input channels classes (int): number of model classes autoshape (bool): apply YOLOv3 .autoshape() wrapper to model verbose (bool): print all information to screen device (str, torch.device, None): device to use for model parameters Returns: YOLOv3 pytorch model """ from pathlib import Path from models.yolo import Model, attempt_load from utils.general import check_requirements, set_logging from utils.google_utils import attempt_download from utils.torch_utils import select_device check_requirements(Path(__file__).parent / 'requirements.txt', exclude=('tensorboard', 'pycocotools', 'thop')) set_logging(verbose=verbose) fname = Path(name).with_suffix('.pt') # checkpoint filename try: if pretrained and channels == 3 and classes == 80: model = attempt_load(fname, map_location=torch.device('cpu')) # download/load FP32 model else: cfg = list((Path(__file__).parent / 'models').rglob(f'{name}.yaml'))[0] # model.yaml path model = Model(cfg, channels, classes) # create model if pretrained: ckpt = torch.load(attempt_download(fname), map_location=torch.device('cpu')) # load msd = model.state_dict() # model state_dict csd = ckpt['model'].float().state_dict() # checkpoint state_dict as FP32 csd = {k: v for k, v in csd.items() if msd[k].shape == v.shape} # filter model.load_state_dict(csd, strict=False) # load if len(ckpt['model'].names) == classes: model.names = ckpt['model'].names # set class names attribute if autoshape: model = model.autoshape() # for file/URI/PIL/cv2/np inputs and NMS device = select_device('0' if torch.cuda.is_available() else 'cpu') if device is None else torch.device(device) return model.to(device) except Exception as e: help_url = 'https://github.com/ultralytics/yolov5/issues/36' s = 'Cache may be out of date, try `force_reload=True`. See %s for help.' % help_url raise Exception(s) from e
# Load PyTorch model # ckpt = torch.load(opt.weights, map_location=lambda storage, loc: storage.cuda())['model'].float() with open('./models/configs/yolo3d_5m.yaml') as f: cfg = yaml.load(f, Loader=yaml.FullLoader) cfg.update(opt.__dict__) print(cfg) model = Model(cfg) model.eval() checkpointer = model_utils.CheckPointer(model, device=device) checkpointer.load(opt.weights, load_solver=False) # model.to(device) if half: model.half() # to FP16 img = img.half() else: model.to(torch.float32) model.to(device) # _print_weights_(model) _ = model(img) # dry run model.model[-1].export = True # set Detect() layer export=True # TorchScript export try: print('\nStarting TorchScript export with torch %s...' % torch.__version__) f = opt.weights.replace('.pt', '.torchscript.pt') # filename ts = torch.jit.trace(model, img) ts.save(f) print('TorchScript export success, saved as %s' % f) except Exception as e: print('TorchScript export failure: %s' % e)
def detect(): weights, imgsz, output = opt.checkpoint, opt.img_size, opt.output_video # Initialize device = torch_utils.select_device(opt.device) half = device.type != 'cpu' # half precision only supported on CUDA with open(opt.project) as f: data_dict = yaml.load(f, Loader=yaml.FullLoader) names = data_dict['names'] dest_object = data_dict[ 'dest_object'] if 'dest_object' in data_dict else names if 'name_map' in data_dict: name_map = data_dict['name_map'] else: name_map = defaultdict() for x in dest_object: name_map[x] = x draw_names = set() for k, v in name_map.items(): draw_names.add(v) draw_names = list(draw_names) colors = get_all_colors(len(draw_names)) model = Model(data_dict).to(device) model.load_state_dict(torch.load(weights, map_location=device)) model.to(device).eval() if half: model.half() if opt.video_type == 'camera': rtsp = 'rtsp://*****:*****@' + opt.input_video + ':554/h264/chCH/sub/av_stream' cap = cv2.VideoCapture(rtsp) elif opt.video_type == 'video': cap = cv2.VideoCapture(opt.input_video) else: print('Input video type ERROR!') return fps = int(cap.get(cv2.CAP_PROP_FPS)) width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) print(fps, width, height) fourcc = cv2.VideoWriter_fourcc('m', 'p', '4', 'v') outstream = cv2.VideoWriter(output, fourcc, fps, (width, height)) t0 = time.time() img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img _ = model(img.half() if half else img ) if device.type != 'cpu' else None # run once while (cap.isOpened()): ret, im0s = cap.read() if ret: pred = forward_one(model, im0s, imgsz, device, half, opt) for i, det in enumerate(pred): # detections per image if det is not None and len(det): for *xyxy, conf, cls in det: if names[int(cls)] in dest_object: draw_str = name_map[names[int(cls)]] color = colors[draw_names.index(draw_str)] label = '%s %.2f' % (draw_str, conf) xmin, ymin, xmax, ymax = xyxy cv2.rectangle(im0s, (xmin, ymin), (xmax, ymax), color=color, thickness=1) cv2.putText(im0s, label, (xmin, ymax), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 1) if opt.need_view: cv2.imshow('real-time', im0s) key = cv2.waitKey(delay=1) if key == ord('q'): break else: outstream.write(im0s) cv2.destroyAllWindows() print('Done. (%.3fs)' % (time.time() - t0))
def detect(): source, weights, view_img, save_txt, imgsz = \ opt.input_images, opt.checkpoint, opt.view_img, opt.save_txt, opt.img_size webcam = source == '0' or source.startswith('rtsp') or source.startswith('http') or source.endswith('.txt') # Initialize device = torch_utils.select_device(opt.device) if os.path.exists(opt.output_images): shutil.rmtree(opt.output_images) # delete output folder os.makedirs(opt.output_images) # make new output folder half = device.type != 'cpu' # half precision only supported on CUDA with open(opt.project) as f: data_dict = yaml.load(f, Loader=yaml.FullLoader) model = Model(data_dict).to(device) model.load_state_dict(torch.load(weights, map_location=device)) model.names = data_dict['names'] # Load model #google_utils.attempt_download(weights) #model = torch.load(weights, map_location=device)['model'].float() # load to FP32 #model = torch.load(weights, map_location=device).float() # load to FP32 # torch.save(torch.load(weights, map_location=device), weights) # update model if SourceChangeWarning # model.fuse() model.to(device).eval() if half: model.half() # to FP16 # 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(len(names))] colors = get_all_colors(len(names)) # Run inference t0 = time.time() img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img _ = model(img.half() if half else img) if device.type != 'cpu' else None # run once #for path, img, im0s, vid_cap in dataset: for f in os.listdir(source): t1 = time.time() path = os.path.join(source, f) im0s = cv2.imread(path) img = letterbox(im0s, new_shape=imgsz)[0] img = img[:, :, ::-1].transpose(2, 0, 1) img = np.ascontiguousarray(img) 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 = torch_utils.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 = torch_utils.time_synchronized() # Process detections for i, det in enumerate(pred): # detections per image p, s, im0 = path, '', im0s s += '%gx%g ' % img.shape[2:] # print string gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh 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 += '%g %ss, ' % (n, names[int(c)]) # add to string # Write results for *xyxy, conf, cls in det: label = '%s %.2f' % (names[int(cls)], conf) xmin, ymin, xmax, ymax = xyxy color = colors[int(cls)] cv2.rectangle(im0, (xmin, ymin), (xmax, ymax), color=color, thickness=2) cv2.putText(im0, label, (xmin, ymax), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 1) cv2.imwrite(os.path.join(opt.output_images, f), im0) print('%s Done. (%.3fs)' % (path, time.time() - t1)) print('Done. (%.3fs)' % (time.time() - t0))
def test(data, weights=None, batch_size=16, imgsz=640, conf_thres=0.001, iou_thres=0.6, # for NMS save_json=False, single_cls=False, augment=False, verbose=False, model=None, dataloader=None, fast=False): # Initialize/load model and set device if model is None: training = False device = torch_utils.select_device(opt.device, batch_size=batch_size) half = device.type != 'cpu' # half precision only supported on CUDA # Remove previous for f in glob.glob('test_batch*.jpg'): os.remove(f) # Load model google_utils.attempt_download(weights) # model = torch.load(weights, map_location=device)['state_dict'].float() # load to FP32 model = Model(model_cfg='/home/ai/yulu/yolov5/models/yolov5s.yaml') if os.path.exists(opt.weights): ckpt = torch.load(opt.weights, map_location=device) state_dict = {key: ckpt['state_dict'][key] for key in model.state_dict().keys()} model.load_state_dict(state_dict) # torch_utils.model_info(model) model.fuse() model.to(device) if half: model.half() # to FP16 # Multi-GPU disabled, incompatible with .half() # if device.type != 'cpu' and torch.cuda.device_count() > 1: # model = nn.DataParallel(model) else: # called by train.py training = True device = next(model.parameters()).device # get model device # half disabled https://github.com/ultralytics/yolov5/issues/99 half = False # device.type != 'cpu' and torch.cuda.device_count() == 1 if half: model.half() # to FP16 # Configure model.eval() with open(data) as f: data = yaml.load(f, Loader=yaml.FullLoader) # model dict nc = 1 if single_cls else int(data['nc']) # number of classes iouv = torch.linspace(0.5, 0.95, 10).to(device) # iou vector for [email protected]:0.95 # iouv = iouv[0].view(1) # comment for [email protected]:0.95 niou = iouv.numel() # Dataloader if dataloader is None: # not training img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img _ = model(img.half() if half else img) if device.type != 'cpu' else None # run once fast |= conf_thres > 0.001 # enable fast mode path = data['test'] if opt.task == 'test' else data['val'] # path to val/test images dataset = LoadImagesAndLabels(path, imgsz, batch_size, rect=True, # rectangular inference single_cls=opt.single_cls, # single class mode pad=0.5) # padding batch_size = min(batch_size, len(dataset)) nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, 8]) # number of workers dataloader = DataLoader(dataset, batch_size=batch_size, num_workers=nw, pin_memory=True, collate_fn=dataset.collate_fn) seen = 0 # names = model.names if hasattr(model, 'names') else model.module.names coco91class = coco80_to_coco91_class() s = ('%20s' + '%12s' * 6) % ('Class', 'Images', 'Targets', 'P', 'R', '[email protected]', '[email protected]:.95') p, r, f1, mp, mr, map50, map, t0, t1 = 0., 0., 0., 0., 0., 0., 0., 0., 0. loss = torch.zeros(3, device=device) jdict, stats, ap, ap_class = [], [], [], [] for batch_i, (img, targets, paths, shapes) in enumerate(tqdm(dataloader, desc=s)): img = 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 targets = targets.to(device) nb, _, height, width = img.shape # batch size, channels, height, width whwh = torch.Tensor([width, height, width, height]).to(device) # Disable gradients with torch.no_grad(): # Run model t = torch_utils.time_synchronized() inf_out, train_out = model(img, augment=augment) # inference and training outputs t0 += torch_utils.time_synchronized() - t # Compute loss if training: # if model has loss hyperparameters loss += compute_loss([x.float() for x in train_out], targets, model)[1][:3] # GIoU, obj, cls # Run NMS t = torch_utils.time_synchronized() output = non_max_suppression(inf_out, conf_thres=conf_thres, iou_thres=iou_thres, fast=fast) # ?????????????? t1 += torch_utils.time_synchronized() - t # Statistics per image for si, pred in enumerate(output): labels = targets[targets[:, 0] == si, 1:] nl = len(labels) tcls = labels[:, 0].tolist() if nl else [] # target class seen += 1 if pred is None: if nl: stats.append((torch.zeros(0, niou, dtype=torch.bool), torch.Tensor(), torch.Tensor(), tcls)) continue # Append to text file # with open('test.txt', 'a') as file: # [file.write('%11.5g' * 7 % tuple(x) + '\n') for x in pred] # Clip boxes to image bounds clip_coords(pred, (height, width)) # Append to pycocotools JSON dictionary if save_json: # [{"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}, ... image_id = int(Path(paths[si]).stem.split('_')[-1]) box = pred[:, :4].clone() # xyxy scale_coords(img[si].shape[1:], box, shapes[si][0], shapes[si][1]) # to original shape box = xyxy2xywh(box) # xywh box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner for p, b in zip(pred.tolist(), box.tolist()): jdict.append({'image_id': image_id, 'category_id': coco91class[int(p[5])], 'bbox': [round(x, 3) for x in b], 'score': round(p[4], 5)}) # Assign all predictions as incorrect correct = torch.zeros(pred.shape[0], niou, dtype=torch.bool, device=device) if nl: detected = [] # target indices tcls_tensor = labels[:, 0] # target boxes tbox = xywh2xyxy(labels[:, 1:5]) * whwh # Per target class for cls in torch.unique(tcls_tensor): ti = (cls == tcls_tensor).nonzero().view(-1) # prediction indices pi = (cls == pred[:, 5]).nonzero().view(-1) # target indices # Search for detections if pi.shape[0]: # Prediction to target ious ious, i = box_iou(pred[pi, :4], tbox[ti]).max(1) # best ious, indices # Append detections for j in (ious > iouv[0]).nonzero(): d = ti[i[j]] # detected target if d not in detected: detected.append(d) correct[pi[j]] = ious[j] > iouv # iou_thres is 1xn if len(detected) == nl: # all targets already located in image break # Append statistics (correct, conf, pcls, tcls) stats.append((correct.cpu(), pred[:, 4].cpu(), pred[:, 5].cpu(), tcls)) # Plot images if batch_i < 1: f = 'test_batch%g_gt.jpg' % batch_i # filename # plot_images(img, targets, paths, f, names) # ground truth plot_images(img, targets, paths, f, 'yolov5s') f = 'test_batch%g_pred.jpg' % batch_i # plot_images(img, output_to_target(output, width, height), paths, f, names) # predictions plot_images(img, output_to_target(output, width, height), paths, f, 'yolov5s') # Compute statistics stats = [np.concatenate(x, 0) for x in zip(*stats)] # to numpy if len(stats): p, r, ap, f1, ap_class = ap_per_class(*stats) p, r, ap50, ap = p[:, 0], r[:, 0], ap[:, 0], ap.mean(1) # [P, R, [email protected], [email protected]:0.95] mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean() nt = np.bincount(stats[3].astype(np.int64), minlength=nc) # number of targets per class else: nt = torch.zeros(1) # Print results pf = '%20s' + '%12.3g' * 6 # print format print(pf % ('all', seen, nt.sum(), mp, mr, map50, map)) # Print results per class if verbose and nc > 1 and len(stats): for i, c in enumerate(ap_class): print(pf % (names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i])) # Print speeds t = tuple(x / seen * 1E3 for x in (t0, t1, t0 + t1)) + (imgsz, imgsz, batch_size) # tuple if not training: print('Speed: %.1f/%.1f/%.1f ms inference/NMS/total per %gx%g image at batch-size %g' % t) # Save JSON if save_json and map50 and len(jdict): imgIds = [int(Path(x).stem.split('_')[-1]) for x in dataloader.dataset.img_files] f = 'detections_val2017_%s_results.json' % \ (weights.split(os.sep)[-1].replace('.pt', '') if weights else '') # filename print('\nCOCO mAP with pycocotools... saving %s...' % f) with open(f, 'w') as file: json.dump(jdict, file) try: from pycocotools.coco import COCO from pycocotools.cocoeval import COCOeval # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb cocoGt = COCO(glob.glob('../coco/annotations/instances_val*.json')[0]) # initialize COCO ground truth api cocoDt = cocoGt.loadRes(f) # initialize COCO pred api cocoEval = COCOeval(cocoGt, cocoDt, 'bbox') cocoEval.params.imgIds = imgIds # image IDs to evaluate cocoEval.evaluate() cocoEval.accumulate() cocoEval.summarize() map, map50 = cocoEval.stats[:2] # update results ([email protected]:0.95, [email protected]) except: print('WARNING: pycocotools must be installed with numpy==1.17 to run correctly. ' 'See https://github.com/cocodataset/cocoapi/issues/356') # Return results maps = np.zeros(nc) + map for i, c in enumerate(ap_class): maps[c] = ap[i] return (mp, mr, map50, map, *(loss.cpu() / len(dataloader)).tolist()), maps, t
def train(hyp, opt, device, tb_writer=None, wandb=None): logger.info(f'Hyperparameters {hyp}') save_dir, epochs, batch_size, total_batch_size, weights, rank = \ Path(opt.save_dir), opt.epochs, opt.batch_size, opt.total_batch_size, opt.weights, opt.global_rank # Directories wdir = save_dir / 'weights' wdir.mkdir(parents=True, exist_ok=True) # make dir last = wdir / 'last.pt' best = wdir / 'best.pt' results_file = save_dir / 'results.txt' # Save run settings with open(save_dir / 'hyp.yaml', 'w') as f: yaml.dump(hyp, f, sort_keys=False) with open(save_dir / 'opt.yaml', 'w') as f: yaml.dump(vars(opt), f, sort_keys=False) # Configure plots = not opt.evolve # create plots cuda = device.type != 'cpu' init_seeds(2 + rank) with open(opt.data) as f: data_dict = yaml.load(f, Loader=yaml.FullLoader) # data dict with torch_distributed_zero_first(rank): check_dataset(data_dict) # check train_path = data_dict['train'] test_path = data_dict['val'] nc, names = (1, ['item']) if opt.single_cls else (int( data_dict['nc']), data_dict['names']) # number classes, names assert len(names) == nc, '%g names found for nc=%g dataset in %s' % ( len(names), nc, opt.data) # check # Model pretrained = weights.endswith('.pt') if pretrained: with torch_distributed_zero_first(rank): attempt_download(weights) # download if not found locally ckpt = torch.load(weights, map_location=device) # load checkpoint if hyp.get('anchors'): ckpt['model'].yaml['anchors'] = round( hyp['anchors']) # force autoanchor model = Model(opt.cfg or ckpt['model'].yaml, ch=3, nc=nc).to(device) # create exclude = ['anchor'] if opt.cfg or hyp.get('anchors') else [ ] # exclude keys state_dict = ckpt['model'].float().state_dict() # to FP32 state_dict = intersect_dicts(state_dict, model.state_dict(), exclude=exclude) # intersect model.load_state_dict(state_dict, strict=False) # load logger.info( 'Transferred %g/%g items from %s' % (len(state_dict), len(model.state_dict()), weights)) # report else: model = Model(opt.cfg, ch=3, nc=nc).to(device) # create # replace the original model to a quantized model if opt.quantization: print(model) model = model.cpu() print("opt.cfg => {}".format(opt.cfg)) if opt.cfg.split("/")[-1].split(".")[-2] == "yolov3-tiny": print("model: yolov3-tiny") layer_layer_count = 12 else: layer_layer_count = 100000 # quantize model model = tflite.replace(model=model, quantization_bits=opt.quantization_bits, scale_bits=opt.scale_bits, bias_bits=opt.bias_bits, layer_layer_count=layer_layer_count) model = model.to(device) # Freeze freeze = [] # parameter names to freeze (full or partial) for k, v in model.named_parameters(): v.requires_grad = True # train all layers if any(x in k for x in freeze): print('freezing %s' % k) v.requires_grad = False # Optimizer nbs = 64 # nominal batch size accumulate = max(round(nbs / total_batch_size), 1) # accumulate loss before optimizing hyp['weight_decay'] *= total_batch_size * accumulate / nbs # scale weight_decay pg0, pg1, pg2 = [], [], [] # optimizer parameter groups for k, v in model.named_modules(): if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter): pg2.append(v.bias) # biases if isinstance(v, nn.BatchNorm2d): pg0.append(v.weight) # no decay elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter): pg1.append(v.weight) # apply decay if opt.adam: optimizer = optim.Adam(pg0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999)) # adjust beta1 to momentum else: optimizer = optim.SGD(pg0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True) optimizer.add_param_group({ 'params': pg1, 'weight_decay': hyp['weight_decay'] }) # add pg1 with weight_decay optimizer.add_param_group({'params': pg2}) # add pg2 (biases) logger.info('Optimizer groups: %g .bias, %g conv.weight, %g other' % (len(pg2), len(pg1), len(pg0))) del pg0, pg1, pg2 # Scheduler https://arxiv.org/pdf/1812.01187.pdf # https://pytorch.org/docs/stable/_modules/torch/optim/lr_scheduler.html#OneCycleLR lf = lambda x: ((1 + math.cos(x * math.pi / epochs)) / 2) * (1 - hyp[ 'lrf']) + hyp['lrf'] # cosine scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf) # plot_lr_scheduler(optimizer, scheduler, epochs) # Logging if wandb and wandb.run is None: opt.hyp = hyp # add hyperparameters wandb_run = wandb.init( config=opt, resume="allow", project='YOLOv3' if opt.project == 'runs/train' else Path(opt.project).stem, name=save_dir.stem, id=ckpt.get('wandb_id') if 'ckpt' in locals() else None) loggers = {'wandb': wandb} # loggers dict # Resume start_epoch, best_fitness = 0, 0.0 if pretrained: # Optimizer if ckpt['optimizer'] is not None: optimizer.load_state_dict(ckpt['optimizer']) best_fitness = ckpt['best_fitness'] # Results if ckpt.get('training_results') is not None: with open(results_file, 'w') as file: file.write(ckpt['training_results']) # write results.txt # Epochs start_epoch = ckpt['epoch'] + 1 if opt.resume: assert start_epoch > 0, '%s training to %g epochs is finished, nothing to resume.' % ( weights, epochs) if epochs < start_epoch: logger.info( '%s has been trained for %g epochs. Fine-tuning for %g additional epochs.' % (weights, ckpt['epoch'], epochs)) epochs += ckpt['epoch'] # finetune additional epochs del ckpt, state_dict # Image sizes gs = int(max(model.stride)) # grid size (max stride) imgsz, imgsz_test = [check_img_size(x, gs) for x in opt.img_size ] # verify imgsz are gs-multiples # DP mode if cuda and rank == -1 and torch.cuda.device_count() > 1: model = torch.nn.DataParallel(model) # SyncBatchNorm if opt.sync_bn and cuda and rank != -1: model = torch.nn.SyncBatchNorm.convert_sync_batchnorm(model).to(device) logger.info('Using SyncBatchNorm()') # EMA ema = ModelEMA(model) if rank in [-1, 0] else None # DDP mode if cuda and rank != -1: model = DDP(model, device_ids=[opt.local_rank], output_device=opt.local_rank) # xiezheng add # Set hook to store Conv results for No register_hook(model, hook_conv_results) logger.info("len(oaq_conv_result)={},\noaq_conv_result={}".format( len(oaq_conv_result), oaq_conv_result)) # Trainloader dataloader, dataset = create_dataloader(train_path, imgsz, batch_size, gs, opt, hyp=hyp, augment=True, cache=opt.cache_images, rect=opt.rect, rank=rank, world_size=opt.world_size, workers=opt.workers, image_weights=opt.image_weights) mlc = np.concatenate(dataset.labels, 0)[:, 0].max() # max label class nb = len(dataloader) # number of batches assert mlc < nc, 'Label class %g exceeds nc=%g in %s. Possible class labels are 0-%g' % ( mlc, nc, opt.data, nc - 1) # Process 0 if rank in [-1, 0]: ema.updates = start_epoch * nb // accumulate # set EMA updates testloader = create_dataloader( test_path, imgsz_test, total_batch_size, gs, opt, # testloader hyp=hyp, cache=opt.cache_images and not opt.notest, rect=True, rank=-1, world_size=opt.world_size, workers=opt.workers, pad=0.5)[0] if not opt.resume: labels = np.concatenate(dataset.labels, 0) c = torch.tensor(labels[:, 0]) # classes # cf = torch.bincount(c.long(), minlength=nc) + 1. # frequency # model._initialize_biases(cf.to(device)) if plots: Thread(target=plot_labels, args=(labels, save_dir, loggers), daemon=True).start() if tb_writer: tb_writer.add_histogram('classes', c, 0) # Anchors if not opt.noautoanchor: check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz) # Model parameters hyp['cls'] *= nc / 80. # scale coco-tuned hyp['cls'] to current dataset model.nc = nc # attach number of classes to model model.hyp = hyp # attach hyperparameters to model model.gr = 1.0 # iou loss ratio (obj_loss = 1.0 or iou) model.class_weights = labels_to_class_weights(dataset.labels, nc).to( device) # attach class weights model.names = names # Start training t0 = time.time() nw = max(round(hyp['warmup_epochs'] * nb), 1000) # number of warmup iterations, max(3 epochs, 1k iterations) # nw = min(nw, (epochs - start_epoch) / 2 * nb) # limit warmup to < 1/2 of training maps = np.zeros(nc) # mAP per class results = (0, 0, 0, 0, 0, 0, 0 ) # P, R, [email protected], [email protected], val_loss(box, obj, cls) scheduler.last_epoch = start_epoch - 1 # do not move # scaler = amp.GradScaler(enabled=cuda) logger.info('Image sizes %g train, %g test\n' 'Using %g dataloader workers\nLogging results to %s\n' 'Starting training for %g epochs...' % (imgsz, imgsz_test, dataloader.num_workers, save_dir, epochs)) for epoch in range( start_epoch, epochs ): # epoch ------------------------------------------------------------------ model.train() if opt.quantization: tflite.unfreeze_model(model) # unfreeze the activation range # Update image weights (optional) if opt.image_weights: # Generate indices if rank in [-1, 0]: cw = model.class_weights.cpu().numpy() * ( 1 - maps)**2 # class weights iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw) # image weights dataset.indices = random.choices( range(dataset.n), weights=iw, k=dataset.n) # rand weighted idx # Broadcast if DDP if rank != -1: indices = (torch.tensor(dataset.indices) if rank == 0 else torch.zeros(dataset.n)).int() dist.broadcast(indices, 0) if rank != 0: dataset.indices = indices.cpu().numpy() # Update mosaic border # b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs) # dataset.mosaic_border = [b - imgsz, -b] # height, width borders mloss = torch.zeros(4, device=device) # mean losses if rank != -1: dataloader.sampler.set_epoch(epoch) pbar = enumerate(dataloader) logger.info( ('\n' + '%10s' * 8) % ('Epoch', 'gpu_mem', 'box', 'obj', 'cls', 'total', 'targets', 'img_size')) if rank in [-1, 0]: pbar = tqdm(pbar, total=nb) # progress bar optimizer.zero_grad() for i, ( imgs, targets, paths, _ ) in pbar: # batch ------------------------------------------------------------- ni = i + nb * epoch # number integrated batches (since train start) imgs = imgs.to(device, non_blocking=True).float( ) / 255.0 # uint8 to float32, 0-255 to 0.0-1.0 # Warmup if ni <= nw: xi = [0, nw] # x interp # model.gr = np.interp(ni, xi, [0.0, 1.0]) # iou loss ratio (obj_loss = 1.0 or iou) accumulate = max( 1, np.interp(ni, xi, [1, nbs / total_batch_size]).round()) for j, x in enumerate(optimizer.param_groups): # bias lr falls from 0.1 to lr0, all other lrs rise from 0.0 to lr0 x['lr'] = np.interp(ni, xi, [ hyp['warmup_bias_lr'] if j == 2 else 0.0, x['initial_lr'] * lf(epoch) ]) if 'momentum' in x: x['momentum'] = np.interp( ni, xi, [hyp['warmup_momentum'], hyp['momentum']]) # Multi-scale if opt.multi_scale: sz = random.randrange(imgsz * 0.5, imgsz * 1.5 + gs) // gs * gs # size sf = sz / max(imgs.shape[2:]) # scale factor if sf != 1: ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:] ] # new shape (stretched to gs-multiple) imgs = F.interpolate(imgs, size=ns, mode='bilinear', align_corners=False) # Forward # with amp.autocast(enabled=cuda): pred = model(imgs) # forward loss, loss_items = compute_loss(pred, targets.to(device), model) # loss scaled by batch_size if rank != -1: loss *= opt.world_size # gradient averaged between devices in DDP mode # Backward # scaler.scale(loss).backward() loss.backward() # Optimize if ni % accumulate == 0: # scaler.step(optimizer) # optimizer.step # scaler.update() optimizer.step() optimizer.zero_grad() if ema: ema.update(model) # Print if rank in [-1, 0]: mloss = (mloss * i + loss_items) / (i + 1 ) # update mean losses mem = '%.3gG' % (torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0) # (GB) s = ('%10s' * 2 + '%10.4g' * 6) % ('%g/%g' % (epoch, epochs - 1), mem, *mloss, targets.shape[0], imgs.shape[-1]) pbar.set_description(s) # Plot if plots and ni < 3: f = save_dir / f'train_batch{ni}.jpg' # filename Thread(target=plot_images, args=(imgs, targets, paths, f), daemon=True).start() # if tb_writer: # tb_writer.add_image(f, result, dataformats='HWC', global_step=epoch) # tb_writer.add_graph(model, imgs) # add model to tensorboard elif plots and ni == 3 and wandb: wandb.log({ "Mosaics": [ wandb.Image(str(x), caption=x.name) for x in save_dir.glob('train*.jpg') ] }) # xiezheng add # calculating No and updating alpha are executed every M=50 steps if ni > nw and (ni - nw) % opt.OAQ_m == 0: logger.info("ni={}, calculate No and updata alpha!".format(ni)) No = calculate_No(device, model, oaq_conv_result, opt.conv_accumulator_bits, logger) lr_max = hyp['lr0'] lr_curr = [x['lr'] for x in optimizer.param_groups][1] # weight lr iteration_batch_size = imgs.size(0) update_alpha(device, model, No, iteration_batch_size, lr_max, lr_curr, logger) # Test: calculate No and updata alpha # logger.info("ni={}, calculate No and updata alpha!".format(ni)) # No = calculate_No(device, model, oaq_conv_result, opt.conv_accumulator_bits, logger) # lr_max = hyp['lr0'] # lr_curr = [x['lr'] for x in optimizer.param_groups][1] # weight lr # iteration_batch_size = imgs.size(0) # update_alpha(device, model, No, iteration_batch_size, lr_max, lr_curr, logger) # xiezheng add # clear hook results oaq_conv_result.clear() # end batch ------------------------------------------------------------------------------------------------ # end epoch ---------------------------------------------------------------------------------------------------- # Scheduler lr = [x['lr'] for x in optimizer.param_groups] # for tensorboard scheduler.step() # DDP process 0 or single-GPU if rank in [-1, 0]: # mAP if ema: ema.update_attr( model, include=['yaml', 'nc', 'hyp', 'gr', 'names', 'stride']) final_epoch = epoch + 1 == epochs if not opt.notest or final_epoch: # Calculate mAP if opt.quantization: tflite.freeze_model(ema.ema) results, maps, times = test.test( opt.data, batch_size=total_batch_size, imgsz=imgsz_test, model=ema.ema, single_cls=opt.single_cls, dataloader=testloader, save_dir=save_dir, plots=plots and final_epoch, log_imgs=opt.log_imgs if wandb else 0) # Write with open(results_file, 'a') as f: f.write( s + '%10.4g' * 7 % results + '\n') # P, R, [email protected], [email protected], val_loss(box, obj, cls) if len(opt.name) and opt.bucket: os.system('gsutil cp %s gs://%s/results/results%s.txt' % (results_file, opt.bucket, opt.name)) # Log tags = [ 'train/box_loss', 'train/obj_loss', 'train/cls_loss', # train loss 'metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95', 'val/box_loss', 'val/obj_loss', 'val/cls_loss', # val loss 'x/lr0', 'x/lr1', 'x/lr2' ] # params for x, tag in zip(list(mloss[:-1]) + list(results) + lr, tags): if tb_writer: tb_writer.add_scalar(tag, x, epoch) # tensorboard if wandb: wandb.log({tag: x}) # W&B # Update best mAP fi = fitness(np.array(results).reshape( 1, -1)) # weighted combination of [P, R, [email protected], [email protected]] if fi > best_fitness: best_fitness = fi # Save model save = (not opt.nosave) or (final_epoch and not opt.evolve) if save: with open(results_file, 'r') as f: # create checkpoint ckpt = { 'epoch': epoch, 'best_fitness': best_fitness, 'training_results': f.read(), 'model': ema.ema, 'optimizer': None if final_epoch else optimizer.state_dict(), 'wandb_id': wandb_run.id if wandb else None } # Save last, best and delete torch.save(ckpt, last) if best_fitness == fi: torch.save(ckpt, best) del ckpt # end epoch ---------------------------------------------------------------------------------------------------- # end training if rank in [-1, 0]: # Strip optimizers for f in [last, best]: if f.exists(): # is *.pt strip_optimizer(f) # strip optimizer os.system('gsutil cp %s gs://%s/weights' % (f, opt.bucket)) if opt.bucket else None # upload # Plots if plots: plot_results(save_dir=save_dir) # save as results.png if wandb: files = [ 'results.png', 'precision_recall_curve.png', 'confusion_matrix.png' ] wandb.log({ "Results": [ wandb.Image(str(save_dir / f), caption=f) for f in files if (save_dir / f).exists() ] }) logger.info('%g epochs completed in %.3f hours.\n' % (epoch - start_epoch + 1, (time.time() - t0) / 3600)) # Test best.pt if opt.data.endswith('coco.yaml') and nc == 80: # if COCO if opt.quantization: model_test = torch.load( best if best.exists() else last, map_location=device)['model'].float().eval() tflite.freeze_model(model_test) else: model_test = attempt_load(best if best.exists() else last, device) results, _, _ = test.test( opt.data, batch_size=total_batch_size, imgsz=imgsz_test, model=model_test, single_cls=opt.single_cls, dataloader=testloader, save_dir=save_dir, save_json=True, # use pycocotools plots=False) # Write with open(results_file, 'a') as f: f.write("Test best.pt:\n") f.write( s + '%10.4g' * 7 % results + '\n') # P, R, [email protected], [email protected], val_loss(box, obj, cls) else: dist.destroy_process_group() wandb.run.finish() if wandb and wandb.run else None torch.cuda.empty_cache() return results
def detect(): weights, imgsz = opt.checkpoint, opt.img_size # Initialize device = torch_utils.select_device(opt.device) if os.path.exists(opt.output_images): shutil.rmtree(opt.output_images) # delete output folder os.makedirs(opt.output_images) # make new output folder half = device.type != 'cpu' # half precision only supported on CUDA with open(opt.project) as f: data_dict = yaml.load(f, Loader=yaml.FullLoader) names = data_dict['names'] dest_object = data_dict[ 'dest_object'] if 'dest_object' in data_dict else names if 'name_map' in data_dict: name_map = data_dict['name_map'] else: name_map = defaultdict() for x in dest_object: name_map[x] = x draw_names = set() for k, v in name_map.items(): draw_names.add(v) draw_names = list(draw_names) colors = get_all_colors(len(draw_names)) model = Model(data_dict).to(device) model.load_state_dict(torch.load(weights, map_location=device)) model.to(device).eval() if half: model.half() # to FP16 t0 = time.time() img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img _ = model(img.half() if half else img ) if device.type != 'cpu' else None # run once for f in os.listdir(opt.input_images): t1 = time.time() path = os.path.join(opt.input_images, f) im0s = cv2.imread(path) pred = forward_one(model, im0s, imgsz, device, half, opt) for i, det in enumerate(pred): # detections per image if det is not None and len(det): for *xyxy, conf, cls in det: if names[int(cls)] in dest_object: draw_str = name_map[names[int(cls)]] color = colors[draw_names.index(draw_str)] label = '%s %.2f' % (draw_str, conf) xmin, ymin, xmax, ymax = xyxy cv2.rectangle(im0s, (xmin, ymin), (xmax, ymax), color=color, thickness=1) cv2.putText(im0s, label, (xmin, ymax), cv2.FONT_HERSHEY_SIMPLEX, 0.6, color, 1) cv2.imwrite(os.path.join(opt.output_images, f), im0s) print('%s Done. (%.3fs)' % (f, time.time() - t1)) print('Done. (%.3fs)' % (time.time() - t0))
def detect(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('http') or source.endswith('.txt') # Initialize device = torch_utils.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 google_utils.attempt_download(weights) # model = torch.load(weights, map_location=device)['model'].float() # load to FP32 from models.yolo import Model model = Model(model_cfg='/home/ai/yulu/yolov5/models/yolov5s.yaml').to(device) if os.path.exists(opt.weights): ckpt = torch.load('/home/ai/yulu/yolov5/weights/best0722.pt', map_location=device) state_dict = {key:ckpt['state_dict'][key] for key in model.state_dict().keys()} model.load_state_dict(state_dict) # ################ # model = torch.load(weights, map_location=device) # ckpt = {'epoch': model['epoch'], # 'best_fitness': model['best_fitness'], # 'training_results': model['training_results'], # 'state_dict': model['model'].state_dict(), # 'optimizer': None} # # # # # # # Save last, best and delete # torch.save(ckpt, 'yolov5_adult_kid.pt') # ################ # model.fuse() model.to(device).eval() if half: model.half() # to FP16 # Second-stage classifier classify = False if classify: modelc = torch_utils.load_classifier(name='resnet101', n=2) # initialize modelc.load_state_dict(torch.load('weights/resnet101.pt', map_location=device)['model']) # load weights modelc.to(device).eval() # Set Dataloader vid_path, vid_writer = None, None if webcam: view_img = True torch.backends.cudnn.benchmark = True # set True to speed up constant image size inference dataset = LoadStreams(source, img_size=imgsz) else: save_img = True dataset = LoadImages(source, img_size=imgsz) # Get names and colors names = ['child', 'adult'] # names = model.names if hasattr(model, 'names') else model.modules.names colors = [[random.randint(0, 255) for _ in range(3)] for _ in range(len(names))] # Run inference t0 = time.time() # img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img # _ = model(img.half() if half else img) if device.type != 'cpu' else None # run once for path, img, im0s, vid_cap in dataset: 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 = torch_utils.time_synchronized() # img_input = torch.cat([img[..., ::2, ::2], img[..., 1::2, ::2], img[..., ::2, 1::2], img[..., 1::2, 1::2]], 1) pred = model(img, augment=opt.augment)[0] # Apply NMS # pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, # fast=True, classes=opt.classes, agnostic=opt.agnostic_nms) pred = non_max_suppression(pred, opt.conf_thres, opt.iou_thres, fast=False, classes=opt.classes, agnostic=True) t2 = torch_utils.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 = path[i], '%g: ' % i, im0s[i].copy() else: p, s, im0 = path, '', im0s save_path = str(Path(out) / Path(p).name) s += '%gx%g ' % img.shape[2:] # print string gn = torch.tensor(im0.shape)[[1, 0, 1, 0]].float() # normalization gain whwh 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 += '%g %ss, ' % (n, names[int(c)]) # add to string # Write results for *xyxy, conf, cls in det: if save_txt: # Write to file xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh with open(save_path[:save_path.rfind('.')] + '.txt', 'a') as file: file.write(('%g %g %g %g %g' + '\n') % (cls, *xywh)) # label format if save_img or view_img: # Add bbox to image label = '%s %.2f' % (names[int(cls)], conf) plot_one_box(xyxy, im0, label=label, color=colors[int(cls)], line_thickness=3) # Print time (inference + NMS) print('%sDone. (%.3fs)' % (s, t2 - t1)) # Stream results if view_img: cv2.imshow(p, im0) if cv2.waitKey(1) == ord('q'): # q to quit raise StopIteration # Save results (image with detections) if save_img: if dataset.mode == 'images': cv2.imwrite(save_path, im0) else: 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))