diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 48e752f448f1..f3f687c2fc94 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -36,12 +36,11 @@ repos: - id: isort name: Sort imports - # TODO - #- repo: https://github.com/pre-commit/mirrors-yapf - # rev: v0.31.0 - # hooks: - # - id: yapf - # name: formatting + - repo: https://github.com/psf/black + rev: 21.7b0 + hooks: + - id: black + name: Black code # TODO #- repo: https://github.com/executablebooks/mdformat diff --git a/detect.py b/detect.py index 661a0b86bc99..fde14c1695a7 100644 --- a/detect.py +++ b/detect.py @@ -31,105 +31,127 @@ from models.experimental import attempt_load from utils.datasets import IMG_FORMATS, VID_FORMATS, LoadImages, LoadStreams -from utils.general import (LOGGER, apply_classifier, check_file, check_img_size, check_imshow, check_requirements, - check_suffix, colorstr, increment_path, non_max_suppression, print_args, scale_coords, - strip_optimizer, xyxy2xywh) +from utils.general import ( + LOGGER, + apply_classifier, + check_file, + check_img_size, + check_imshow, + check_requirements, + check_suffix, + colorstr, + increment_path, + non_max_suppression, + print_args, + scale_coords, + strip_optimizer, + xyxy2xywh, +) from utils.plots import Annotator, colors, save_one_box from utils.torch_utils import load_classifier, select_device, time_sync @torch.no_grad() -def run(weights=ROOT / 'yolov5s.pt', # model.pt path(s) - source=ROOT / 'data/images', # file/dir/URL/glob, 0 for webcam - imgsz=640, # inference size (pixels) - conf_thres=0.25, # confidence threshold - iou_thres=0.45, # NMS IOU threshold - max_det=1000, # maximum detections per image - device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu - view_img=False, # show results - save_txt=False, # save results to *.txt - save_conf=False, # save confidences in --save-txt labels - save_crop=False, # save cropped prediction boxes - nosave=False, # do not save images/videos - classes=None, # filter by class: --class 0, or --class 0 2 3 - agnostic_nms=False, # class-agnostic NMS - augment=False, # augmented inference - visualize=False, # visualize features - update=False, # update all models - project=ROOT / 'runs/detect', # save results to project/name - name='exp', # save results to project/name - exist_ok=False, # existing project/name ok, do not increment - line_thickness=3, # bounding box thickness (pixels) - hide_labels=False, # hide labels - hide_conf=False, # hide confidences - half=False, # use FP16 half-precision inference - dnn=False, # use OpenCV DNN for ONNX inference - ): +def run( + weights=ROOT / "yolov5s.pt", # model.pt path(s) + source=ROOT / "data/images", # file/dir/URL/glob, 0 for webcam + imgsz=640, # inference size (pixels) + conf_thres=0.25, # confidence threshold + iou_thres=0.45, # NMS IOU threshold + max_det=1000, # maximum detections per image + device="", # cuda device, i.e. 0 or 0,1,2,3 or cpu + view_img=False, # show results + save_txt=False, # save results to *.txt + save_conf=False, # save confidences in --save-txt labels + save_crop=False, # save cropped prediction boxes + nosave=False, # do not save images/videos + classes=None, # filter by class: --class 0, or --class 0 2 3 + agnostic_nms=False, # class-agnostic NMS + augment=False, # augmented inference + visualize=False, # visualize features + update=False, # update all models + project=ROOT / "runs/detect", # save results to project/name + name="exp", # save results to project/name + exist_ok=False, # existing project/name ok, do not increment + line_thickness=3, # bounding box thickness (pixels) + hide_labels=False, # hide labels + hide_conf=False, # hide confidences + half=False, # use FP16 half-precision inference + dnn=False, # use OpenCV DNN for ONNX inference +): source = str(source) - save_img = not nosave and not source.endswith('.txt') # save inference images + save_img = not nosave and not source.endswith(".txt") # save inference images is_file = Path(source).suffix[1:] in (IMG_FORMATS + VID_FORMATS) - is_url = source.lower().startswith(('rtsp://', 'rtmp://', 'http://', 'https://')) - webcam = source.isnumeric() or source.endswith('.txt') or (is_url and not is_file) + is_url = source.lower().startswith(("rtsp://", "rtmp://", "http://", "https://")) + webcam = source.isnumeric() or source.endswith(".txt") or (is_url and not is_file) if is_url and is_file: source = check_file(source) # download # Directories save_dir = increment_path(Path(project) / name, exist_ok=exist_ok) # increment run - (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir + (save_dir / "labels" if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir # Initialize device = select_device(device) - half &= device.type != 'cpu' # half precision only supported on CUDA + half &= device.type != "cpu" # half precision only supported on CUDA # Load model w = str(weights[0] if isinstance(weights, list) else weights) - classify, suffix, suffixes = False, Path(w).suffix.lower(), ['.pt', '.onnx', '.tflite', '.pb', ''] + classify, suffix, suffixes = False, Path(w).suffix.lower(), [".pt", ".onnx", ".tflite", ".pb", ""] check_suffix(w, suffixes) # check weights have acceptable suffix pt, onnx, tflite, pb, saved_model = (suffix == x for x in suffixes) # backend booleans - stride, names = 64, [f'class{i}' for i in range(1000)] # assign defaults + stride, names = 64, [f"class{i}" for i in range(1000)] # assign defaults if pt: - model = torch.jit.load(w) if 'torchscript' in w else attempt_load(weights, map_location=device) + model = torch.jit.load(w) if "torchscript" in w else attempt_load(weights, map_location=device) stride = int(model.stride.max()) # model stride - names = model.module.names if hasattr(model, 'module') else model.names # get class names + names = model.module.names if hasattr(model, "module") else model.names # get class names if half: model.half() # to FP16 if classify: # second-stage classifier - modelc = load_classifier(name='resnet50', n=2) # initialize - modelc.load_state_dict(torch.load('resnet50.pt', map_location=device)['model']).to(device).eval() + modelc = load_classifier(name="resnet50", n=2) # initialize + modelc.load_state_dict(torch.load("resnet50.pt", map_location=device)["model"]).to(device).eval() elif onnx: if dnn: - check_requirements(('opencv-python>=4.5.4',)) + check_requirements(("opencv-python>=4.5.4",)) net = cv2.dnn.readNetFromONNX(w) else: - check_requirements(('onnx', 'onnxruntime-gpu' if torch.has_cuda else 'onnxruntime')) + check_requirements(("onnx", "onnxruntime-gpu" if torch.has_cuda else "onnxruntime")) import onnxruntime + session = onnxruntime.InferenceSession(w, None) else: # TensorFlow models import tensorflow as tf + if pb: # https://www.tensorflow.org/guide/migrate#a_graphpb_or_graphpbtxt + def wrap_frozen_graph(gd, inputs, outputs): x = tf.compat.v1.wrap_function(lambda: tf.compat.v1.import_graph_def(gd, name=""), []) # wrapped import - return x.prune(tf.nest.map_structure(x.graph.as_graph_element, inputs), - tf.nest.map_structure(x.graph.as_graph_element, outputs)) + return x.prune( + tf.nest.map_structure(x.graph.as_graph_element, inputs), + tf.nest.map_structure(x.graph.as_graph_element, outputs), + ) graph_def = tf.Graph().as_graph_def() - graph_def.ParseFromString(open(w, 'rb').read()) + graph_def.ParseFromString(open(w, "rb").read()) frozen_func = wrap_frozen_graph(gd=graph_def, inputs="x:0", outputs="Identity:0") elif saved_model: model = tf.keras.models.load_model(w) elif tflite: if "edgetpu" in w: # https://www.tensorflow.org/lite/guide/python#install_tensorflow_lite_for_python import tflite_runtime.interpreter as tflri - delegate = {'Linux': 'libedgetpu.so.1', # install libedgetpu https://coral.ai/software/#edgetpu-runtime - 'Darwin': 'libedgetpu.1.dylib', - 'Windows': 'edgetpu.dll'}[platform.system()] + + delegate = { + "Linux": "libedgetpu.so.1", # install libedgetpu https://coral.ai/software/#edgetpu-runtime + "Darwin": "libedgetpu.1.dylib", + "Windows": "edgetpu.dll", + }[platform.system()] interpreter = tflri.Interpreter(model_path=w, experimental_delegates=[tflri.load_delegate(delegate)]) else: interpreter = tf.lite.Interpreter(model_path=w) # load TFLite model interpreter.allocate_tensors() # allocate input_details = interpreter.get_input_details() # inputs output_details = interpreter.get_output_details() # outputs - int8 = input_details[0]['dtype'] == np.uint8 # is TFLite quantized uint8 model + int8 = input_details[0]["dtype"] == np.uint8 # is TFLite quantized uint8 model imgsz = check_img_size(imgsz, s=stride) # check image size # Dataloader @@ -144,13 +166,13 @@ def wrap_frozen_graph(gd, inputs, outputs): vid_path, vid_writer = [None] * bs, [None] * bs # Run inference - if pt and device.type != 'cpu': + if pt and device.type != "cpu": model(torch.zeros(1, 3, *imgsz).to(device).type_as(next(model.parameters()))) # run once dt, seen = [0.0, 0.0, 0.0], 0 for path, img, im0s, vid_cap, s in dataset: t1 = time_sync() if onnx: - img = img.astype('float32') + img = img.astype("float32") else: img = torch.from_numpy(img).to(device) img = img.half() if half else img.float() # uint8 to fp16/32 @@ -178,13 +200,13 @@ def wrap_frozen_graph(gd, inputs, outputs): pred = model(imn, training=False).numpy() elif tflite: if int8: - scale, zero_point = input_details[0]['quantization'] + scale, zero_point = input_details[0]["quantization"] imn = (imn / scale + zero_point).astype(np.uint8) # de-scale - interpreter.set_tensor(input_details[0]['index'], imn) + interpreter.set_tensor(input_details[0]["index"], imn) interpreter.invoke() - pred = interpreter.get_tensor(output_details[0]['index']) + pred = interpreter.get_tensor(output_details[0]["index"]) if int8: - scale, zero_point = output_details[0]['quantization'] + scale, zero_point = output_details[0]["quantization"] pred = (pred.astype(np.float32) - zero_point) * scale # re-scale pred[..., 0] *= imgsz[1] # x pred[..., 1] *= imgsz[0] # y @@ -207,14 +229,14 @@ def wrap_frozen_graph(gd, inputs, outputs): seen += 1 if webcam: # batch_size >= 1 p, im0, frame = path[i], im0s[i].copy(), dataset.count - s += f'{i}: ' + s += f"{i}: " else: - p, im0, frame = path, im0s.copy(), getattr(dataset, 'frame', 0) + p, im0, frame = path, im0s.copy(), getattr(dataset, "frame", 0) p = Path(p) # to Path save_path = str(save_dir / p.name) # img.jpg - txt_path = str(save_dir / 'labels' / p.stem) + ('' if dataset.mode == 'image' else f'_{frame}') # img.txt - s += '%gx%g ' % img.shape[2:] # print string + txt_path = str(save_dir / "labels" / p.stem) + ("" if dataset.mode == "image" else f"_{frame}") # img.txt + s += "%gx%g " % img.shape[2:] # print string gn = torch.tensor(im0.shape)[[1, 0, 1, 0]] # normalization gain whwh imc = im0.copy() if save_crop else im0 # for save_crop annotator = Annotator(im0, line_width=line_thickness, example=str(names)) @@ -232,18 +254,18 @@ def wrap_frozen_graph(gd, inputs, outputs): if save_txt: # Write to file xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh line = (cls, *xywh, conf) if save_conf else (cls, *xywh) # label format - with open(txt_path + '.txt', 'a') as f: - f.write(('%g ' * len(line)).rstrip() % line + '\n') + with open(txt_path + ".txt", "a") as f: + f.write(("%g " * len(line)).rstrip() % line + "\n") if save_img or save_crop or view_img: # Add bbox to image c = int(cls) # integer class - label = None if hide_labels else (names[c] if hide_conf else f'{names[c]} {conf:.2f}') + label = None if hide_labels else (names[c] if hide_conf else f"{names[c]} {conf:.2f}") annotator.box_label(xyxy, label, color=colors(c, True)) if save_crop: - save_one_box(xyxy, imc, file=save_dir / 'crops' / names[c] / f'{p.stem}.jpg', BGR=True) + save_one_box(xyxy, imc, file=save_dir / "crops" / names[c] / f"{p.stem}.jpg", BGR=True) # Print time (inference-only) - LOGGER.info(f'{s}Done. ({t3 - t2:.3f}s)') + LOGGER.info(f"{s}Done. ({t3 - t2:.3f}s)") # Stream results im0 = annotator.result() @@ -253,7 +275,7 @@ def wrap_frozen_graph(gd, inputs, outputs): # Save results (image with detections) if save_img: - if dataset.mode == 'image': + if dataset.mode == "image": cv2.imwrite(save_path, im0) else: # 'video' or 'stream' if vid_path[i] != save_path: # new video @@ -266,15 +288,15 @@ def wrap_frozen_graph(gd, inputs, outputs): 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[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h)) + save_path += ".mp4" + vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h)) vid_writer[i].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 per image at shape {(1, 3, *imgsz)}' % t) + t = tuple(x / seen * 1e3 for x in dt) # speeds per image + LOGGER.info(f"Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {(1, 3, *imgsz)}" % t) if save_txt or save_img: - s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else '' + s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else "" LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}") if update: strip_optimizer(weights) # update model (to fix SourceChangeWarning) @@ -282,31 +304,31 @@ def wrap_frozen_graph(gd, inputs, outputs): def parse_opt(): parser = argparse.ArgumentParser() - parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model path(s)') - parser.add_argument('--source', type=str, default=ROOT / 'data/images', help='file/dir/URL/glob, 0 for webcam') - 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.25, help='confidence threshold') - parser.add_argument('--iou-thres', type=float, default=0.45, help='NMS IoU threshold') - parser.add_argument('--max-det', type=int, default=1000, help='maximum detections per image') - 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='show results') - parser.add_argument('--save-txt', action='store_true', help='save results to *.txt') - parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels') - parser.add_argument('--save-crop', action='store_true', help='save cropped prediction boxes') - parser.add_argument('--nosave', action='store_true', help='do not save images/videos') - parser.add_argument('--classes', nargs='+', type=int, help='filter by class: --classes 0, or --classes 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('--visualize', action='store_true', help='visualize features') - parser.add_argument('--update', action='store_true', help='update all models') - parser.add_argument('--project', default=ROOT / 'runs/detect', 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') - parser.add_argument('--line-thickness', default=3, type=int, help='bounding box thickness (pixels)') - parser.add_argument('--hide-labels', default=False, action='store_true', help='hide labels') - parser.add_argument('--hide-conf', default=False, action='store_true', help='hide confidences') - parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference') - parser.add_argument('--dnn', action='store_true', help='use OpenCV DNN for ONNX inference') + parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "yolov5s.pt", help="model path(s)") + parser.add_argument("--source", type=str, default=ROOT / "data/images", help="file/dir/URL/glob, 0 for webcam") + 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.25, help="confidence threshold") + parser.add_argument("--iou-thres", type=float, default=0.45, help="NMS IoU threshold") + parser.add_argument("--max-det", type=int, default=1000, help="maximum detections per image") + 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="show results") + parser.add_argument("--save-txt", action="store_true", help="save results to *.txt") + parser.add_argument("--save-conf", action="store_true", help="save confidences in --save-txt labels") + parser.add_argument("--save-crop", action="store_true", help="save cropped prediction boxes") + parser.add_argument("--nosave", action="store_true", help="do not save images/videos") + parser.add_argument("--classes", nargs="+", type=int, help="filter by class: --classes 0, or --classes 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("--visualize", action="store_true", help="visualize features") + parser.add_argument("--update", action="store_true", help="update all models") + parser.add_argument("--project", default=ROOT / "runs/detect", 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") + parser.add_argument("--line-thickness", default=3, type=int, help="bounding box thickness (pixels)") + parser.add_argument("--hide-labels", default=False, action="store_true", help="hide labels") + parser.add_argument("--hide-conf", default=False, action="store_true", help="hide confidences") + parser.add_argument("--half", action="store_true", help="use FP16 half-precision inference") + parser.add_argument("--dnn", action="store_true", help="use OpenCV DNN for ONNX inference") opt = parser.parse_args() opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1 # expand print_args(FILE.stem, opt) @@ -314,7 +336,7 @@ def parse_opt(): def main(opt): - check_requirements(exclude=('tensorboard', 'thop')) + check_requirements(exclude=("tensorboard", "thop")) run(**vars(opt)) diff --git a/export.py b/export.py index f5eb487045b0..80d187fd5425 100644 --- a/export.py +++ b/export.py @@ -42,42 +42,59 @@ from models.yolo import Detect from utils.activations import SiLU from utils.datasets import LoadImages -from utils.general import (LOGGER, check_dataset, check_img_size, check_requirements, colorstr, file_size, print_args, - url2file) +from utils.general import ( + LOGGER, + check_dataset, + check_img_size, + check_requirements, + colorstr, + file_size, + print_args, + url2file, +) from utils.torch_utils import select_device -def export_torchscript(model, im, file, optimize, prefix=colorstr('TorchScript:')): +def export_torchscript(model, im, file, optimize, prefix=colorstr("TorchScript:")): # YOLOv5 TorchScript model export try: - LOGGER.info(f'\n{prefix} starting export with torch {torch.__version__}...') - f = file.with_suffix('.torchscript.pt') + LOGGER.info(f"\n{prefix} starting export with torch {torch.__version__}...") + f = file.with_suffix(".torchscript.pt") ts = torch.jit.trace(model, im, strict=False) (optimize_for_mobile(ts) if optimize else ts).save(f) - LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') + LOGGER.info(f"{prefix} export success, saved as {f} ({file_size(f):.1f} MB)") except Exception as e: - LOGGER.info(f'{prefix} export failure: {e}') + LOGGER.info(f"{prefix} export failure: {e}") -def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorstr('ONNX:')): +def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorstr("ONNX:")): # YOLOv5 ONNX export try: - check_requirements(('onnx',)) + check_requirements(("onnx",)) import onnx - LOGGER.info(f'\n{prefix} starting export with onnx {onnx.__version__}...') - f = file.with_suffix('.onnx') - - torch.onnx.export(model, im, f, verbose=False, opset_version=opset, - training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL, - do_constant_folding=not train, - input_names=['images'], - output_names=['output'], - dynamic_axes={'images': {0: 'batch', 2: 'height', 3: 'width'}, # shape(1,3,640,640) - 'output': {0: 'batch', 1: 'anchors'} # shape(1,25200,85) - } if dynamic else None) + LOGGER.info(f"\n{prefix} starting export with onnx {onnx.__version__}...") + f = file.with_suffix(".onnx") + + torch.onnx.export( + model, + im, + f, + verbose=False, + opset_version=opset, + training=torch.onnx.TrainingMode.TRAINING if train else torch.onnx.TrainingMode.EVAL, + do_constant_folding=not train, + input_names=["images"], + output_names=["output"], + dynamic_axes={ + "images": {0: "batch", 2: "height", 3: "width"}, # shape(1,3,640,640) + "output": {0: "batch", 1: "anchors"}, # shape(1,25200,85) + } + if dynamic + else None, + ) # Checks model_onnx = onnx.load(f) # load onnx model @@ -87,49 +104,60 @@ def export_onnx(model, im, file, opset, train, dynamic, simplify, prefix=colorst # Simplify if simplify: try: - check_requirements(('onnx-simplifier',)) + check_requirements(("onnx-simplifier",)) import onnxsim - LOGGER.info(f'{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...') + LOGGER.info(f"{prefix} simplifying with onnx-simplifier {onnxsim.__version__}...") model_onnx, check = onnxsim.simplify( model_onnx, dynamic_input_shape=dynamic, - input_shapes={'images': list(im.shape)} if dynamic else None) - assert check, 'assert check failed' + input_shapes={"images": list(im.shape)} if dynamic else None, + ) + assert check, "assert check failed" onnx.save(model_onnx, f) except Exception as e: - LOGGER.info(f'{prefix} simplifier failure: {e}') - LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') + LOGGER.info(f"{prefix} simplifier failure: {e}") + LOGGER.info(f"{prefix} export success, saved as {f} ({file_size(f):.1f} MB)") LOGGER.info(f"{prefix} run --dynamic ONNX model inference with: 'python detect.py --weights {f}'") except Exception as e: - LOGGER.info(f'{prefix} export failure: {e}') + LOGGER.info(f"{prefix} export failure: {e}") -def export_coreml(model, im, file, prefix=colorstr('CoreML:')): +def export_coreml(model, im, file, prefix=colorstr("CoreML:")): # YOLOv5 CoreML export ct_model = None try: - check_requirements(('coremltools',)) + check_requirements(("coremltools",)) import coremltools as ct - LOGGER.info(f'\n{prefix} starting export with coremltools {ct.__version__}...') - f = file.with_suffix('.mlmodel') + LOGGER.info(f"\n{prefix} starting export with coremltools {ct.__version__}...") + f = file.with_suffix(".mlmodel") model.train() # CoreML exports should be placed in model.train() mode ts = torch.jit.trace(model, im, strict=False) # TorchScript model - ct_model = ct.convert(ts, inputs=[ct.ImageType('image', shape=im.shape, scale=1 / 255, bias=[0, 0, 0])]) + ct_model = ct.convert(ts, inputs=[ct.ImageType("image", shape=im.shape, scale=1 / 255, bias=[0, 0, 0])]) ct_model.save(f) - LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') + LOGGER.info(f"{prefix} export success, saved as {f} ({file_size(f):.1f} MB)") except Exception as e: - LOGGER.info(f'\n{prefix} export failure: {e}') + LOGGER.info(f"\n{prefix} export failure: {e}") return ct_model -def export_saved_model(model, im, file, dynamic, - tf_nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100, iou_thres=0.45, - conf_thres=0.25, prefix=colorstr('TensorFlow saved_model:')): +def export_saved_model( + model, + im, + file, + dynamic, + tf_nms=False, + agnostic_nms=False, + topk_per_class=100, + topk_all=100, + iou_thres=0.45, + conf_thres=0.25, + prefix=colorstr("TensorFlow saved_model:"), +): # YOLOv5 TensorFlow saved_model export keras_model = None try: @@ -138,8 +166,8 @@ def export_saved_model(model, im, file, dynamic, from models.tf import TFDetect, TFModel - LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...') - f = str(file).replace('.pt', '_saved_model') + LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...") + f = str(file).replace(".pt", "_saved_model") batch_size, ch, *imgsz = list(im.shape) # BCHW tf_model = TFModel(cfg=model.yaml, model=model, nc=model.nc, imgsz=imgsz) @@ -150,23 +178,23 @@ def export_saved_model(model, im, file, dynamic, keras_model = keras.Model(inputs=inputs, outputs=outputs) keras_model.trainable = False keras_model.summary() - keras_model.save(f, save_format='tf') + keras_model.save(f, save_format="tf") - LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') + LOGGER.info(f"{prefix} export success, saved as {f} ({file_size(f):.1f} MB)") except Exception as e: - LOGGER.info(f'\n{prefix} export failure: {e}') + LOGGER.info(f"\n{prefix} export failure: {e}") return keras_model -def export_pb(keras_model, im, file, prefix=colorstr('TensorFlow GraphDef:')): +def export_pb(keras_model, im, file, prefix=colorstr("TensorFlow GraphDef:")): # YOLOv5 TensorFlow GraphDef *.pb export https://github.com/leimao/Frozen_Graph_TensorFlow try: import tensorflow as tf from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2 - LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...') - f = file.with_suffix('.pb') + LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...") + f = file.with_suffix(".pb") m = tf.function(lambda x: keras_model(x)) # full model m = m.get_concrete_function(tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype)) @@ -174,63 +202,65 @@ def export_pb(keras_model, im, file, prefix=colorstr('TensorFlow GraphDef:')): frozen_func.graph.as_graph_def() tf.io.write_graph(graph_or_graph_def=frozen_func.graph, logdir=str(f.parent), name=f.name, as_text=False) - LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') + LOGGER.info(f"{prefix} export success, saved as {f} ({file_size(f):.1f} MB)") except Exception as e: - LOGGER.info(f'\n{prefix} export failure: {e}') + LOGGER.info(f"\n{prefix} export failure: {e}") -def export_tflite(keras_model, im, file, int8, data, ncalib, prefix=colorstr('TensorFlow Lite:')): +def export_tflite(keras_model, im, file, int8, data, ncalib, prefix=colorstr("TensorFlow Lite:")): # YOLOv5 TensorFlow Lite export try: import tensorflow as tf from models.tf import representative_dataset_gen - LOGGER.info(f'\n{prefix} starting export with tensorflow {tf.__version__}...') + LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...") batch_size, ch, *imgsz = list(im.shape) # BCHW - f = str(file).replace('.pt', '-fp16.tflite') + f = str(file).replace(".pt", "-fp16.tflite") converter = tf.lite.TFLiteConverter.from_keras_model(keras_model) converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS] converter.target_spec.supported_types = [tf.float16] converter.optimizations = [tf.lite.Optimize.DEFAULT] if int8: - dataset = LoadImages(check_dataset(data)['train'], img_size=imgsz, auto=False) # representative data + dataset = LoadImages(check_dataset(data)["train"], img_size=imgsz, auto=False) # representative data converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib) converter.target_spec.supported_ops = [tf.lite.OpsSet.TFLITE_BUILTINS_INT8] converter.target_spec.supported_types = [] converter.inference_input_type = tf.uint8 # or tf.int8 converter.inference_output_type = tf.uint8 # or tf.int8 converter.experimental_new_quantizer = False - f = str(file).replace('.pt', '-int8.tflite') + f = str(file).replace(".pt", "-int8.tflite") tflite_model = converter.convert() open(f, "wb").write(tflite_model) - LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') + LOGGER.info(f"{prefix} export success, saved as {f} ({file_size(f):.1f} MB)") except Exception as e: - LOGGER.info(f'\n{prefix} export failure: {e}') + LOGGER.info(f"\n{prefix} export failure: {e}") -def export_tfjs(keras_model, im, file, prefix=colorstr('TensorFlow.js:')): +def export_tfjs(keras_model, im, file, prefix=colorstr("TensorFlow.js:")): # YOLOv5 TensorFlow.js export try: - check_requirements(('tensorflowjs',)) + check_requirements(("tensorflowjs",)) import re import tensorflowjs as tfjs - LOGGER.info(f'\n{prefix} starting export with tensorflowjs {tfjs.__version__}...') - f = str(file).replace('.pt', '_web_model') # js dir - f_pb = file.with_suffix('.pb') # *.pb path - f_json = f + '/model.json' # *.json path + LOGGER.info(f"\n{prefix} starting export with tensorflowjs {tfjs.__version__}...") + f = str(file).replace(".pt", "_web_model") # js dir + f_pb = file.with_suffix(".pb") # *.pb path + f_json = f + "/model.json" # *.json path - cmd = f"tensorflowjs_converter --input_format=tf_frozen_model " \ - f"--output_node_names='Identity,Identity_1,Identity_2,Identity_3' {f_pb} {f}" + cmd = ( + f"tensorflowjs_converter --input_format=tf_frozen_model " + f"--output_node_names='Identity,Identity_1,Identity_2,Identity_3' {f_pb} {f}" + ) subprocess.run(cmd, shell=True) json = open(f_json).read() - with open(f_json, 'w') as j: # sort JSON Identity_* in ascending order + with open(f_json, "w") as j: # sort JSON Identity_* in ascending order subst = re.sub( r'{"outputs": {"Identity.?.?": {"name": "Identity.?.?"}, ' r'"Identity.?.?": {"name": "Identity.?.?"}, ' @@ -240,43 +270,45 @@ def export_tfjs(keras_model, im, file, prefix=colorstr('TensorFlow.js:')): r'"Identity_1": {"name": "Identity_1"}, ' r'"Identity_2": {"name": "Identity_2"}, ' r'"Identity_3": {"name": "Identity_3"}}}', - json) + json, + ) j.write(subst) - LOGGER.info(f'{prefix} export success, saved as {f} ({file_size(f):.1f} MB)') + LOGGER.info(f"{prefix} export success, saved as {f} ({file_size(f):.1f} MB)") except Exception as e: - LOGGER.info(f'\n{prefix} export failure: {e}') + LOGGER.info(f"\n{prefix} export failure: {e}") @torch.no_grad() -def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path' - weights=ROOT / 'yolov5s.pt', # weights path - imgsz=(640, 640), # image (height, width) - batch_size=1, # batch size - device='cpu', # cuda device, i.e. 0 or 0,1,2,3 or cpu - include=('torchscript', 'onnx', 'coreml'), # include formats - half=False, # FP16 half-precision export - inplace=False, # set YOLOv5 Detect() inplace=True - train=False, # model.train() mode - optimize=False, # TorchScript: optimize for mobile - int8=False, # CoreML/TF INT8 quantization - dynamic=False, # ONNX/TF: dynamic axes - simplify=False, # ONNX: simplify model - opset=12, # ONNX: opset version - topk_per_class=100, # TF.js NMS: topk per class to keep - topk_all=100, # TF.js NMS: topk for all classes to keep - iou_thres=0.45, # TF.js NMS: IoU threshold - conf_thres=0.25 # TF.js NMS: confidence threshold - ): +def run( + data=ROOT / "data/coco128.yaml", # 'dataset.yaml path' + weights=ROOT / "yolov5s.pt", # weights path + imgsz=(640, 640), # image (height, width) + batch_size=1, # batch size + device="cpu", # cuda device, i.e. 0 or 0,1,2,3 or cpu + include=("torchscript", "onnx", "coreml"), # include formats + half=False, # FP16 half-precision export + inplace=False, # set YOLOv5 Detect() inplace=True + train=False, # model.train() mode + optimize=False, # TorchScript: optimize for mobile + int8=False, # CoreML/TF INT8 quantization + dynamic=False, # ONNX/TF: dynamic axes + simplify=False, # ONNX: simplify model + opset=12, # ONNX: opset version + topk_per_class=100, # TF.js NMS: topk per class to keep + topk_all=100, # TF.js NMS: topk for all classes to keep + iou_thres=0.45, # TF.js NMS: IoU threshold + conf_thres=0.25, # TF.js NMS: confidence threshold +): t = time.time() include = [x.lower() for x in include] - tf_exports = list(x in include for x in ('saved_model', 'pb', 'tflite', 'tfjs')) # TensorFlow exports + tf_exports = list(x in include for x in ("saved_model", "pb", "tflite", "tfjs")) # TensorFlow exports imgsz *= 2 if len(imgsz) == 1 else 1 # expand - file = Path(url2file(weights) if str(weights).startswith(('http:/', 'https:/')) else weights) + file = Path(url2file(weights) if str(weights).startswith(("http:/", "https:/")) else weights) # Load PyTorch model device = select_device(device) - assert not (device.type == 'cpu' and half), '--half only compatible with GPU export, i.e. use --device 0' + assert not (device.type == "cpu" and half), "--half only compatible with GPU export, i.e. use --device 0" model = attempt_load(weights, map_location=device, inplace=True, fuse=True) # load FP32 model nc, names = model.nc, model.names # number of classes, class names @@ -303,20 +335,29 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path' LOGGER.info(f"\n{colorstr('PyTorch:')} starting from {file} ({file_size(file):.1f} MB)") # Exports - if 'torchscript' in include: + if "torchscript" in include: export_torchscript(model, im, file, optimize) - if 'onnx' in include: + if "onnx" in include: export_onnx(model, im, file, opset, train, dynamic, simplify) - if 'coreml' in include: + if "coreml" in include: export_coreml(model, im, file) # TensorFlow Exports if any(tf_exports): pb, tflite, tfjs = tf_exports[1:] - assert not (tflite and tfjs), 'TFLite and TF.js models must be exported separately, please pass only one type.' - model = export_saved_model(model, im, file, dynamic, tf_nms=tfjs, agnostic_nms=tfjs, - topk_per_class=topk_per_class, topk_all=topk_all, conf_thres=conf_thres, - iou_thres=iou_thres) # keras model + assert not (tflite and tfjs), "TFLite and TF.js models must be exported separately, please pass only one type." + model = export_saved_model( + model, + im, + file, + dynamic, + tf_nms=tfjs, + agnostic_nms=tfjs, + topk_per_class=topk_per_class, + topk_all=topk_all, + conf_thres=conf_thres, + iou_thres=iou_thres, + ) # keras model if pb or tfjs: # pb prerequisite to tfjs export_pb(model, im, file) if tflite: @@ -325,33 +366,38 @@ def run(data=ROOT / 'data/coco128.yaml', # 'dataset.yaml path' export_tfjs(model, im, file) # Finish - LOGGER.info(f'\nExport complete ({time.time() - t:.2f}s)' - f"\nResults saved to {colorstr('bold', file.parent.resolve())}" - f'\nVisualize with https://netron.app') + LOGGER.info( + f"\nExport complete ({time.time() - t:.2f}s)" + f"\nResults saved to {colorstr('bold', file.parent.resolve())}" + f"\nVisualize with https://netron.app" + ) def parse_opt(): parser = argparse.ArgumentParser() - parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path') - parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='weights path') - parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640, 640], help='image (h, w)') - parser.add_argument('--batch-size', type=int, default=1, help='batch size') - parser.add_argument('--device', default='cpu', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') - parser.add_argument('--half', action='store_true', help='FP16 half-precision export') - parser.add_argument('--inplace', action='store_true', help='set YOLOv5 Detect() inplace=True') - parser.add_argument('--train', action='store_true', help='model.train() mode') - parser.add_argument('--optimize', action='store_true', help='TorchScript: optimize for mobile') - parser.add_argument('--int8', action='store_true', help='CoreML/TF INT8 quantization') - parser.add_argument('--dynamic', action='store_true', help='ONNX/TF: dynamic axes') - parser.add_argument('--simplify', action='store_true', help='ONNX: simplify model') - parser.add_argument('--opset', type=int, default=13, help='ONNX: opset version') - parser.add_argument('--topk-per-class', type=int, default=100, help='TF.js NMS: topk per class to keep') - parser.add_argument('--topk-all', type=int, default=100, help='TF.js NMS: topk for all classes to keep') - parser.add_argument('--iou-thres', type=float, default=0.45, help='TF.js NMS: IoU threshold') - parser.add_argument('--conf-thres', type=float, default=0.25, help='TF.js NMS: confidence threshold') - parser.add_argument('--include', nargs='+', - default=['torchscript', 'onnx'], - help='available formats are (torchscript, onnx, coreml, saved_model, pb, tflite, tfjs)') + parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="dataset.yaml path") + parser.add_argument("--weights", type=str, default=ROOT / "yolov5s.pt", help="weights path") + parser.add_argument("--imgsz", "--img", "--img-size", nargs="+", type=int, default=[640, 640], help="image (h, w)") + parser.add_argument("--batch-size", type=int, default=1, help="batch size") + parser.add_argument("--device", default="cpu", help="cuda device, i.e. 0 or 0,1,2,3 or cpu") + parser.add_argument("--half", action="store_true", help="FP16 half-precision export") + parser.add_argument("--inplace", action="store_true", help="set YOLOv5 Detect() inplace=True") + parser.add_argument("--train", action="store_true", help="model.train() mode") + parser.add_argument("--optimize", action="store_true", help="TorchScript: optimize for mobile") + parser.add_argument("--int8", action="store_true", help="CoreML/TF INT8 quantization") + parser.add_argument("--dynamic", action="store_true", help="ONNX/TF: dynamic axes") + parser.add_argument("--simplify", action="store_true", help="ONNX: simplify model") + parser.add_argument("--opset", type=int, default=13, help="ONNX: opset version") + parser.add_argument("--topk-per-class", type=int, default=100, help="TF.js NMS: topk per class to keep") + parser.add_argument("--topk-all", type=int, default=100, help="TF.js NMS: topk for all classes to keep") + parser.add_argument("--iou-thres", type=float, default=0.45, help="TF.js NMS: IoU threshold") + parser.add_argument("--conf-thres", type=float, default=0.25, help="TF.js NMS: confidence threshold") + parser.add_argument( + "--include", + nargs="+", + default=["torchscript", "onnx"], + help="available formats are (torchscript, onnx, coreml, saved_model, pb, tflite, tfjs)", + ) opt = parser.parse_args() print_args(FILE.stem, opt) return opt diff --git a/hubconf.py b/hubconf.py index 3488fef76ac5..2166da0dafd3 100644 --- a/hubconf.py +++ b/hubconf.py @@ -34,93 +34,93 @@ def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbo from utils.torch_utils import select_device file = Path(__file__).resolve() - check_requirements(exclude=('tensorboard', 'thop', 'opencv-python')) + 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 + 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) + 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 + 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 + 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 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 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 + 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 custom(path='path/to/model.pt', autoshape=True, verbose=True, device=None): +def custom(path="path/to/model.pt", autoshape=True, verbose=True, device=None): # YOLOv5 custom or local model return _create(path, autoshape=autoshape, verbose=verbose, device=device) def yolov5n(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): # YOLOv5-nano model https://github.com/ultralytics/yolov5 - return _create('yolov5n', pretrained, channels, classes, autoshape, verbose, device) + return _create("yolov5n", pretrained, channels, classes, autoshape, verbose, device) def yolov5s(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): # YOLOv5-small model https://github.com/ultralytics/yolov5 - return _create('yolov5s', pretrained, channels, classes, autoshape, verbose, device) + return _create("yolov5s", pretrained, channels, classes, autoshape, verbose, device) def yolov5m(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): # YOLOv5-medium model https://github.com/ultralytics/yolov5 - return _create('yolov5m', pretrained, channels, classes, autoshape, verbose, device) + return _create("yolov5m", pretrained, channels, classes, autoshape, verbose, device) def yolov5l(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): # YOLOv5-large model https://github.com/ultralytics/yolov5 - return _create('yolov5l', pretrained, channels, classes, autoshape, verbose, device) + return _create("yolov5l", pretrained, channels, classes, autoshape, verbose, device) def yolov5x(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): # YOLOv5-xlarge model https://github.com/ultralytics/yolov5 - return _create('yolov5x', pretrained, channels, classes, autoshape, verbose, device) + return _create("yolov5x", pretrained, channels, classes, autoshape, verbose, device) def yolov5n6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): # YOLOv5-nano-P6 model https://github.com/ultralytics/yolov5 - return _create('yolov5n6', pretrained, channels, classes, autoshape, verbose, device) + return _create("yolov5n6", pretrained, channels, classes, autoshape, verbose, device) def yolov5s6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): # YOLOv5-small-P6 model https://github.com/ultralytics/yolov5 - return _create('yolov5s6', pretrained, channels, classes, autoshape, verbose, device) + return _create("yolov5s6", pretrained, channels, classes, autoshape, verbose, device) def yolov5m6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): # YOLOv5-medium-P6 model https://github.com/ultralytics/yolov5 - return _create('yolov5m6', pretrained, channels, classes, autoshape, verbose, device) + return _create("yolov5m6", pretrained, channels, classes, autoshape, verbose, device) def yolov5l6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): # YOLOv5-large-P6 model https://github.com/ultralytics/yolov5 - return _create('yolov5l6', pretrained, channels, classes, autoshape, verbose, device) + return _create("yolov5l6", pretrained, channels, classes, autoshape, verbose, device) def yolov5x6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None): # YOLOv5-xlarge-P6 model https://github.com/ultralytics/yolov5 - return _create('yolov5x6', pretrained, channels, classes, autoshape, verbose, device) + return _create("yolov5x6", pretrained, channels, classes, autoshape, verbose, device) -if __name__ == '__main__': - model = _create(name='yolov5s', pretrained=True, channels=3, classes=80, autoshape=True, verbose=True) # pretrained +if __name__ == "__main__": + model = _create(name="yolov5s", pretrained=True, channels=3, classes=80, autoshape=True, verbose=True) # pretrained # model = custom(path='path/to/model.pt') # custom # Verify inference @@ -130,12 +130,14 @@ def yolov5x6(pretrained=True, channels=3, classes=80, autoshape=True, verbose=Tr import numpy as np from PIL import Image - imgs = ['data/images/zidane.jpg', # filename - Path('data/images/zidane.jpg'), # Path - 'https://ultralytics.com/images/zidane.jpg', # URI - cv2.imread('data/images/bus.jpg')[:, :, ::-1], # OpenCV - Image.open('data/images/bus.jpg'), # PIL - np.zeros((320, 640, 3))] # numpy + imgs = [ + "data/images/zidane.jpg", # filename + Path("data/images/zidane.jpg"), # Path + "https://ultralytics.com/images/zidane.jpg", # URI + cv2.imread("data/images/bus.jpg")[:, :, ::-1], # OpenCV + Image.open("data/images/bus.jpg"), # PIL + np.zeros((320, 640, 3)), + ] # numpy results = model(imgs) # batched inference results.print() diff --git a/models/common.py b/models/common.py index f9e4fc69f006..f5aa3bea6b40 100644 --- a/models/common.py +++ b/models/common.py @@ -169,7 +169,7 @@ def __init__(self, c1, c2, k=(5, 9, 13)): def forward(self, x): x = self.cv1(x) with warnings.catch_warnings(): - warnings.simplefilter('ignore') # suppress torch 1.9.0 max_pool2d() warning + warnings.simplefilter("ignore") # suppress torch 1.9.0 max_pool2d() warning return self.cv2(torch.cat([x] + [m(x) for m in self.m], 1)) @@ -185,7 +185,7 @@ def __init__(self, c1, c2, k=5): # equivalent to SPP(k=(5, 9, 13)) def forward(self, x): x = self.cv1(x) with warnings.catch_warnings(): - warnings.simplefilter('ignore') # suppress torch 1.9.0 max_pool2d() warning + warnings.simplefilter("ignore") # suppress torch 1.9.0 max_pool2d() warning y1 = self.m(x) y2 = self.m(y1) return self.cv2(torch.cat([x, y1, y2, self.m(y2)], 1)) @@ -221,11 +221,14 @@ class GhostBottleneck(nn.Module): def __init__(self, c1, c2, k=3, s=1): # ch_in, ch_out, kernel, stride super().__init__() c_ = c2 // 2 - self.conv = nn.Sequential(GhostConv(c1, c_, 1, 1), # pw - DWConv(c_, c_, k, s, act=False) if s == 2 else nn.Identity(), # dw - GhostConv(c_, c2, 1, 1, act=False)) # pw-linear - self.shortcut = nn.Sequential(DWConv(c1, c1, k, s, act=False), - Conv(c1, c2, 1, 1, act=False)) if s == 2 else nn.Identity() + self.conv = nn.Sequential( + GhostConv(c1, c_, 1, 1), # pw + DWConv(c_, c_, k, s, act=False) if s == 2 else nn.Identity(), # dw + GhostConv(c_, c2, 1, 1, act=False), + ) # pw-linear + self.shortcut = ( + nn.Sequential(DWConv(c1, c1, k, s, act=False), Conv(c1, c2, 1, 1, act=False)) if s == 2 else nn.Identity() + ) def forward(self, x): return self.conv(x) + self.shortcut(x) @@ -282,7 +285,7 @@ def __init__(self, model): self.model = model.eval() def autoshape(self): - LOGGER.info('AutoShape already enabled, skipping... ') # model already converted to model.autoshape() + LOGGER.info("AutoShape already enabled, skipping... ") # model already converted to model.autoshape() return self def _apply(self, fn): @@ -309,26 +312,26 @@ def forward(self, imgs, size=640, augment=False, profile=False): t = [time_sync()] p = next(self.model.parameters()) # for device and type if isinstance(imgs, torch.Tensor): # torch - with amp.autocast(enabled=p.device.type != 'cpu'): + with amp.autocast(enabled=p.device.type != "cpu"): return self.model(imgs.to(p.device).type_as(p), augment, profile) # inference # Pre-process n, imgs = (len(imgs), imgs) if isinstance(imgs, list) else (1, [imgs]) # number of images, list of images shape0, shape1, files = [], [], [] # image and inference shapes, filenames for i, im in enumerate(imgs): - f = f'image{i}' # filename + f = f"image{i}" # filename if isinstance(im, (str, Path)): # filename or uri - im, f = Image.open(requests.get(im, stream=True).raw if str(im).startswith('http') else im), im + im, f = Image.open(requests.get(im, stream=True).raw if str(im).startswith("http") else im), im im = np.asarray(exif_transpose(im)) elif isinstance(im, Image.Image): # PIL Image - im, f = np.asarray(exif_transpose(im)), getattr(im, 'filename', f) or f - files.append(Path(f).with_suffix('.jpg').name) + im, f = np.asarray(exif_transpose(im)), getattr(im, "filename", f) or f + files.append(Path(f).with_suffix(".jpg").name) if im.shape[0] < 5: # image in CHW im = im.transpose((1, 2, 0)) # reverse dataloader .transpose(2, 0, 1) im = im[..., :3] if im.ndim == 3 else np.tile(im[..., None], 3) # enforce 3ch input s = im.shape[:2] # HWC shape0.append(s) # image shape - g = (size / max(s)) # gain + g = size / max(s) # gain shape1.append([y * g for y in s]) imgs[i] = im if im.data.contiguous else np.ascontiguousarray(im) # update shape1 = [make_divisible(x, int(self.stride.max())) for x in np.stack(shape1, 0).max(0)] # inference shape @@ -338,14 +341,20 @@ def forward(self, imgs, size=640, augment=False, profile=False): x = torch.from_numpy(x).to(p.device).type_as(p) / 255 # uint8 to fp16/32 t.append(time_sync()) - with amp.autocast(enabled=p.device.type != 'cpu'): + with amp.autocast(enabled=p.device.type != "cpu"): # Inference y = self.model(x, augment, profile)[0] # forward t.append(time_sync()) # Post-process - y = non_max_suppression(y, self.conf, iou_thres=self.iou, classes=self.classes, - multi_label=self.multi_label, max_det=self.max_det) # NMS + y = non_max_suppression( + y, + self.conf, + iou_thres=self.iou, + classes=self.classes, + multi_label=self.multi_label, + max_det=self.max_det, + ) # NMS for i in range(n): scale_coords(shape1, y[i][:, :4], shape0[i]) @@ -371,10 +380,10 @@ def __init__(self, imgs, pred, files, times=None, names=None, shape=None): self.t = tuple((times[i + 1] - times[i]) * 1000 / self.n for i in range(3)) # timestamps (ms) self.s = shape # inference BCHW shape - def display(self, pprint=False, show=False, save=False, crop=False, render=False, save_dir=Path('')): + def display(self, pprint=False, show=False, save=False, crop=False, render=False, save_dir=Path("")): crops = [] for i, (im, pred) in enumerate(zip(self.imgs, self.pred)): - s = f'image {i + 1}/{len(self.pred)}: {im.shape[0]}x{im.shape[1]} ' # string + s = f"image {i + 1}/{len(self.pred)}: {im.shape[0]}x{im.shape[1]} " # string if pred.shape[0]: for c in pred[:, -1].unique(): n = (pred[:, -1] == c).sum() # detections per class @@ -382,20 +391,27 @@ def display(self, pprint=False, show=False, save=False, crop=False, render=False if show or save or render or crop: annotator = Annotator(im, example=str(self.names)) for *box, conf, cls in reversed(pred): # xyxy, confidence, class - label = f'{self.names[int(cls)]} {conf:.2f}' + label = f"{self.names[int(cls)]} {conf:.2f}" if crop: - file = save_dir / 'crops' / self.names[int(cls)] / self.files[i] if save else None - crops.append({'box': box, 'conf': conf, 'cls': cls, 'label': label, - 'im': save_one_box(box, im, file=file, save=save)}) + file = save_dir / "crops" / self.names[int(cls)] / self.files[i] if save else None + crops.append( + { + "box": box, + "conf": conf, + "cls": cls, + "label": label, + "im": save_one_box(box, im, file=file, save=save), + } + ) else: # all others annotator.box_label(box, label, color=colors(cls)) im = annotator.im else: - s += '(no detections)' + s += "(no detections)" im = Image.fromarray(im.astype(np.uint8)) if isinstance(im, np.ndarray) else im # from np if pprint: - LOGGER.info(s.rstrip(', ')) + LOGGER.info(s.rstrip(", ")) if show: im.show(self.files[i]) # show if save: @@ -407,23 +423,24 @@ def display(self, pprint=False, show=False, save=False, crop=False, render=False self.imgs[i] = np.asarray(im) if crop: if save: - LOGGER.info(f'Saved results to {save_dir}\n') + LOGGER.info(f"Saved results to {save_dir}\n") return crops def print(self): self.display(pprint=True) # print results - LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {tuple(self.s)}' % - self.t) + LOGGER.info( + f"Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {tuple(self.s)}" % self.t + ) def show(self): self.display(show=True) # show results - def save(self, save_dir='runs/detect/exp'): - save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/detect/exp', mkdir=True) # increment save_dir + def save(self, save_dir="runs/detect/exp"): + save_dir = increment_path(save_dir, exist_ok=save_dir != "runs/detect/exp", mkdir=True) # increment save_dir self.display(save=True, save_dir=save_dir) # save results - def crop(self, save=True, save_dir='runs/detect/exp'): - save_dir = increment_path(save_dir, exist_ok=save_dir != 'runs/detect/exp', mkdir=True) if save else None + def crop(self, save=True, save_dir="runs/detect/exp"): + save_dir = increment_path(save_dir, exist_ok=save_dir != "runs/detect/exp", mkdir=True) if save else None return self.display(crop=True, save=save, save_dir=save_dir) # crop results def render(self): @@ -433,9 +450,9 @@ def render(self): def pandas(self): # return detections as pandas DataFrames, i.e. print(results.pandas().xyxy[0]) new = copy(self) # return copy - ca = 'xmin', 'ymin', 'xmax', 'ymax', 'confidence', 'class', 'name' # xyxy columns - cb = 'xcenter', 'ycenter', 'width', 'height', 'confidence', 'class', 'name' # xywh columns - for k, c in zip(['xyxy', 'xyxyn', 'xywh', 'xywhn'], [ca, ca, cb, cb]): + ca = "xmin", "ymin", "xmax", "ymax", "confidence", "class", "name" # xyxy columns + cb = "xcenter", "ycenter", "width", "height", "confidence", "class", "name" # xywh columns + for k, c in zip(["xyxy", "xyxyn", "xywh", "xywhn"], [ca, ca, cb, cb]): a = [[x[:5] + [int(x[5]), self.names[int(x[5])]] for x in x.tolist()] for x in getattr(self, k)] # update setattr(new, k, [pd.DataFrame(x, columns=c) for x in a]) return new @@ -444,7 +461,7 @@ def tolist(self): # return a list of Detections objects, i.e. 'for result in results.tolist():' x = [Detections([self.imgs[i]], [self.pred[i]], self.names, self.s) for i in range(self.n)] for d in x: - for k in ['imgs', 'pred', 'xyxy', 'xyxyn', 'xywh', 'xywhn']: + for k in ["imgs", "pred", "xyxy", "xyxyn", "xywh", "xywhn"]: setattr(d, k, getattr(d, k)[0]) # pop out of list return x diff --git a/models/experimental.py b/models/experimental.py index 463e5514a06e..241ca4878edf 100644 --- a/models/experimental.py +++ b/models/experimental.py @@ -53,7 +53,7 @@ def __init__(self, c1, c2, k=(1, 3), s=1, equal_ch=True): # ch_in, ch_out, kern super().__init__() n = len(k) # number of convolutions if equal_ch: # equal c_ per group - i = torch.linspace(0, n - 1E-6, c2).floor() # c2 indices + i = torch.linspace(0, n - 1e-6, c2).floor() # c2 indices c_ = [(i == g).sum() for g in range(n)] # intermediate channels else: # equal weight.numel() per group b = [c2] + [0] * n @@ -64,7 +64,8 @@ def __init__(self, c1, c2, k=(1, 3), s=1, equal_ch=True): # ch_in, ch_out, kern c_ = np.linalg.lstsq(a, b, rcond=None)[0].round() # solve for equal weight indices, ax = b self.m = nn.ModuleList( - [nn.Conv2d(c1, int(c_), k, s, k // 2, groups=math.gcd(c1, int(c_)), bias=False) for k, c_ in zip(k, c_)]) + [nn.Conv2d(c1, int(c_), k, s, k // 2, groups=math.gcd(c1, int(c_)), bias=False) for k, c_ in zip(k, c_)] + ) self.bn = nn.BatchNorm2d(c2) self.act = nn.SiLU() @@ -95,9 +96,9 @@ def attempt_load(weights, map_location=None, inplace=True, fuse=True): for w in weights if isinstance(weights, list) else [weights]: ckpt = torch.load(attempt_download(w), map_location=map_location) # load if fuse: - model.append(ckpt['ema' if ckpt.get('ema') else 'model'].float().fuse().eval()) # FP32 model + model.append(ckpt["ema" if ckpt.get("ema") else "model"].float().fuse().eval()) # FP32 model else: - model.append(ckpt['ema' if ckpt.get('ema') else 'model'].float().eval()) # without layer fuse + model.append(ckpt["ema" if ckpt.get("ema") else "model"].float().eval()) # without layer fuse # Compatibility updates for m in model.modules(): @@ -105,16 +106,16 @@ def attempt_load(weights, map_location=None, inplace=True, fuse=True): m.inplace = inplace # pytorch 1.7.0 compatibility if type(m) is Detect: if not isinstance(m.anchor_grid, list): # new Detect Layer compatibility - delattr(m, 'anchor_grid') - setattr(m, 'anchor_grid', [torch.zeros(1)] * m.nl) + delattr(m, "anchor_grid") + setattr(m, "anchor_grid", [torch.zeros(1)] * m.nl) elif type(m) is Conv: m._non_persistent_buffers_set = set() # pytorch 1.6.0 compatibility if len(model) == 1: return model[-1] # return model else: - print(f'Ensemble created with {weights}\n') - for k in ['names']: + print(f"Ensemble created with {weights}\n") + for k in ["names"]: setattr(model, k, getattr(model[-1], k)) model.stride = model[torch.argmax(torch.tensor([m.stride.max() for m in model])).int()].stride # max stride return model # return ensemble diff --git a/models/tf.py b/models/tf.py index 6de0245cfe50..a0b230300b7a 100644 --- a/models/tf.py +++ b/models/tf.py @@ -44,7 +44,8 @@ def __init__(self, w=None): gamma_initializer=keras.initializers.Constant(w.weight.numpy()), moving_mean_initializer=keras.initializers.Constant(w.running_mean.numpy()), moving_variance_initializer=keras.initializers.Constant(w.running_var.numpy()), - epsilon=w.eps) + epsilon=w.eps, + ) def call(self, inputs): return self.bn(inputs) @@ -56,7 +57,7 @@ def __init__(self, pad): self.pad = tf.constant([[0, 0], [pad, pad], [pad, pad], [0, 0]]) def call(self, inputs): - return tf.pad(inputs, self.pad, mode='constant', constant_values=0) + return tf.pad(inputs, self.pad, mode="constant", constant_values=0) class TFConv(keras.layers.Layer): @@ -70,11 +71,16 @@ def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True, w=None): # see https://stackoverflow.com/questions/52975843/comparing-conv2d-with-padding-between-tensorflow-and-pytorch conv = keras.layers.Conv2D( - c2, k, s, 'SAME' if s == 1 else 'VALID', use_bias=False if hasattr(w, 'bn') else True, + c2, + k, + s, + "SAME" if s == 1 else "VALID", + use_bias=False if hasattr(w, "bn") else True, kernel_initializer=keras.initializers.Constant(w.conv.weight.permute(2, 3, 1, 0).numpy()), - bias_initializer='zeros' if hasattr(w, 'bn') else keras.initializers.Constant(w.conv.bias.numpy())) + bias_initializer="zeros" if hasattr(w, "bn") else keras.initializers.Constant(w.conv.bias.numpy()), + ) self.conv = conv if s == 1 else keras.Sequential([TFPad(autopad(k, p)), conv]) - self.bn = TFBN(w.bn) if hasattr(w, 'bn') else tf.identity + self.bn = TFBN(w.bn) if hasattr(w, "bn") else tf.identity # YOLOv5 activations if isinstance(w.act, nn.LeakyReLU): @@ -84,7 +90,7 @@ def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True, w=None): elif isinstance(w.act, (nn.SiLU, SiLU)): self.act = (lambda x: keras.activations.swish(x)) if act else tf.identity else: - raise Exception(f'no matching TensorFlow activation found for {w.act}') + raise Exception(f"no matching TensorFlow activation found for {w.act}") def call(self, inputs): return self.act(self.bn(self.conv(inputs))) @@ -99,10 +105,11 @@ def __init__(self, c1, c2, k=1, s=1, p=None, g=1, act=True, w=None): def call(self, inputs): # x(b,w,h,c) -> y(b,w/2,h/2,4c) # inputs = inputs / 255 # normalize 0-255 to 0-1 - return self.conv(tf.concat([inputs[:, ::2, ::2, :], - inputs[:, 1::2, ::2, :], - inputs[:, ::2, 1::2, :], - inputs[:, 1::2, 1::2, :]], 3)) + return self.conv( + tf.concat( + [inputs[:, ::2, ::2, :], inputs[:, 1::2, ::2, :], inputs[:, ::2, 1::2, :], inputs[:, 1::2, 1::2, :]], 3 + ) + ) class TFBottleneck(keras.layers.Layer): @@ -124,9 +131,14 @@ def __init__(self, c1, c2, k, s=1, g=1, bias=True, w=None): super().__init__() assert g == 1, "TF v2.2 Conv2D does not support 'groups' argument" self.conv = keras.layers.Conv2D( - c2, k, s, 'VALID', use_bias=bias, + c2, + k, + s, + "VALID", + use_bias=bias, kernel_initializer=keras.initializers.Constant(w.weight.permute(2, 3, 1, 0).numpy()), - bias_initializer=keras.initializers.Constant(w.bias.numpy()) if bias else None, ) + bias_initializer=keras.initializers.Constant(w.bias.numpy()) if bias else None, + ) def call(self, inputs): return self.conv(inputs) @@ -174,7 +186,7 @@ def __init__(self, c1, c2, k=(5, 9, 13), w=None): c_ = c1 // 2 # hidden channels self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1) self.cv2 = TFConv(c_ * (len(k) + 1), c2, 1, 1, w=w.cv2) - self.m = [keras.layers.MaxPool2D(pool_size=x, strides=1, padding='SAME') for x in k] + self.m = [keras.layers.MaxPool2D(pool_size=x, strides=1, padding="SAME") for x in k] def call(self, inputs): x = self.cv1(inputs) @@ -188,7 +200,7 @@ def __init__(self, c1, c2, k=5, w=None): c_ = c1 // 2 # hidden channels self.cv1 = TFConv(c1, c_, 1, 1, w=w.cv1) self.cv2 = TFConv(c_ * 4, c2, 1, 1, w=w.cv2) - self.m = keras.layers.MaxPool2D(pool_size=k, strides=1, padding='SAME') + self.m = keras.layers.MaxPool2D(pool_size=k, strides=1, padding="SAME") def call(self, inputs): x = self.cv1(inputs) @@ -207,8 +219,7 @@ def __init__(self, nc=80, anchors=(), ch=(), imgsz=(640, 640), w=None): # detec self.na = len(anchors[0]) // 2 # number of anchors self.grid = [tf.zeros(1)] * self.nl # init grid self.anchors = tf.convert_to_tensor(w.anchors.numpy(), dtype=tf.float32) - self.anchor_grid = tf.reshape(self.anchors * tf.reshape(self.stride, [self.nl, 1, 1]), - [self.nl, 1, -1, 1, 2]) + self.anchor_grid = tf.reshape(self.anchors * tf.reshape(self.stride, [self.nl, 1, 1]), [self.nl, 1, -1, 1, 2]) self.m = [TFConv2d(x, self.no * self.na, 1, w=w.m[i]) for i, x in enumerate(ch)] self.training = False # set to False after building model self.imgsz = imgsz @@ -271,12 +282,12 @@ def call(self, inputs): def parse_model(d, ch, model, imgsz): # model_dict, input_channels(3) LOGGER.info(f"\n{'':>3}{'from':>18}{'n':>3}{'params':>10} {'module':<40}{'arguments':<30}") - anchors, nc, gd, gw = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple'] + anchors, nc, gd, gw = d["anchors"], d["nc"], d["depth_multiple"], d["width_multiple"] na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors # number of anchors no = na * (nc + 5) # number of outputs = anchors * (classes + 5) layers, save, c2 = [], [], ch[-1] # layers, savelist, ch out - for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']): # from, number, module, args + for i, (f, n, m, args) in enumerate(d["backbone"] + d["head"]): # from, number, module, args m_str = m m = eval(m) if isinstance(m, str) else m # eval strings for j, a in enumerate(args): @@ -306,15 +317,18 @@ def parse_model(d, ch, model, imgsz): # model_dict, input_channels(3) else: c2 = ch[f] - tf_m = eval('TF' + m_str.replace('nn.', '')) - m_ = keras.Sequential([tf_m(*args, w=model.model[i][j]) for j in range(n)]) if n > 1 \ - else tf_m(*args, w=model.model[i]) # module + tf_m = eval("TF" + m_str.replace("nn.", "")) + m_ = ( + keras.Sequential([tf_m(*args, w=model.model[i][j]) for j in range(n)]) + if n > 1 + else tf_m(*args, w=model.model[i]) + ) # module torch_m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args) # module - t = str(m)[8:-2].replace('__main__.', '') # module type + t = str(m)[8:-2].replace("__main__.", "") # module type np = sum(x.numel() for x in torch_m_.parameters()) # number params m_.i, m_.f, m_.type, m_.np = i, f, t, np # attach index, 'from' index, type, number params - LOGGER.info(f'{i:>3}{str(f):>18}{str(n):>3}{np:>10} {t:<40}{str(args):<30}') # print + LOGGER.info(f"{i:>3}{str(f):>18}{str(n):>3}{np:>10} {t:<40}{str(args):<30}") # print save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist layers.append(m_) ch.append(c2) @@ -322,24 +336,33 @@ def parse_model(d, ch, model, imgsz): # model_dict, input_channels(3) class TFModel: - def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None, model=None, imgsz=(640, 640)): # model, channels, classes + def __init__(self, cfg="yolov5s.yaml", ch=3, nc=None, model=None, imgsz=(640, 640)): # model, channels, classes super().__init__() if isinstance(cfg, dict): self.yaml = cfg # model dict else: # is *.yaml import yaml # for torch hub + self.yaml_file = Path(cfg).name with open(cfg) as f: self.yaml = yaml.load(f, Loader=yaml.FullLoader) # model dict # Define model - if nc and nc != self.yaml['nc']: + if nc and nc != self.yaml["nc"]: LOGGER.info(f"Overriding {cfg} nc={self.yaml['nc']} with nc={nc}") - self.yaml['nc'] = nc # override yaml value + self.yaml["nc"] = nc # override yaml value self.model, self.savelist = parse_model(deepcopy(self.yaml), ch=[ch], model=model, imgsz=imgsz) - def predict(self, inputs, tf_nms=False, agnostic_nms=False, topk_per_class=100, topk_all=100, iou_thres=0.45, - conf_thres=0.25): + def predict( + self, + inputs, + tf_nms=False, + agnostic_nms=False, + topk_per_class=100, + topk_all=100, + iou_thres=0.45, + conf_thres=0.25, + ): y = [] # outputs x = inputs for i, m in enumerate(self.model.layers): @@ -361,7 +384,8 @@ def predict(self, inputs, tf_nms=False, agnostic_nms=False, topk_per_class=100, else: boxes = tf.expand_dims(boxes, 2) nms = tf.image.combined_non_max_suppression( - boxes, scores, topk_per_class, topk_all, iou_thres, conf_thres, clip_boxes=False) + boxes, scores, topk_per_class, topk_all, iou_thres, conf_thres, clip_boxes=False + ) return nms, x[1] return x[0] # output only first tensor [1,6300,85] = [xywh, conf, class0, class1, ...] @@ -382,9 +406,12 @@ class AgnosticNMS(keras.layers.Layer): # TF Agnostic NMS def call(self, input, topk_all, iou_thres, conf_thres): # wrap map_fn to avoid TypeSpec related error https://stackoverflow.com/a/65809989/3036450 - return tf.map_fn(lambda x: self._nms(x, topk_all, iou_thres, conf_thres), input, - fn_output_signature=(tf.float32, tf.float32, tf.float32, tf.int32), - name='agnostic_nms') + return tf.map_fn( + lambda x: self._nms(x, topk_all, iou_thres, conf_thres), + input, + fn_output_signature=(tf.float32, tf.float32, tf.float32, tf.int32), + name="agnostic_nms", + ) @staticmethod def _nms(x, topk_all=100, iou_thres=0.45, conf_thres=0.25): # agnostic NMS @@ -392,19 +419,29 @@ def _nms(x, topk_all=100, iou_thres=0.45, conf_thres=0.25): # agnostic NMS class_inds = tf.cast(tf.argmax(classes, axis=-1), tf.float32) scores_inp = tf.reduce_max(scores, -1) selected_inds = tf.image.non_max_suppression( - boxes, scores_inp, max_output_size=topk_all, iou_threshold=iou_thres, score_threshold=conf_thres) + boxes, scores_inp, max_output_size=topk_all, iou_threshold=iou_thres, score_threshold=conf_thres + ) selected_boxes = tf.gather(boxes, selected_inds) - padded_boxes = tf.pad(selected_boxes, - paddings=[[0, topk_all - tf.shape(selected_boxes)[0]], [0, 0]], - mode="CONSTANT", constant_values=0.0) + padded_boxes = tf.pad( + selected_boxes, + paddings=[[0, topk_all - tf.shape(selected_boxes)[0]], [0, 0]], + mode="CONSTANT", + constant_values=0.0, + ) selected_scores = tf.gather(scores_inp, selected_inds) - padded_scores = tf.pad(selected_scores, - paddings=[[0, topk_all - tf.shape(selected_boxes)[0]]], - mode="CONSTANT", constant_values=-1.0) + padded_scores = tf.pad( + selected_scores, + paddings=[[0, topk_all - tf.shape(selected_boxes)[0]]], + mode="CONSTANT", + constant_values=-1.0, + ) selected_classes = tf.gather(class_inds, selected_inds) - padded_classes = tf.pad(selected_classes, - paddings=[[0, topk_all - tf.shape(selected_boxes)[0]]], - mode="CONSTANT", constant_values=-1.0) + padded_classes = tf.pad( + selected_classes, + paddings=[[0, topk_all - tf.shape(selected_boxes)[0]]], + mode="CONSTANT", + constant_values=-1.0, + ) valid_detections = tf.shape(selected_inds)[0] return padded_boxes, padded_scores, padded_classes, valid_detections @@ -420,14 +457,15 @@ def representative_dataset_gen(dataset, ncalib=100): break -def run(weights=ROOT / 'yolov5s.pt', # weights path - imgsz=(640, 640), # inference size h,w - batch_size=1, # batch size - dynamic=False, # dynamic batch size - ): +def run( + weights=ROOT / "yolov5s.pt", # weights path + imgsz=(640, 640), # inference size h,w + batch_size=1, # batch size + dynamic=False, # dynamic batch size +): # PyTorch model im = torch.zeros((batch_size, 3, *imgsz)) # BCHW image - model = attempt_load(weights, map_location=torch.device('cpu'), inplace=True, fuse=False) + model = attempt_load(weights, map_location=torch.device("cpu"), inplace=True, fuse=False) y = model(im) # inference model.info() @@ -441,15 +479,15 @@ def run(weights=ROOT / 'yolov5s.pt', # weights path keras_model = keras.Model(inputs=im, outputs=tf_model.predict(im)) keras_model.summary() - LOGGER.info('PyTorch, TensorFlow and Keras models successfully verified.\nUse export.py for TF model export.') + LOGGER.info("PyTorch, TensorFlow and Keras models successfully verified.\nUse export.py for TF model export.") def parse_opt(): parser = argparse.ArgumentParser() - parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='weights path') - parser.add_argument('--imgsz', '--img', '--img-size', nargs='+', type=int, default=[640], help='inference size h,w') - parser.add_argument('--batch-size', type=int, default=1, help='batch size') - parser.add_argument('--dynamic', action='store_true', help='dynamic batch size') + parser.add_argument("--weights", type=str, default=ROOT / "yolov5s.pt", help="weights path") + parser.add_argument("--imgsz", "--img", "--img-size", nargs="+", type=int, default=[640], help="inference size h,w") + parser.add_argument("--batch-size", type=int, default=1, help="batch size") + parser.add_argument("--dynamic", action="store_true", help="dynamic batch size") opt = parser.parse_args() opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1 # expand print_args(FILE.stem, opt) diff --git a/models/yolo.py b/models/yolo.py index 305f0ca0cc88..6b610a99f699 100644 --- a/models/yolo.py +++ b/models/yolo.py @@ -22,8 +22,15 @@ from utils.autoanchor import check_anchor_order from utils.general import LOGGER, check_version, check_yaml, make_divisible, print_args from utils.plots import feature_visualization -from utils.torch_utils import (copy_attr, fuse_conv_and_bn, initialize_weights, model_info, scale_img, select_device, - time_sync) +from utils.torch_utils import ( + copy_attr, + fuse_conv_and_bn, + initialize_weights, + model_info, + scale_img, + select_device, + time_sync, +) try: import thop # for FLOPs computation @@ -43,7 +50,7 @@ def __init__(self, nc=80, anchors=(), ch=(), inplace=True): # detection layer self.na = len(anchors[0]) // 2 # number of anchors self.grid = [torch.zeros(1)] * self.nl # init grid self.anchor_grid = [torch.zeros(1)] * self.nl # init anchor grid - self.register_buffer('anchors', torch.tensor(anchors).float().view(self.nl, -1, 2)) # shape(nl,na,2) + self.register_buffer("anchors", torch.tensor(anchors).float().view(self.nl, -1, 2)) # shape(nl,na,2) self.m = nn.ModuleList(nn.Conv2d(x, self.no * self.na, 1) for x in ch) # output conv self.inplace = inplace # use in-place ops (e.g. slice assignment) @@ -72,38 +79,43 @@ def forward(self, x): def _make_grid(self, nx=20, ny=20, i=0): d = self.anchors[i].device - if check_version(torch.__version__, '1.10.0'): # torch>=1.10.0 meshgrid workaround for torch>=0.7 compatibility - yv, xv = torch.meshgrid([torch.arange(ny).to(d), torch.arange(nx).to(d)], indexing='ij') + if check_version(torch.__version__, "1.10.0"): # torch>=1.10.0 meshgrid workaround for torch>=0.7 compatibility + yv, xv = torch.meshgrid([torch.arange(ny).to(d), torch.arange(nx).to(d)], indexing="ij") else: yv, xv = torch.meshgrid([torch.arange(ny).to(d), torch.arange(nx).to(d)]) grid = torch.stack((xv, yv), 2).expand((1, self.na, ny, nx, 2)).float() - anchor_grid = (self.anchors[i].clone() * self.stride[i]) \ - .view((1, self.na, 1, 1, 2)).expand((1, self.na, ny, nx, 2)).float() + anchor_grid = ( + (self.anchors[i].clone() * self.stride[i]) + .view((1, self.na, 1, 1, 2)) + .expand((1, self.na, ny, nx, 2)) + .float() + ) return grid, anchor_grid class Model(nn.Module): - def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None, anchors=None): # model, input channels, number of classes + def __init__(self, cfg="yolov5s.yaml", ch=3, nc=None, anchors=None): # model, input channels, number of classes super().__init__() if isinstance(cfg, dict): self.yaml = cfg # model dict else: # is *.yaml import yaml # for torch hub + self.yaml_file = Path(cfg).name - with open(cfg, encoding='ascii', errors='ignore') as f: + with open(cfg, encoding="ascii", errors="ignore") as f: self.yaml = yaml.safe_load(f) # model dict # Define model - ch = self.yaml['ch'] = self.yaml.get('ch', ch) # input channels - if nc and nc != self.yaml['nc']: + ch = self.yaml["ch"] = self.yaml.get("ch", ch) # input channels + if nc and nc != self.yaml["nc"]: LOGGER.info(f"Overriding model.yaml nc={self.yaml['nc']} with nc={nc}") - self.yaml['nc'] = nc # override yaml value + self.yaml["nc"] = nc # override yaml value if anchors: - LOGGER.info(f'Overriding model.yaml anchors with anchors={anchors}') - self.yaml['anchors'] = round(anchors) # override yaml value + LOGGER.info(f"Overriding model.yaml anchors with anchors={anchors}") + self.yaml["anchors"] = round(anchors) # override yaml value self.model, self.save = parse_model(deepcopy(self.yaml), ch=[ch]) # model, savelist - self.names = [str(i) for i in range(self.yaml['nc'])] # default names - self.inplace = self.yaml.get('inplace', True) + self.names = [str(i) for i in range(self.yaml["nc"])] # default names + self.inplace = self.yaml.get("inplace", True) # Build strides, anchors m = self.model[-1] # Detect() @@ -119,7 +131,7 @@ def __init__(self, cfg='yolov5s.yaml', ch=3, nc=None, anchors=None): # model, i # Init weights, biases initialize_weights(self) self.info() - LOGGER.info('') + LOGGER.info("") def forward(self, x, augment=False, profile=False, visualize=False): if augment: @@ -183,14 +195,14 @@ def _clip_augmented(self, y): def _profile_one_layer(self, m, x, dt): c = isinstance(m, Detect) # is final layer, copy input as inplace fix - o = thop.profile(m, inputs=(x.copy() if c else x,), verbose=False)[0] / 1E9 * 2 if thop else 0 # FLOPs + o = thop.profile(m, inputs=(x.copy() if c else x,), verbose=False)[0] / 1e9 * 2 if thop else 0 # FLOPs t = time_sync() for _ in range(10): m(x.copy() if c else x) dt.append((time_sync() - t) * 100) if m == self.model[0]: LOGGER.info(f"{'time (ms)':>10s} {'GFLOPs':>10s} {'params':>10s} {'module'}") - LOGGER.info(f'{dt[-1]:10.2f} {o:10.2f} {m.np:10.0f} {m.type}') + LOGGER.info(f"{dt[-1]:10.2f} {o:10.2f} {m.np:10.0f} {m.type}") if c: LOGGER.info(f"{sum(dt):10.2f} {'-':>10s} {'-':>10s} Total") @@ -209,7 +221,8 @@ def _print_biases(self): for mi in m.m: # from b = mi.bias.detach().view(m.na, -1).T # conv.bias(255) to (3,85) LOGGER.info( - ('%6g Conv2d.bias:' + '%10.3g' * 6) % (mi.weight.shape[1], *b[:5].mean(1).tolist(), b[5:].mean())) + ("%6g Conv2d.bias:" + "%10.3g" * 6) % (mi.weight.shape[1], *b[:5].mean(1).tolist(), b[5:].mean()) + ) # def _print_weights(self): # for m in self.model.modules(): @@ -217,19 +230,19 @@ def _print_biases(self): # LOGGER.info('%10.3g' % (m.w.detach().sigmoid() * 2)) # shortcut weights def fuse(self): # fuse model Conv2d() + BatchNorm2d() layers - LOGGER.info('Fusing layers... ') + LOGGER.info("Fusing layers... ") for m in self.model.modules(): - if isinstance(m, (Conv, DWConv)) and hasattr(m, 'bn'): + if isinstance(m, (Conv, DWConv)) and hasattr(m, "bn"): m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv - delattr(m, 'bn') # remove batchnorm + delattr(m, "bn") # remove batchnorm m.forward = m.forward_fuse # update forward self.info() return self def autoshape(self): # add AutoShape module - LOGGER.info('Adding AutoShape... ') + LOGGER.info("Adding AutoShape... ") m = AutoShape(self) # wrap model - copy_attr(m, self, include=('yaml', 'nc', 'hyp', 'names', 'stride'), exclude=()) # copy attributes + copy_attr(m, self, include=("yaml", "nc", "hyp", "names", "stride"), exclude=()) # copy attributes return m def info(self, verbose=False, img_size=640): # print model information @@ -249,12 +262,12 @@ def _apply(self, fn): def parse_model(d, ch): # model_dict, input_channels(3) LOGGER.info(f"\n{'':>3}{'from':>18}{'n':>3}{'params':>10} {'module':<40}{'arguments':<30}") - anchors, nc, gd, gw = d['anchors'], d['nc'], d['depth_multiple'], d['width_multiple'] + anchors, nc, gd, gw = d["anchors"], d["nc"], d["depth_multiple"], d["width_multiple"] na = (len(anchors[0]) // 2) if isinstance(anchors, list) else anchors # number of anchors no = na * (nc + 5) # number of outputs = anchors * (classes + 5) layers, save, c2 = [], [], ch[-1] # layers, savelist, ch out - for i, (f, n, m, args) in enumerate(d['backbone'] + d['head']): # from, number, module, args + for i, (f, n, m, args) in enumerate(d["backbone"] + d["head"]): # from, number, module, args m = eval(m) if isinstance(m, str) else m # eval strings for j, a in enumerate(args): try: @@ -263,8 +276,23 @@ def parse_model(d, ch): # model_dict, input_channels(3) pass n = n_ = max(round(n * gd), 1) if n > 1 else n # depth gain - if m in [Conv, GhostConv, Bottleneck, GhostBottleneck, SPP, SPPF, DWConv, MixConv2d, Focus, CrossConv, - BottleneckCSP, C3, C3TR, C3SPP, C3Ghost]: + if m in [ + Conv, + GhostConv, + Bottleneck, + GhostBottleneck, + SPP, + SPPF, + DWConv, + MixConv2d, + Focus, + CrossConv, + BottleneckCSP, + C3, + C3TR, + C3SPP, + C3Ghost, + ]: c1, c2 = ch[f], args[0] if c2 != no: # if not output c2 = make_divisible(c2 * gw, 8) @@ -289,10 +317,10 @@ def parse_model(d, ch): # model_dict, input_channels(3) c2 = ch[f] m_ = nn.Sequential(*(m(*args) for _ in range(n))) if n > 1 else m(*args) # module - t = str(m)[8:-2].replace('__main__.', '') # module type + t = str(m)[8:-2].replace("__main__.", "") # module type np = sum(x.numel() for x in m_.parameters()) # number params m_.i, m_.f, m_.type, m_.np = i, f, t, np # attach index, 'from' index, type, number params - LOGGER.info(f'{i:>3}{str(f):>18}{n_:>3}{np:10.0f} {t:<40}{str(args):<30}') # print + LOGGER.info(f"{i:>3}{str(f):>18}{n_:>3}{np:10.0f} {t:<40}{str(args):<30}") # print save.extend(x % i for x in ([f] if isinstance(f, int) else f) if x != -1) # append to savelist layers.append(m_) if i == 0: @@ -301,12 +329,12 @@ def parse_model(d, ch): # model_dict, input_channels(3) return nn.Sequential(*layers), sorted(save) -if __name__ == '__main__': +if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument('--cfg', type=str, default='yolov5s.yaml', help='model.yaml') - parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') - parser.add_argument('--profile', action='store_true', help='profile model speed') - parser.add_argument('--test', action='store_true', help='test all yolo*.yaml') + parser.add_argument("--cfg", type=str, default="yolov5s.yaml", help="model.yaml") + parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu") + parser.add_argument("--profile", action="store_true", help="profile model speed") + parser.add_argument("--test", action="store_true", help="test all yolo*.yaml") opt = parser.parse_args() opt.cfg = check_yaml(opt.cfg) # check YAML print_args(FILE.stem, opt) @@ -323,11 +351,11 @@ def parse_model(d, ch): # model_dict, input_channels(3) # Test all models if opt.test: - for cfg in Path(ROOT / 'models').rglob('yolo*.yaml'): + for cfg in Path(ROOT / "models").rglob("yolo*.yaml"): try: _ = Model(cfg) except Exception as e: - print(f'Error in {cfg}: {e}') + print(f"Error in {cfg}: {e}") # Tensorboard (not working https://github.com/ultralytics/yolov5/issues/2898) # from torch.utils.tensorboard import SummaryWriter diff --git a/pyproject.toml b/pyproject.toml new file mode 100644 index 000000000000..743134c7e99a --- /dev/null +++ b/pyproject.toml @@ -0,0 +1,4 @@ +[tool.black] +# https://github.com/psf/black +line-length = 120 +exclude = "(.eggs|.git|.hg|.mypy_cache|.venv|_build|buck-out|build|dist)" diff --git a/setup.cfg b/setup.cfg index 4ca0f0d7aabb..602e6e6279f1 100644 --- a/setup.cfg +++ b/setup.cfg @@ -43,6 +43,8 @@ ignore = F403 E302 F541 + W503 + E203 [isort] diff --git a/train.py b/train.py index fedc55d8be5c..8c3ffce76b83 100644 --- a/train.py +++ b/train.py @@ -40,10 +40,28 @@ from utils.callbacks import Callbacks from utils.datasets import create_dataloader from utils.downloads import attempt_download -from utils.general import (LOGGER, check_dataset, check_file, check_git_status, check_img_size, check_requirements, - check_suffix, check_yaml, colorstr, get_latest_run, increment_path, init_seeds, - intersect_dicts, labels_to_class_weights, labels_to_image_weights, methods, one_cycle, - print_args, print_mutation, strip_optimizer) +from utils.general import ( + LOGGER, + check_dataset, + check_file, + check_git_status, + check_img_size, + check_requirements, + check_suffix, + check_yaml, + colorstr, + get_latest_run, + increment_path, + init_seeds, + intersect_dicts, + labels_to_class_weights, + labels_to_image_weights, + methods, + one_cycle, + print_args, + print_mutation, + strip_optimizer, +) from utils.loggers import Loggers from utils.loggers.wandb.wandb_utils import check_wandb_resume from utils.loss import ComputeLoss @@ -51,35 +69,43 @@ from utils.plots import plot_evolve, plot_labels from utils.torch_utils import EarlyStopping, ModelEMA, de_parallel, select_device, torch_distributed_zero_first -LOCAL_RANK = int(os.getenv('LOCAL_RANK', -1)) # https://pytorch.org/docs/stable/elastic/run.html -RANK = int(os.getenv('RANK', -1)) -WORLD_SIZE = int(os.getenv('WORLD_SIZE', 1)) - - -def train(hyp, # path/to/hyp.yaml or hyp dictionary - opt, - device, - callbacks - ): - save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, freeze, = \ - Path(opt.save_dir), opt.epochs, opt.batch_size, opt.weights, opt.single_cls, opt.evolve, opt.data, opt.cfg, \ - opt.resume, opt.noval, opt.nosave, opt.workers, opt.freeze +LOCAL_RANK = int(os.getenv("LOCAL_RANK", -1)) # https://pytorch.org/docs/stable/elastic/run.html +RANK = int(os.getenv("RANK", -1)) +WORLD_SIZE = int(os.getenv("WORLD_SIZE", 1)) + + +def train(hyp, opt, device, callbacks): # path/to/hyp.yaml or hyp dictionary + save_dir, epochs, batch_size, weights, single_cls, evolve, data, cfg, resume, noval, nosave, workers, freeze, = ( + Path(opt.save_dir), + opt.epochs, + opt.batch_size, + opt.weights, + opt.single_cls, + opt.evolve, + opt.data, + opt.cfg, + opt.resume, + opt.noval, + opt.nosave, + opt.workers, + opt.freeze, + ) # Directories - w = save_dir / 'weights' # weights dir + w = save_dir / "weights" # weights dir (w.parent if evolve else w).mkdir(parents=True, exist_ok=True) # make dir - last, best = w / 'last.pt', w / 'best.pt' + last, best = w / "last.pt", w / "best.pt" # Hyperparameters if isinstance(hyp, str): - with open(hyp, errors='ignore') as f: + with open(hyp, errors="ignore") as f: hyp = yaml.safe_load(f) # load hyps dict - LOGGER.info(colorstr('hyperparameters: ') + ', '.join(f'{k}={v}' for k, v in hyp.items())) + LOGGER.info(colorstr("hyperparameters: ") + ", ".join(f"{k}={v}" for k, v in hyp.items())) # Save run settings - with open(save_dir / 'hyp.yaml', 'w') as f: + with open(save_dir / "hyp.yaml", "w") as f: yaml.safe_dump(hyp, f, sort_keys=False) - with open(save_dir / 'opt.yaml', 'w') as f: + with open(save_dir / "opt.yaml", "w") as f: yaml.safe_dump(vars(opt), f, sort_keys=False) data_dict = None @@ -97,38 +123,38 @@ def train(hyp, # path/to/hyp.yaml or hyp dictionary # Config plots = not evolve # create plots - cuda = device.type != 'cpu' + cuda = device.type != "cpu" init_seeds(1 + RANK) with torch_distributed_zero_first(LOCAL_RANK): data_dict = data_dict or check_dataset(data) # check if None - train_path, val_path = data_dict['train'], data_dict['val'] - nc = 1 if single_cls else int(data_dict['nc']) # number of classes - names = ['item'] if single_cls and len(data_dict['names']) != 1 else data_dict['names'] # class names - assert len(names) == nc, f'{len(names)} names found for nc={nc} dataset in {data}' # check - is_coco = isinstance(val_path, str) and val_path.endswith('coco/val2017.txt') # COCO dataset + train_path, val_path = data_dict["train"], data_dict["val"] + nc = 1 if single_cls else int(data_dict["nc"]) # number of classes + names = ["item"] if single_cls and len(data_dict["names"]) != 1 else data_dict["names"] # class names + assert len(names) == nc, f"{len(names)} names found for nc={nc} dataset in {data}" # check + is_coco = isinstance(val_path, str) and val_path.endswith("coco/val2017.txt") # COCO dataset # Model - check_suffix(weights, '.pt') # check weights - pretrained = weights.endswith('.pt') + check_suffix(weights, ".pt") # check weights + pretrained = weights.endswith(".pt") if pretrained: with torch_distributed_zero_first(LOCAL_RANK): weights = attempt_download(weights) # download if not found locally ckpt = torch.load(weights, map_location=device) # load checkpoint - model = Model(cfg or ckpt['model'].yaml, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create - exclude = ['anchor'] if (cfg or hyp.get('anchors')) and not resume else [] # exclude keys - csd = ckpt['model'].float().state_dict() # checkpoint state_dict as FP32 + model = Model(cfg or ckpt["model"].yaml, ch=3, nc=nc, anchors=hyp.get("anchors")).to(device) # create + exclude = ["anchor"] if (cfg or hyp.get("anchors")) and not resume else [] # exclude keys + csd = ckpt["model"].float().state_dict() # checkpoint state_dict as FP32 csd = intersect_dicts(csd, model.state_dict(), exclude=exclude) # intersect model.load_state_dict(csd, strict=False) # load - LOGGER.info(f'Transferred {len(csd)}/{len(model.state_dict())} items from {weights}') # report + LOGGER.info(f"Transferred {len(csd)}/{len(model.state_dict())} items from {weights}") # report else: - model = Model(cfg, ch=3, nc=nc, anchors=hyp.get('anchors')).to(device) # create + model = Model(cfg, ch=3, nc=nc, anchors=hyp.get("anchors")).to(device) # create # Freeze - freeze = [f'model.{x}.' for x in range(freeze)] # layers to freeze + freeze = [f"model.{x}." for x in range(freeze)] # layers to freeze for k, v in model.named_parameters(): v.requires_grad = True # train all layers if any(x in k for x in freeze): - LOGGER.info(f'freezing {k}') + LOGGER.info(f"freezing {k}") v.requires_grad = False # Image size @@ -142,34 +168,36 @@ def train(hyp, # path/to/hyp.yaml or hyp dictionary # Optimizer nbs = 64 # nominal batch size accumulate = max(round(nbs / batch_size), 1) # accumulate loss before optimizing - hyp['weight_decay'] *= batch_size * accumulate / nbs # scale weight_decay + hyp["weight_decay"] *= batch_size * accumulate / nbs # scale weight_decay LOGGER.info(f"Scaled weight_decay = {hyp['weight_decay']}") g0, g1, g2 = [], [], [] # optimizer parameter groups for v in model.modules(): - if hasattr(v, 'bias') and isinstance(v.bias, nn.Parameter): # bias + if hasattr(v, "bias") and isinstance(v.bias, nn.Parameter): # bias g2.append(v.bias) if isinstance(v, nn.BatchNorm2d): # weight (no decay) g0.append(v.weight) - elif hasattr(v, 'weight') and isinstance(v.weight, nn.Parameter): # weight (with decay) + elif hasattr(v, "weight") and isinstance(v.weight, nn.Parameter): # weight (with decay) g1.append(v.weight) if opt.adam: - optimizer = Adam(g0, lr=hyp['lr0'], betas=(hyp['momentum'], 0.999)) # adjust beta1 to momentum + optimizer = Adam(g0, lr=hyp["lr0"], betas=(hyp["momentum"], 0.999)) # adjust beta1 to momentum else: - optimizer = SGD(g0, lr=hyp['lr0'], momentum=hyp['momentum'], nesterov=True) - - optimizer.add_param_group({'params': g1, 'weight_decay': hyp['weight_decay']}) # add g1 with weight_decay - optimizer.add_param_group({'params': g2}) # add g2 (biases) - LOGGER.info(f"{colorstr('optimizer:')} {type(optimizer).__name__} with parameter groups " - f"{len(g0)} weight, {len(g1)} weight (no decay), {len(g2)} bias") + optimizer = SGD(g0, lr=hyp["lr0"], momentum=hyp["momentum"], nesterov=True) + + optimizer.add_param_group({"params": g1, "weight_decay": hyp["weight_decay"]}) # add g1 with weight_decay + optimizer.add_param_group({"params": g2}) # add g2 (biases) + LOGGER.info( + f"{colorstr('optimizer:')} {type(optimizer).__name__} with parameter groups " + f"{len(g0)} weight, {len(g1)} weight (no decay), {len(g2)} bias" + ) del g0, g1, g2 # Scheduler if opt.linear_lr: - lf = lambda x: (1 - x / (epochs - 1)) * (1.0 - hyp['lrf']) + hyp['lrf'] # linear + lf = lambda x: (1 - x / (epochs - 1)) * (1.0 - hyp["lrf"]) + hyp["lrf"] # linear else: - lf = one_cycle(1, hyp['lrf'], epochs) # cosine 1->hyp['lrf'] + lf = one_cycle(1, hyp["lrf"], epochs) # cosine 1->hyp['lrf'] scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf) # plot_lr_scheduler(optimizer, scheduler, epochs) # EMA @@ -179,51 +207,75 @@ def train(hyp, # path/to/hyp.yaml or hyp dictionary 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'] + if ckpt["optimizer"] is not None: + optimizer.load_state_dict(ckpt["optimizer"]) + best_fitness = ckpt["best_fitness"] # EMA - if ema and ckpt.get('ema'): - ema.ema.load_state_dict(ckpt['ema'].float().state_dict()) - ema.updates = ckpt['updates'] + if ema and ckpt.get("ema"): + ema.ema.load_state_dict(ckpt["ema"].float().state_dict()) + ema.updates = ckpt["updates"] # Epochs - start_epoch = ckpt['epoch'] + 1 + start_epoch = ckpt["epoch"] + 1 if resume: - assert start_epoch > 0, f'{weights} training to {epochs} epochs is finished, nothing to resume.' + assert start_epoch > 0, f"{weights} training to {epochs} epochs is finished, nothing to resume." if epochs < start_epoch: LOGGER.info(f"{weights} has been trained for {ckpt['epoch']} epochs. Fine-tuning for {epochs} more epochs.") - epochs += ckpt['epoch'] # finetune additional epochs + epochs += ckpt["epoch"] # finetune additional epochs del ckpt, csd # DP mode if cuda and RANK == -1 and torch.cuda.device_count() > 1: - LOGGER.warning('WARNING: DP not recommended, use torch.distributed.run for best DDP Multi-GPU results.\n' - 'See Multi-GPU Tutorial at https://github.com/ultralytics/yolov5/issues/475 to get started.') + LOGGER.warning( + "WARNING: DP not recommended, use torch.distributed.run for best DDP Multi-GPU results.\n" + "See Multi-GPU Tutorial at https://github.com/ultralytics/yolov5/issues/475 to get started." + ) 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()') + LOGGER.info("Using SyncBatchNorm()") # Trainloader - train_loader, dataset = create_dataloader(train_path, imgsz, batch_size // WORLD_SIZE, gs, single_cls, - hyp=hyp, augment=True, cache=opt.cache, rect=opt.rect, rank=LOCAL_RANK, - workers=workers, image_weights=opt.image_weights, quad=opt.quad, - prefix=colorstr('train: ')) + train_loader, dataset = create_dataloader( + train_path, + imgsz, + batch_size // WORLD_SIZE, + gs, + single_cls, + hyp=hyp, + augment=True, + cache=opt.cache, + rect=opt.rect, + rank=LOCAL_RANK, + workers=workers, + image_weights=opt.image_weights, + quad=opt.quad, + prefix=colorstr("train: "), + ) mlc = int(np.concatenate(dataset.labels, 0)[:, 0].max()) # max label class nb = len(train_loader) # number of batches - assert mlc < nc, f'Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}' + assert mlc < nc, f"Label class {mlc} exceeds nc={nc} in {data}. Possible class labels are 0-{nc - 1}" # Process 0 if RANK in [-1, 0]: - val_loader = create_dataloader(val_path, imgsz, batch_size // WORLD_SIZE * 2, gs, single_cls, - hyp=hyp, cache=None if noval else opt.cache, rect=True, rank=-1, - workers=workers, pad=0.5, - prefix=colorstr('val: '))[0] + val_loader = create_dataloader( + val_path, + imgsz, + batch_size // WORLD_SIZE * 2, + gs, + single_cls, + hyp=hyp, + cache=None if noval else opt.cache, + rect=True, + rank=-1, + workers=workers, + pad=0.5, + prefix=colorstr("val: "), + )[0] if not resume: labels = np.concatenate(dataset.labels, 0) @@ -235,10 +287,10 @@ def train(hyp, # path/to/hyp.yaml or hyp dictionary # Anchors if not opt.noautoanchor: - check_anchors(dataset, model=model, thr=hyp['anchor_t'], imgsz=imgsz) + check_anchors(dataset, model=model, thr=hyp["anchor_t"], imgsz=imgsz) model.half().float() # pre-reduce anchor precision - callbacks.run('on_pretrain_routine_end') + callbacks.run("on_pretrain_routine_end") # DDP mode if cuda and RANK != -1: @@ -246,10 +298,10 @@ def train(hyp, # path/to/hyp.yaml or hyp dictionary # Model parameters nl = de_parallel(model).model[-1].nl # number of detection layers (to scale hyps) - hyp['box'] *= 3 / nl # scale to layers - hyp['cls'] *= nc / 80 * 3 / nl # scale to classes and layers - hyp['obj'] *= (imgsz / 640) ** 2 * 3 / nl # scale to image size and layers - hyp['label_smoothing'] = opt.label_smoothing + hyp["box"] *= 3 / nl # scale to layers + hyp["cls"] *= nc / 80 * 3 / nl # scale to classes and layers + hyp["obj"] *= (imgsz / 640) ** 2 * 3 / nl # scale to image size and layers + hyp["label_smoothing"] = opt.label_smoothing model.nc = nc # attach number of classes to model model.hyp = hyp # attach hyperparameters to model model.class_weights = labels_to_class_weights(dataset.labels, nc).to(device) * nc # attach class weights @@ -257,7 +309,7 @@ def train(hyp, # path/to/hyp.yaml or hyp dictionary # Start training t0 = time.time() - nw = max(round(hyp['warmup_epochs'] * nb), 1000) # number of warmup iterations, max(3 epochs, 1k iterations) + 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 last_opt_step = -1 maps = np.zeros(nc) # mAP per class @@ -266,10 +318,12 @@ def train(hyp, # path/to/hyp.yaml or hyp dictionary scaler = amp.GradScaler(enabled=cuda) stopper = EarlyStopping(patience=opt.patience) compute_loss = ComputeLoss(model) # init loss class - LOGGER.info(f'Image sizes {imgsz} train, {imgsz} val\n' - f'Using {train_loader.num_workers} dataloader workers\n' - f"Logging results to {colorstr('bold', save_dir)}\n" - f'Starting training for {epochs} epochs...') + LOGGER.info( + f"Image sizes {imgsz} train, {imgsz} val\n" + f"Using {train_loader.num_workers} dataloader workers\n" + f"Logging results to {colorstr('bold', save_dir)}\n" + f"Starting training for {epochs} epochs..." + ) for epoch in range(start_epoch, epochs): # epoch ------------------------------------------------------------------ model.train() @@ -287,7 +341,7 @@ def train(hyp, # path/to/hyp.yaml or hyp dictionary if RANK != -1: train_loader.sampler.set_epoch(epoch) pbar = enumerate(train_loader) - LOGGER.info(('\n' + '%10s' * 7) % ('Epoch', 'gpu_mem', 'box', 'obj', 'cls', 'labels', 'img_size')) + LOGGER.info(("\n" + "%10s" * 7) % ("Epoch", "gpu_mem", "box", "obj", "cls", "labels", "img_size")) if RANK in [-1, 0]: pbar = tqdm(pbar, total=nb) # progress bar optimizer.zero_grad() @@ -302,9 +356,9 @@ def train(hyp, # path/to/hyp.yaml or hyp dictionary accumulate = max(1, np.interp(ni, xi, [1, nbs / 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']]) + 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: @@ -312,7 +366,7 @@ def train(hyp, # path/to/hyp.yaml or hyp dictionary 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 = nn.functional.interpolate(imgs, size=ns, mode='bilinear', align_corners=False) + imgs = nn.functional.interpolate(imgs, size=ns, mode="bilinear", align_corners=False) # Forward with amp.autocast(enabled=cuda): @@ -321,7 +375,7 @@ def train(hyp, # path/to/hyp.yaml or hyp dictionary if RANK != -1: loss *= WORLD_SIZE # gradient averaged between devices in DDP mode if opt.quad: - loss *= 4. + loss *= 4.0 # Backward scaler.scale(loss).backward() @@ -338,59 +392,65 @@ def train(hyp, # path/to/hyp.yaml or hyp dictionary # Log if RANK in [-1, 0]: mloss = (mloss * i + loss_items) / (i + 1) # update mean losses - mem = f'{torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0:.3g}G' # (GB) - pbar.set_description(('%10s' * 2 + '%10.4g' * 5) % ( - f'{epoch}/{epochs - 1}', mem, *mloss, targets.shape[0], imgs.shape[-1])) - callbacks.run('on_train_batch_end', ni, model, imgs, targets, paths, plots, opt.sync_bn) + mem = f"{torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0:.3g}G" # (GB) + pbar.set_description( + ("%10s" * 2 + "%10.4g" * 5) + % (f"{epoch}/{epochs - 1}", mem, *mloss, targets.shape[0], imgs.shape[-1]) + ) + callbacks.run("on_train_batch_end", ni, model, imgs, targets, paths, plots, opt.sync_bn) # end batch ------------------------------------------------------------------------------------------------ # Scheduler - lr = [x['lr'] for x in optimizer.param_groups] # for loggers + lr = [x["lr"] for x in optimizer.param_groups] # for loggers scheduler.step() if RANK in [-1, 0]: # mAP - callbacks.run('on_train_epoch_end', epoch=epoch) - ema.update_attr(model, include=['yaml', 'nc', 'hyp', 'names', 'stride', 'class_weights']) + callbacks.run("on_train_epoch_end", epoch=epoch) + ema.update_attr(model, include=["yaml", "nc", "hyp", "names", "stride", "class_weights"]) final_epoch = (epoch + 1 == epochs) or stopper.possible_stop if not noval or final_epoch: # Calculate mAP - results, maps, _ = val.run(data_dict, - batch_size=batch_size // WORLD_SIZE * 2, - imgsz=imgsz, - model=ema.ema, - single_cls=single_cls, - dataloader=val_loader, - save_dir=save_dir, - plots=False, - callbacks=callbacks, - compute_loss=compute_loss) + results, maps, _ = val.run( + data_dict, + batch_size=batch_size // WORLD_SIZE * 2, + imgsz=imgsz, + model=ema.ema, + single_cls=single_cls, + dataloader=val_loader, + save_dir=save_dir, + plots=False, + callbacks=callbacks, + compute_loss=compute_loss, + ) # Update best mAP fi = fitness(np.array(results).reshape(1, -1)) # weighted combination of [P, R, mAP@.5, mAP@.5-.95] if fi > best_fitness: best_fitness = fi log_vals = list(mloss) + list(results) + lr - callbacks.run('on_fit_epoch_end', log_vals, epoch, best_fitness, fi) + callbacks.run("on_fit_epoch_end", log_vals, epoch, best_fitness, fi) # Save model if (not nosave) or (final_epoch and not evolve): # if save - ckpt = {'epoch': epoch, - 'best_fitness': best_fitness, - 'model': deepcopy(de_parallel(model)).half(), - 'ema': deepcopy(ema.ema).half(), - 'updates': ema.updates, - 'optimizer': optimizer.state_dict(), - 'wandb_id': loggers.wandb.wandb_run.id if loggers.wandb else None, - 'date': datetime.now().isoformat()} + ckpt = { + "epoch": epoch, + "best_fitness": best_fitness, + "model": deepcopy(de_parallel(model)).half(), + "ema": deepcopy(ema.ema).half(), + "updates": ema.updates, + "optimizer": optimizer.state_dict(), + "wandb_id": loggers.wandb.wandb_run.id if loggers.wandb else None, + "date": datetime.now().isoformat(), + } # Save last, best and delete torch.save(ckpt, last) if best_fitness == fi: torch.save(ckpt, best) if (epoch > 0) and (opt.save_period > 0) and (epoch % opt.save_period == 0): - torch.save(ckpt, w / f'epoch{epoch}.pt') + torch.save(ckpt, w / f"epoch{epoch}.pt") del ckpt - callbacks.run('on_model_save', last, epoch, final_epoch, best_fitness, fi) + callbacks.run("on_model_save", last, epoch, final_epoch, best_fitness, fi) # Stop Single-GPU if RANK == -1 and stopper(epoch=epoch, fitness=fi): @@ -409,29 +469,31 @@ def train(hyp, # path/to/hyp.yaml or hyp dictionary # end epoch ---------------------------------------------------------------------------------------------------- # end training ----------------------------------------------------------------------------------------------------- if RANK in [-1, 0]: - LOGGER.info(f'\n{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours.') + LOGGER.info(f"\n{epoch - start_epoch + 1} epochs completed in {(time.time() - t0) / 3600:.3f} hours.") for f in last, best: if f.exists(): strip_optimizer(f) # strip optimizers if f is best: - LOGGER.info(f'\nValidating {f}...') - results, _, _ = val.run(data_dict, - batch_size=batch_size // WORLD_SIZE * 2, - imgsz=imgsz, - model=attempt_load(f, device).half(), - iou_thres=0.65 if is_coco else 0.60, # best pycocotools results at 0.65 - single_cls=single_cls, - dataloader=val_loader, - save_dir=save_dir, - save_json=is_coco, - verbose=True, - plots=True, - callbacks=callbacks, - compute_loss=compute_loss) # val best model with plots + LOGGER.info(f"\nValidating {f}...") + results, _, _ = val.run( + data_dict, + batch_size=batch_size // WORLD_SIZE * 2, + imgsz=imgsz, + model=attempt_load(f, device).half(), + iou_thres=0.65 if is_coco else 0.60, # best pycocotools results at 0.65 + single_cls=single_cls, + dataloader=val_loader, + save_dir=save_dir, + save_json=is_coco, + verbose=True, + plots=True, + callbacks=callbacks, + compute_loss=compute_loss, + ) # val best model with plots if is_coco: - callbacks.run('on_fit_epoch_end', list(mloss) + list(results) + lr, epoch, best_fitness, fi) + callbacks.run("on_fit_epoch_end", list(mloss) + list(results) + lr, epoch, best_fitness, fi) - callbacks.run('on_train_end', last, best, plots, epoch, results) + callbacks.run("on_train_end", last, best, plots, epoch, results) LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}") torch.cuda.empty_cache() @@ -440,44 +502,44 @@ def train(hyp, # path/to/hyp.yaml or hyp dictionary def parse_opt(known=False): parser = argparse.ArgumentParser() - parser.add_argument('--weights', type=str, default=ROOT / 'yolov5s.pt', help='initial weights path') - parser.add_argument('--cfg', type=str, default='', help='model.yaml path') - parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path') - parser.add_argument('--hyp', type=str, default=ROOT / 'data/hyps/hyp.scratch.yaml', help='hyperparameters path') - parser.add_argument('--epochs', type=int, default=300) - parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs, -1 for autobatch') - parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='train, val image size (pixels)') - parser.add_argument('--rect', action='store_true', help='rectangular training') - parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training') - parser.add_argument('--nosave', action='store_true', help='only save final checkpoint') - parser.add_argument('--noval', action='store_true', help='only validate final epoch') - parser.add_argument('--noautoanchor', action='store_true', help='disable autoanchor check') - parser.add_argument('--evolve', type=int, nargs='?', const=300, help='evolve hyperparameters for x generations') - parser.add_argument('--bucket', type=str, default='', help='gsutil bucket') - parser.add_argument('--cache', type=str, nargs='?', const='ram', help='--cache images in "ram" (default) or "disk"') - parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training') - parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') - parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%') - parser.add_argument('--single-cls', action='store_true', help='train multi-class data as single-class') - parser.add_argument('--adam', action='store_true', help='use torch.optim.Adam() optimizer') - parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode') - parser.add_argument('--workers', type=int, default=8, help='maximum number of dataloader workers') - parser.add_argument('--project', default=ROOT / 'runs/train', help='save to project/name') - parser.add_argument('--name', default='exp', help='save to project/name') - parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment') - parser.add_argument('--quad', action='store_true', help='quad dataloader') - parser.add_argument('--linear-lr', action='store_true', help='linear LR') - parser.add_argument('--label-smoothing', type=float, default=0.0, help='Label smoothing epsilon') - parser.add_argument('--patience', type=int, default=100, help='EarlyStopping patience (epochs without improvement)') - parser.add_argument('--freeze', type=int, default=0, help='Number of layers to freeze. backbone=10, all=24') - parser.add_argument('--save-period', type=int, default=-1, help='Save checkpoint every x epochs (disabled if < 1)') - parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter, do not modify') + parser.add_argument("--weights", type=str, default=ROOT / "yolov5s.pt", help="initial weights path") + parser.add_argument("--cfg", type=str, default="", help="model.yaml path") + parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="dataset.yaml path") + parser.add_argument("--hyp", type=str, default=ROOT / "data/hyps/hyp.scratch.yaml", help="hyperparameters path") + parser.add_argument("--epochs", type=int, default=300) + parser.add_argument("--batch-size", type=int, default=16, help="total batch size for all GPUs, -1 for autobatch") + parser.add_argument("--imgsz", "--img", "--img-size", type=int, default=640, help="train, val image size (pixels)") + parser.add_argument("--rect", action="store_true", help="rectangular training") + parser.add_argument("--resume", nargs="?", const=True, default=False, help="resume most recent training") + parser.add_argument("--nosave", action="store_true", help="only save final checkpoint") + parser.add_argument("--noval", action="store_true", help="only validate final epoch") + parser.add_argument("--noautoanchor", action="store_true", help="disable autoanchor check") + parser.add_argument("--evolve", type=int, nargs="?", const=300, help="evolve hyperparameters for x generations") + parser.add_argument("--bucket", type=str, default="", help="gsutil bucket") + parser.add_argument("--cache", type=str, nargs="?", const="ram", help='--cache images in "ram" (default) or "disk"') + parser.add_argument("--image-weights", action="store_true", help="use weighted image selection for training") + parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu") + parser.add_argument("--multi-scale", action="store_true", help="vary img-size +/- 50%%") + parser.add_argument("--single-cls", action="store_true", help="train multi-class data as single-class") + parser.add_argument("--adam", action="store_true", help="use torch.optim.Adam() optimizer") + parser.add_argument("--sync-bn", action="store_true", help="use SyncBatchNorm, only available in DDP mode") + parser.add_argument("--workers", type=int, default=8, help="maximum number of dataloader workers") + parser.add_argument("--project", default=ROOT / "runs/train", help="save to project/name") + parser.add_argument("--name", default="exp", help="save to project/name") + parser.add_argument("--exist-ok", action="store_true", help="existing project/name ok, do not increment") + parser.add_argument("--quad", action="store_true", help="quad dataloader") + parser.add_argument("--linear-lr", action="store_true", help="linear LR") + parser.add_argument("--label-smoothing", type=float, default=0.0, help="Label smoothing epsilon") + parser.add_argument("--patience", type=int, default=100, help="EarlyStopping patience (epochs without improvement)") + parser.add_argument("--freeze", type=int, default=0, help="Number of layers to freeze. backbone=10, all=24") + parser.add_argument("--save-period", type=int, default=-1, help="Save checkpoint every x epochs (disabled if < 1)") + parser.add_argument("--local_rank", type=int, default=-1, help="DDP parameter, do not modify") # Weights & Biases arguments - parser.add_argument('--entity', default=None, help='W&B: Entity') - parser.add_argument('--upload_dataset', action='store_true', help='W&B: Upload dataset as artifact table') - parser.add_argument('--bbox_interval', type=int, default=-1, help='W&B: Set bounding-box image logging interval') - parser.add_argument('--artifact_alias', type=str, default='latest', help='W&B: Version of dataset artifact to use') + parser.add_argument("--entity", default=None, help="W&B: Entity") + parser.add_argument("--upload_dataset", action="store_true", help="W&B: Upload dataset as artifact table") + parser.add_argument("--bbox_interval", type=int, default=-1, help="W&B: Set bounding-box image logging interval") + parser.add_argument("--artifact_alias", type=str, default="latest", help="W&B: Version of dataset artifact to use") opt = parser.parse_known_args()[0] if known else parser.parse_args() return opt @@ -488,98 +550,105 @@ def main(opt, callbacks=Callbacks()): if RANK in [-1, 0]: print_args(FILE.stem, opt) check_git_status() - check_requirements(exclude=['thop']) + check_requirements(exclude=["thop"]) # Resume if opt.resume and not check_wandb_resume(opt) and not opt.evolve: # resume an interrupted run ckpt = opt.resume if isinstance(opt.resume, str) else get_latest_run() # specified or most recent path - assert os.path.isfile(ckpt), 'ERROR: --resume checkpoint does not exist' - with open(Path(ckpt).parent.parent / 'opt.yaml', errors='ignore') as f: + assert os.path.isfile(ckpt), "ERROR: --resume checkpoint does not exist" + with open(Path(ckpt).parent.parent / "opt.yaml", errors="ignore") as f: opt = argparse.Namespace(**yaml.safe_load(f)) # replace - opt.cfg, opt.weights, opt.resume = '', ckpt, True # reinstate - LOGGER.info(f'Resuming training from {ckpt}') + opt.cfg, opt.weights, opt.resume = "", ckpt, True # reinstate + LOGGER.info(f"Resuming training from {ckpt}") else: - opt.data, opt.cfg, opt.hyp, opt.weights, opt.project = \ - check_file(opt.data), check_yaml(opt.cfg), check_yaml(opt.hyp), str(opt.weights), str(opt.project) # checks - assert len(opt.cfg) or len(opt.weights), 'either --cfg or --weights must be specified' + opt.data, opt.cfg, opt.hyp, opt.weights, opt.project = ( + check_file(opt.data), + check_yaml(opt.cfg), + check_yaml(opt.hyp), + str(opt.weights), + str(opt.project), + ) # checks + assert len(opt.cfg) or len(opt.weights), "either --cfg or --weights must be specified" if opt.evolve: - opt.project = str(ROOT / 'runs/evolve') + opt.project = str(ROOT / "runs/evolve") opt.exist_ok, opt.resume = opt.resume, False # pass resume to exist_ok and disable resume opt.save_dir = str(increment_path(Path(opt.project) / opt.name, exist_ok=opt.exist_ok)) # DDP mode device = select_device(opt.device, batch_size=opt.batch_size) if LOCAL_RANK != -1: - assert torch.cuda.device_count() > LOCAL_RANK, 'insufficient CUDA devices for DDP command' - assert opt.batch_size % WORLD_SIZE == 0, '--batch-size must be multiple of CUDA device count' - assert not opt.image_weights, '--image-weights argument is not compatible with DDP training' - assert not opt.evolve, '--evolve argument is not compatible with DDP training' + assert torch.cuda.device_count() > LOCAL_RANK, "insufficient CUDA devices for DDP command" + assert opt.batch_size % WORLD_SIZE == 0, "--batch-size must be multiple of CUDA device count" + assert not opt.image_weights, "--image-weights argument is not compatible with DDP training" + assert not opt.evolve, "--evolve argument is not compatible with DDP training" torch.cuda.set_device(LOCAL_RANK) - device = torch.device('cuda', LOCAL_RANK) + device = torch.device("cuda", LOCAL_RANK) dist.init_process_group(backend="nccl" if dist.is_nccl_available() else "gloo") # Train if not opt.evolve: train(opt.hyp, opt, device, callbacks) if WORLD_SIZE > 1 and RANK == 0: - LOGGER.info('Destroying process group... ') + LOGGER.info("Destroying process group... ") dist.destroy_process_group() # Evolve hyperparameters (optional) else: # Hyperparameter evolution metadata (mutation scale 0-1, lower_limit, upper_limit) - meta = {'lr0': (1, 1e-5, 1e-1), # initial learning rate (SGD=1E-2, Adam=1E-3) - 'lrf': (1, 0.01, 1.0), # final OneCycleLR learning rate (lr0 * lrf) - 'momentum': (0.3, 0.6, 0.98), # SGD momentum/Adam beta1 - 'weight_decay': (1, 0.0, 0.001), # optimizer weight decay - 'warmup_epochs': (1, 0.0, 5.0), # warmup epochs (fractions ok) - 'warmup_momentum': (1, 0.0, 0.95), # warmup initial momentum - 'warmup_bias_lr': (1, 0.0, 0.2), # warmup initial bias lr - 'box': (1, 0.02, 0.2), # box loss gain - 'cls': (1, 0.2, 4.0), # cls loss gain - 'cls_pw': (1, 0.5, 2.0), # cls BCELoss positive_weight - 'obj': (1, 0.2, 4.0), # obj loss gain (scale with pixels) - 'obj_pw': (1, 0.5, 2.0), # obj BCELoss positive_weight - 'iou_t': (0, 0.1, 0.7), # IoU training threshold - 'anchor_t': (1, 2.0, 8.0), # anchor-multiple threshold - 'anchors': (2, 2.0, 10.0), # anchors per output grid (0 to ignore) - 'fl_gamma': (0, 0.0, 2.0), # focal loss gamma (efficientDet default gamma=1.5) - 'hsv_h': (1, 0.0, 0.1), # image HSV-Hue augmentation (fraction) - 'hsv_s': (1, 0.0, 0.9), # image HSV-Saturation augmentation (fraction) - 'hsv_v': (1, 0.0, 0.9), # image HSV-Value augmentation (fraction) - 'degrees': (1, 0.0, 45.0), # image rotation (+/- deg) - 'translate': (1, 0.0, 0.9), # image translation (+/- fraction) - 'scale': (1, 0.0, 0.9), # image scale (+/- gain) - 'shear': (1, 0.0, 10.0), # image shear (+/- deg) - 'perspective': (0, 0.0, 0.001), # image perspective (+/- fraction), range 0-0.001 - 'flipud': (1, 0.0, 1.0), # image flip up-down (probability) - 'fliplr': (0, 0.0, 1.0), # image flip left-right (probability) - 'mosaic': (1, 0.0, 1.0), # image mixup (probability) - 'mixup': (1, 0.0, 1.0), # image mixup (probability) - 'copy_paste': (1, 0.0, 1.0)} # segment copy-paste (probability) - - with open(opt.hyp, errors='ignore') as f: + meta = { + "lr0": (1, 1e-5, 1e-1), # initial learning rate (SGD=1E-2, Adam=1E-3) + "lrf": (1, 0.01, 1.0), # final OneCycleLR learning rate (lr0 * lrf) + "momentum": (0.3, 0.6, 0.98), # SGD momentum/Adam beta1 + "weight_decay": (1, 0.0, 0.001), # optimizer weight decay + "warmup_epochs": (1, 0.0, 5.0), # warmup epochs (fractions ok) + "warmup_momentum": (1, 0.0, 0.95), # warmup initial momentum + "warmup_bias_lr": (1, 0.0, 0.2), # warmup initial bias lr + "box": (1, 0.02, 0.2), # box loss gain + "cls": (1, 0.2, 4.0), # cls loss gain + "cls_pw": (1, 0.5, 2.0), # cls BCELoss positive_weight + "obj": (1, 0.2, 4.0), # obj loss gain (scale with pixels) + "obj_pw": (1, 0.5, 2.0), # obj BCELoss positive_weight + "iou_t": (0, 0.1, 0.7), # IoU training threshold + "anchor_t": (1, 2.0, 8.0), # anchor-multiple threshold + "anchors": (2, 2.0, 10.0), # anchors per output grid (0 to ignore) + "fl_gamma": (0, 0.0, 2.0), # focal loss gamma (efficientDet default gamma=1.5) + "hsv_h": (1, 0.0, 0.1), # image HSV-Hue augmentation (fraction) + "hsv_s": (1, 0.0, 0.9), # image HSV-Saturation augmentation (fraction) + "hsv_v": (1, 0.0, 0.9), # image HSV-Value augmentation (fraction) + "degrees": (1, 0.0, 45.0), # image rotation (+/- deg) + "translate": (1, 0.0, 0.9), # image translation (+/- fraction) + "scale": (1, 0.0, 0.9), # image scale (+/- gain) + "shear": (1, 0.0, 10.0), # image shear (+/- deg) + "perspective": (0, 0.0, 0.001), # image perspective (+/- fraction), range 0-0.001 + "flipud": (1, 0.0, 1.0), # image flip up-down (probability) + "fliplr": (0, 0.0, 1.0), # image flip left-right (probability) + "mosaic": (1, 0.0, 1.0), # image mixup (probability) + "mixup": (1, 0.0, 1.0), # image mixup (probability) + "copy_paste": (1, 0.0, 1.0), + } # segment copy-paste (probability) + + with open(opt.hyp, errors="ignore") as f: hyp = yaml.safe_load(f) # load hyps dict - if 'anchors' not in hyp: # anchors commented in hyp.yaml - hyp['anchors'] = 3 + if "anchors" not in hyp: # anchors commented in hyp.yaml + hyp["anchors"] = 3 opt.noval, opt.nosave, save_dir = True, True, Path(opt.save_dir) # only val/save final epoch # ei = [isinstance(x, (int, float)) for x in hyp.values()] # evolvable indices - evolve_yaml, evolve_csv = save_dir / 'hyp_evolve.yaml', save_dir / 'evolve.csv' + evolve_yaml, evolve_csv = save_dir / "hyp_evolve.yaml", save_dir / "evolve.csv" if opt.bucket: - os.system(f'gsutil cp gs://{opt.bucket}/evolve.csv {save_dir}') # download evolve.csv if exists + os.system(f"gsutil cp gs://{opt.bucket}/evolve.csv {save_dir}") # download evolve.csv if exists for _ in range(opt.evolve): # generations to evolve if evolve_csv.exists(): # if evolve.csv exists: select best hyps and mutate # Select parent(s) - parent = 'single' # parent selection method: 'single' or 'weighted' - x = np.loadtxt(evolve_csv, ndmin=2, delimiter=',', skiprows=1) + parent = "single" # parent selection method: 'single' or 'weighted' + x = np.loadtxt(evolve_csv, ndmin=2, delimiter=",", skiprows=1) n = min(5, len(x)) # number of previous results to consider x = x[np.argsort(-fitness(x))][:n] # top n mutations - w = fitness(x) - fitness(x).min() + 1E-6 # weights (sum > 0) - if parent == 'single' or len(x) == 1: + w = fitness(x) - fitness(x).min() + 1e-6 # weights (sum > 0) + if parent == "single" or len(x) == 1: # x = x[random.randint(0, n - 1)] # random selection x = x[random.choices(range(n), weights=w)[0]] # weighted selection - elif parent == 'weighted': + elif parent == "weighted": x = (x * w.reshape(n, 1)).sum(0) / w.sum() # weighted combination # Mutate @@ -608,9 +677,11 @@ def main(opt, callbacks=Callbacks()): # Plot results plot_evolve(evolve_csv) - LOGGER.info(f'Hyperparameter evolution finished\n' - f"Results saved to {colorstr('bold', save_dir)}\n" - f'Use best hyperparameters example: $ python train.py --hyp {evolve_yaml}') + LOGGER.info( + f"Hyperparameter evolution finished\n" + f"Results saved to {colorstr('bold', save_dir)}\n" + f"Use best hyperparameters example: $ python train.py --hyp {evolve_yaml}" + ) def run(**kwargs): diff --git a/utils/activations.py b/utils/activations.py index 4c7d46c32104..5441b8d648d1 100644 --- a/utils/activations.py +++ b/utils/activations.py @@ -60,7 +60,7 @@ def forward(self, x): # ACON https://arxiv.org/pdf/2009.04759.pdf ---------------------------------------------------------------------------- class AconC(nn.Module): - r""" ACON activation (activate or not). + r"""ACON activation (activate or not). AconC: (p1*x-p2*x) * sigmoid(beta*(p1*x-p2*x)) + p2*x, beta is a learnable parameter according to "Activate or Not: Learning Customized Activation" . """ @@ -77,7 +77,7 @@ def forward(self, x): class MetaAconC(nn.Module): - r""" ACON activation (activate or not). + r"""ACON activation (activate or not). MetaAconC: (p1*x-p2*x) * sigmoid(beta*(p1*x-p2*x)) + p2*x, beta is generated by a small network according to "Activate or Not: Learning Customized Activation" . """ diff --git a/utils/augmentations.py b/utils/augmentations.py index 5dcfd49fdd05..bf1fbd159685 100644 --- a/utils/augmentations.py +++ b/utils/augmentations.py @@ -19,28 +19,32 @@ def __init__(self): self.transform = None try: import albumentations as A - check_version(A.__version__, '1.0.3', hard=True) # version requirement - - self.transform = A.Compose([ - A.Blur(p=0.01), - A.MedianBlur(p=0.01), - A.ToGray(p=0.01), - A.CLAHE(p=0.01), - A.RandomBrightnessContrast(p=0.0), - A.RandomGamma(p=0.0), - A.ImageCompression(quality_lower=75, p=0.0)], - bbox_params=A.BboxParams(format='yolo', label_fields=['class_labels'])) - - LOGGER.info(colorstr('albumentations: ') + ', '.join(f'{x}' for x in self.transform.transforms if x.p)) + + check_version(A.__version__, "1.0.3", hard=True) # version requirement + + self.transform = A.Compose( + [ + A.Blur(p=0.01), + A.MedianBlur(p=0.01), + A.ToGray(p=0.01), + A.CLAHE(p=0.01), + A.RandomBrightnessContrast(p=0.0), + A.RandomGamma(p=0.0), + A.ImageCompression(quality_lower=75, p=0.0), + ], + bbox_params=A.BboxParams(format="yolo", label_fields=["class_labels"]), + ) + + LOGGER.info(colorstr("albumentations: ") + ", ".join(f"{x}" for x in self.transform.transforms if x.p)) except ImportError: # package not installed, skip pass except Exception as e: - LOGGER.info(colorstr('albumentations: ') + f'{e}') + LOGGER.info(colorstr("albumentations: ") + f"{e}") def __call__(self, im, labels, p=1.0): if self.transform and random.random() < p: new = self.transform(image=im, bboxes=labels[:, 1:], class_labels=labels[:, 0]) # transformed - im, labels = new['image'], np.array([[c, *b] for c, b in zip(new['class_labels'], new['bboxes'])]) + im, labels = new["image"], np.array([[c, *b] for c, b in zip(new["class_labels"], new["bboxes"])]) return im, labels @@ -77,7 +81,7 @@ def replicate(im, labels): boxes = labels[:, 1:].astype(int) x1, y1, x2, y2 = boxes.T s = ((x2 - x1) + (y2 - y1)) / 2 # side length (pixels) - for i in s.argsort()[:round(s.size * 0.5)]: # smallest indices + for i in s.argsort()[: round(s.size * 0.5)]: # smallest indices x1b, y1b, x2b, y2b = boxes[i] bh, bw = y2b - y1b, x2b - x1b yc, xc = int(random.uniform(0, h - bh)), int(random.uniform(0, w - bw)) # offset x, y @@ -121,8 +125,9 @@ def letterbox(im, new_shape=(640, 640), color=(114, 114, 114), auto=True, scaleF return im, ratio, (dw, dh) -def random_perspective(im, targets=(), segments=(), degrees=10, translate=.1, scale=.1, shear=10, perspective=0.0, - border=(0, 0)): +def random_perspective( + im, targets=(), segments=(), degrees=10, translate=0.1, scale=0.1, shear=10, perspective=0.0, border=(0, 0) +): # torchvision.transforms.RandomAffine(degrees=(-10, 10), translate=(0.1, 0.1), scale=(0.9, 1.1), shear=(-10, 10)) # targets = [cls, xyxy] diff --git a/utils/autoanchor.py b/utils/autoanchor.py index af0aa7de65ac..3e0aecd9af8b 100644 --- a/utils/autoanchor.py +++ b/utils/autoanchor.py @@ -19,15 +19,15 @@ def check_anchor_order(m): da = a[-1] - a[0] # delta a ds = m.stride[-1] - m.stride[0] # delta s if da.sign() != ds.sign(): # same order - print('Reversing anchor order') + print("Reversing anchor order") m.anchors[:] = m.anchors.flip(0) def check_anchors(dataset, model, thr=4.0, imgsz=640): # Check anchor fit to data, recompute if necessary - prefix = colorstr('autoanchor: ') - print(f'\n{prefix}Analyzing anchors... ', end='') - m = model.module.model[-1] if hasattr(model, 'module') else model.model[-1] # Detect() + prefix = colorstr("autoanchor: ") + print(f"\n{prefix}Analyzing anchors... ", end="") + m = model.module.model[-1] if hasattr(model, "module") else model.model[-1] # Detect() shapes = imgsz * dataset.shapes / dataset.shapes.max(1, keepdims=True) scale = np.random.uniform(0.9, 1.1, size=(shapes.shape[0], 1)) # augment scale wh = torch.tensor(np.concatenate([l[:, 3:5] * s for s, l in zip(shapes * scale, dataset.labels)])).float() # wh @@ -42,46 +42,46 @@ def metric(k): # compute metric anchors = m.anchors.clone() * m.stride.to(m.anchors.device).view(-1, 1, 1) # current anchors bpr, aat = metric(anchors.cpu().view(-1, 2)) - print(f'anchors/target = {aat:.2f}, Best Possible Recall (BPR) = {bpr:.4f}', end='') + print(f"anchors/target = {aat:.2f}, Best Possible Recall (BPR) = {bpr:.4f}", end="") if bpr < 0.98: # threshold to recompute - print('. Attempting to improve anchors, please wait...') + print(". Attempting to improve anchors, please wait...") na = m.anchors.numel() // 2 # number of anchors try: anchors = kmean_anchors(dataset, n=na, img_size=imgsz, thr=thr, gen=1000, verbose=False) except Exception as e: - print(f'{prefix}ERROR: {e}') + print(f"{prefix}ERROR: {e}") new_bpr = metric(anchors)[0] if new_bpr > bpr: # replace anchors anchors = torch.tensor(anchors, device=m.anchors.device).type_as(m.anchors) m.anchors[:] = anchors.clone().view_as(m.anchors) / m.stride.to(m.anchors.device).view(-1, 1, 1) # loss check_anchor_order(m) - print(f'{prefix}New anchors saved to model. Update model *.yaml to use these anchors in the future.') + print(f"{prefix}New anchors saved to model. Update model *.yaml to use these anchors in the future.") else: - print(f'{prefix}Original anchors better than new anchors. Proceeding with original anchors.') - print('') # newline + print(f"{prefix}Original anchors better than new anchors. Proceeding with original anchors.") + print("") # newline -def kmean_anchors(dataset='./data/coco128.yaml', n=9, img_size=640, thr=4.0, gen=1000, verbose=True): - """ Creates kmeans-evolved anchors from training dataset +def kmean_anchors(dataset="./data/coco128.yaml", n=9, img_size=640, thr=4.0, gen=1000, verbose=True): + """Creates kmeans-evolved anchors from training dataset - Arguments: - dataset: path to data.yaml, or a loaded dataset - n: number of anchors - img_size: image size used for training - thr: anchor-label wh ratio threshold hyperparameter hyp['anchor_t'] used for training, default=4.0 - gen: generations to evolve anchors using genetic algorithm - verbose: print all results + Arguments: + dataset: path to data.yaml, or a loaded dataset + n: number of anchors + img_size: image size used for training + thr: anchor-label wh ratio threshold hyperparameter hyp['anchor_t'] used for training, default=4.0 + gen: generations to evolve anchors using genetic algorithm + verbose: print all results - Return: - k: kmeans evolved anchors + Return: + k: kmeans evolved anchors - Usage: - from utils.autoanchor import *; _ = kmean_anchors() + Usage: + from utils.autoanchor import *; _ = kmean_anchors() """ from scipy.cluster.vq import kmeans thr = 1 / thr - prefix = colorstr('autoanchor: ') + prefix = colorstr("autoanchor: ") def metric(k, wh): # compute metrics r = wh[:, None] / k[None] @@ -97,18 +97,22 @@ def print_results(k): k = k[np.argsort(k.prod(1))] # sort small to large x, best = metric(k, wh0) bpr, aat = (best > thr).float().mean(), (x > thr).float().mean() * n # best possible recall, anch > thr - print(f'{prefix}thr={thr:.2f}: {bpr:.4f} best possible recall, {aat:.2f} anchors past thr') - print(f'{prefix}n={n}, img_size={img_size}, metric_all={x.mean():.3f}/{best.mean():.3f}-mean/best, ' - f'past_thr={x[x > thr].mean():.3f}-mean: ', end='') + print(f"{prefix}thr={thr:.2f}: {bpr:.4f} best possible recall, {aat:.2f} anchors past thr") + print( + f"{prefix}n={n}, img_size={img_size}, metric_all={x.mean():.3f}/{best.mean():.3f}-mean/best, " + f"past_thr={x[x > thr].mean():.3f}-mean: ", + end="", + ) for i, x in enumerate(k): - print('%i,%i' % (round(x[0]), round(x[1])), end=', ' if i < len(k) - 1 else '\n') # use in *.cfg + print("%i,%i" % (round(x[0]), round(x[1])), end=", " if i < len(k) - 1 else "\n") # use in *.cfg return k if isinstance(dataset, str): # *.yaml file - with open(dataset, errors='ignore') as f: + with open(dataset, errors="ignore") as f: data_dict = yaml.safe_load(f) # model dict from utils.datasets import LoadImagesAndLabels - dataset = LoadImagesAndLabels(data_dict['train'], augment=True, rect=True) + + dataset = LoadImagesAndLabels(data_dict["train"], augment=True, rect=True) # Get label wh shapes = img_size * dataset.shapes / dataset.shapes.max(1, keepdims=True) @@ -117,15 +121,15 @@ def print_results(k): # Filter i = (wh0 < 3.0).any(1).sum() if i: - print(f'{prefix}WARNING: Extremely small objects found. {i} of {len(wh0)} labels are < 3 pixels in size.') + print(f"{prefix}WARNING: Extremely small objects found. {i} of {len(wh0)} labels are < 3 pixels in size.") wh = wh0[(wh0 >= 2.0).any(1)] # filter > 2 pixels # wh = wh * (np.random.rand(wh.shape[0], 1) * 0.9 + 0.1) # multiply by random scale 0-1 # Kmeans calculation - print(f'{prefix}Running kmeans for {n} anchors on {len(wh)} points...') + print(f"{prefix}Running kmeans for {n} anchors on {len(wh)} points...") s = wh.std(0) # sigmas for whitening k, dist = kmeans(wh / s, n, iter=30) # points, mean distance - assert len(k) == n, f'{prefix}ERROR: scipy.cluster.vq.kmeans requested {n} points but returned only {len(k)}' + assert len(k) == n, f"{prefix}ERROR: scipy.cluster.vq.kmeans requested {n} points but returned only {len(k)}" k *= s wh = torch.tensor(wh, dtype=torch.float32) # filtered wh0 = torch.tensor(wh0, dtype=torch.float32) # unfiltered @@ -146,7 +150,7 @@ def print_results(k): # Evolve npr = np.random f, sh, mp, s = anchor_fitness(k), k.shape, 0.9, 0.1 # fitness, generations, mutation prob, sigma - pbar = tqdm(range(gen), desc=f'{prefix}Evolving anchors with Genetic Algorithm:') # progress bar + pbar = tqdm(range(gen), desc=f"{prefix}Evolving anchors with Genetic Algorithm:") # progress bar for _ in pbar: v = np.ones(sh) while (v == 1).all(): # mutate until a change occurs (prevent duplicates) @@ -155,7 +159,7 @@ def print_results(k): fg = anchor_fitness(kg) if fg > f: f, k = fg, kg.copy() - pbar.desc = f'{prefix}Evolving anchors with Genetic Algorithm: fitness = {f:.4f}' + pbar.desc = f"{prefix}Evolving anchors with Genetic Algorithm: fitness = {f:.4f}" if verbose: print_results(k) diff --git a/utils/autobatch.py b/utils/autobatch.py index 3f2b4d1a4c38..f8fd72924386 100644 --- a/utils/autobatch.py +++ b/utils/autobatch.py @@ -27,11 +27,11 @@ def autobatch(model, imgsz=640, fraction=0.9, batch_size=16): # model = torch.hub.load('ultralytics/yolov5', 'yolov5s', autoshape=False) # print(autobatch(model)) - prefix = colorstr('autobatch: ') - print(f'{prefix}Computing optimal batch size for --imgsz {imgsz}') + prefix = colorstr("autobatch: ") + print(f"{prefix}Computing optimal batch size for --imgsz {imgsz}") device = next(model.parameters()).device # get model device - if device.type == 'cpu': - print(f'{prefix}CUDA not detected, using default CPU batch-size {batch_size}') + if device.type == "cpu": + print(f"{prefix}CUDA not detected, using default CPU batch-size {batch_size}") return batch_size d = str(device).upper() # 'CUDA:0' @@ -40,18 +40,18 @@ def autobatch(model, imgsz=640, fraction=0.9, batch_size=16): r = torch.cuda.memory_reserved(device) / 1024 ** 3 # (GiB) a = torch.cuda.memory_allocated(device) / 1024 ** 3 # (GiB) f = t - (r + a) # free inside reserved - print(f'{prefix}{d} ({properties.name}) {t:.2f}G total, {r:.2f}G reserved, {a:.2f}G allocated, {f:.2f}G free') + print(f"{prefix}{d} ({properties.name}) {t:.2f}G total, {r:.2f}G reserved, {a:.2f}G allocated, {f:.2f}G free") batch_sizes = [1, 2, 4, 8, 16] try: img = [torch.zeros(b, 3, imgsz, imgsz) for b in batch_sizes] y = profile(img, model, n=3, device=device) except Exception as e: - print(f'{prefix}{e}') + print(f"{prefix}{e}") y = [x[2] for x in y if x] # memory [2] - batch_sizes = batch_sizes[:len(y)] + batch_sizes = batch_sizes[: len(y)] p = np.polyfit(batch_sizes, y, deg=1) # first degree polynomial fit b = int((f * fraction - p[1]) / p[0]) # y intercept (optimal batch size) - print(f'{prefix}Using batch-size {b} for {d} {t * fraction:.2f}G/{t:.2f}G ({fraction * 100:.0f}%)') + print(f"{prefix}Using batch-size {b} for {d} {t * fraction:.2f}G/{t:.2f}G ({fraction * 100:.0f}%)") return b diff --git a/utils/aws/resume.py b/utils/aws/resume.py index b21731c979a1..4525ba96749a 100644 --- a/utils/aws/resume.py +++ b/utils/aws/resume.py @@ -14,27 +14,27 @@ sys.path.append(str(ROOT)) # add ROOT to PATH port = 0 # --master_port -path = Path('').resolve() -for last in path.rglob('*/**/last.pt'): +path = Path("").resolve() +for last in path.rglob("*/**/last.pt"): ckpt = torch.load(last) - if ckpt['optimizer'] is None: + if ckpt["optimizer"] is None: continue # Load opt.yaml - with open(last.parent.parent / 'opt.yaml', errors='ignore') as f: + with open(last.parent.parent / "opt.yaml", errors="ignore") as f: opt = yaml.safe_load(f) # Get device count - d = opt['device'].split(',') # devices + d = opt["device"].split(",") # devices nd = len(d) # number of devices ddp = nd > 1 or (nd == 0 and torch.cuda.device_count() > 1) # distributed data parallel if ddp: # multi-GPU port += 1 - cmd = f'python -m torch.distributed.run --nproc_per_node {nd} --master_port {port} train.py --resume {last}' + cmd = f"python -m torch.distributed.run --nproc_per_node {nd} --master_port {port} train.py --resume {last}" else: # single-GPU - cmd = f'python train.py --resume {last}' + cmd = f"python train.py --resume {last}" - cmd += ' > /dev/null 2>&1 &' # redirect output to dev/null and run in daemon thread + cmd += " > /dev/null 2>&1 &" # redirect output to dev/null and run in daemon thread print(cmd) os.system(cmd) diff --git a/utils/callbacks.py b/utils/callbacks.py index 327b8639b60c..b3626666f461 100644 --- a/utils/callbacks.py +++ b/utils/callbacks.py @@ -5,37 +5,33 @@ class Callbacks: - """" + """ " Handles all registered callbacks for YOLOv5 Hooks """ # Define the available callbacks _callbacks = { - 'on_pretrain_routine_start': [], - 'on_pretrain_routine_end': [], - - 'on_train_start': [], - 'on_train_epoch_start': [], - 'on_train_batch_start': [], - 'optimizer_step': [], - 'on_before_zero_grad': [], - 'on_train_batch_end': [], - 'on_train_epoch_end': [], - - 'on_val_start': [], - 'on_val_batch_start': [], - 'on_val_image_end': [], - 'on_val_batch_end': [], - 'on_val_end': [], - - 'on_fit_epoch_end': [], # fit = train + val - 'on_model_save': [], - 'on_train_end': [], - - 'teardown': [], + "on_pretrain_routine_start": [], + "on_pretrain_routine_end": [], + "on_train_start": [], + "on_train_epoch_start": [], + "on_train_batch_start": [], + "optimizer_step": [], + "on_before_zero_grad": [], + "on_train_batch_end": [], + "on_train_epoch_end": [], + "on_val_start": [], + "on_val_batch_start": [], + "on_val_image_end": [], + "on_val_batch_end": [], + "on_val_end": [], + "on_fit_epoch_end": [], # fit = train + val + "on_model_save": [], + "on_train_end": [], + "teardown": [], } - def register_action(self, hook, name='', callback=None): + def register_action(self, hook, name="", callback=None): """ Register a new action to a callback hook @@ -46,10 +42,10 @@ def register_action(self, hook, name='', callback=None): """ assert hook in self._callbacks, f"hook '{hook}' not found in callbacks {self._callbacks}" assert callable(callback), f"callback '{callback}' is not callable" - self._callbacks[hook].append({'name': name, 'callback': callback}) + self._callbacks[hook].append({"name": name, "callback": callback}) def get_registered_actions(self, hook=None): - """" + """ " Returns all the registered actions by callback hook Args: @@ -73,4 +69,4 @@ def run(self, hook, *args, **kwargs): assert hook in self._callbacks, f"hook '{hook}' not found in callbacks {self._callbacks}" for logger in self._callbacks[hook]: - logger['callback'](*args, **kwargs) + logger["callback"](*args, **kwargs) diff --git a/utils/datasets.py b/utils/datasets.py index 94acaaa92cd7..65b1fe435720 100755 --- a/utils/datasets.py +++ b/utils/datasets.py @@ -26,19 +26,29 @@ from tqdm import tqdm from utils.augmentations import Albumentations, augment_hsv, copy_paste, letterbox, mixup, random_perspective -from utils.general import (LOGGER, check_dataset, check_requirements, check_yaml, clean_str, segments2boxes, xyn2xy, - xywh2xyxy, xywhn2xyxy, xyxy2xywhn) +from utils.general import ( + LOGGER, + check_dataset, + check_requirements, + check_yaml, + clean_str, + segments2boxes, + xyn2xy, + xywh2xyxy, + xywhn2xyxy, + xyxy2xywhn, +) from utils.torch_utils import torch_distributed_zero_first # Parameters -HELP_URL = 'https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data' -IMG_FORMATS = ['bmp', 'jpg', 'jpeg', 'png', 'tif', 'tiff', 'dng', 'webp', 'mpo'] # acceptable image suffixes -VID_FORMATS = ['mov', 'avi', 'mp4', 'mpg', 'mpeg', 'm4v', 'wmv', 'mkv'] # acceptable video suffixes +HELP_URL = "https://github.com/ultralytics/yolov5/wiki/Train-Custom-Data" +IMG_FORMATS = ["bmp", "jpg", "jpeg", "png", "tif", "tiff", "dng", "webp", "mpo"] # acceptable image suffixes +VID_FORMATS = ["mov", "avi", "mp4", "mpg", "mpeg", "m4v", "wmv", "mkv"] # acceptable video suffixes NUM_THREADS = min(8, os.cpu_count()) # number of multiprocessing threads # Get orientation exif tag for orientation in ExifTags.TAGS.keys(): - if ExifTags.TAGS[orientation] == 'Orientation': + if ExifTags.TAGS[orientation] == "Orientation": break @@ -46,7 +56,7 @@ def get_hash(paths): # Returns a single hash value of a list of paths (files or dirs) size = sum(os.path.getsize(p) for p in paths if os.path.exists(p)) # sizes h = hashlib.md5(str(size).encode()) # hash sizes - h.update(''.join(paths).encode()) # hash paths + h.update("".join(paths).encode()) # hash paths return h.hexdigest() # return hash @@ -76,14 +86,15 @@ def exif_transpose(image): exif = image.getexif() orientation = exif.get(0x0112, 1) # default 1 if orientation > 1: - method = {2: Image.FLIP_LEFT_RIGHT, - 3: Image.ROTATE_180, - 4: Image.FLIP_TOP_BOTTOM, - 5: Image.TRANSPOSE, - 6: Image.ROTATE_270, - 7: Image.TRANSVERSE, - 8: Image.ROTATE_90, - }.get(orientation) + method = { + 2: Image.FLIP_LEFT_RIGHT, + 3: Image.ROTATE_180, + 4: Image.FLIP_TOP_BOTTOM, + 5: Image.TRANSPOSE, + 6: Image.ROTATE_270, + 7: Image.TRANSVERSE, + 8: Image.ROTATE_90, + }.get(orientation) if method is not None: image = image.transpose(method) del exif[0x0112] @@ -91,44 +102,65 @@ def exif_transpose(image): return image -def create_dataloader(path, imgsz, batch_size, stride, single_cls=False, hyp=None, augment=False, cache=False, pad=0.0, - rect=False, rank=-1, workers=8, image_weights=False, quad=False, prefix=''): +def create_dataloader( + path, + imgsz, + batch_size, + stride, + single_cls=False, + hyp=None, + augment=False, + cache=False, + pad=0.0, + rect=False, + rank=-1, + workers=8, + image_weights=False, + quad=False, + prefix="", +): # Make sure only the first process in DDP process the dataset first, and the following others can use the cache with torch_distributed_zero_first(rank): - dataset = LoadImagesAndLabels(path, imgsz, batch_size, - augment=augment, # augment images - hyp=hyp, # augmentation hyperparameters - rect=rect, # rectangular training - cache_images=cache, - single_cls=single_cls, - stride=int(stride), - pad=pad, - image_weights=image_weights, - prefix=prefix) + dataset = LoadImagesAndLabels( + path, + imgsz, + batch_size, + augment=augment, # augment images + hyp=hyp, # augmentation hyperparameters + rect=rect, # rectangular training + cache_images=cache, + single_cls=single_cls, + stride=int(stride), + pad=pad, + image_weights=image_weights, + prefix=prefix, + ) batch_size = min(batch_size, len(dataset)) nw = min([os.cpu_count(), batch_size if batch_size > 1 else 0, workers]) # number of workers sampler = torch.utils.data.distributed.DistributedSampler(dataset) if rank != -1 else None loader = torch.utils.data.DataLoader if image_weights else InfiniteDataLoader # Use torch.utils.data.DataLoader() if dataset.properties will update during training else InfiniteDataLoader() - dataloader = loader(dataset, - batch_size=batch_size, - num_workers=nw, - sampler=sampler, - pin_memory=True, - collate_fn=LoadImagesAndLabels.collate_fn4 if quad else LoadImagesAndLabels.collate_fn) + dataloader = loader( + dataset, + batch_size=batch_size, + num_workers=nw, + sampler=sampler, + pin_memory=True, + collate_fn=LoadImagesAndLabels.collate_fn4 if quad else LoadImagesAndLabels.collate_fn, + ) return dataloader, dataset class InfiniteDataLoader(torch.utils.data.dataloader.DataLoader): - """ Dataloader that reuses workers + """Dataloader that reuses workers Uses same syntax as vanilla DataLoader """ def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) - object.__setattr__(self, 'batch_sampler', _RepeatSampler(self.batch_sampler)) + object.__setattr__(self, "batch_sampler", _RepeatSampler(self.batch_sampler)) self.iterator = super().__iter__() def __len__(self): @@ -140,7 +172,7 @@ def __iter__(self): class _RepeatSampler: - """ Sampler that repeats forever + """Sampler that repeats forever Args: sampler (Sampler) @@ -158,17 +190,17 @@ class LoadImages: # YOLOv5 image/video dataloader, i.e. `python detect.py --source image.jpg/vid.mp4` def __init__(self, path, img_size=640, stride=32, auto=True): p = str(Path(path).resolve()) # os-agnostic absolute path - if '*' in p: + if "*" in p: files = sorted(glob.glob(p, recursive=True)) # glob elif os.path.isdir(p): - files = sorted(glob.glob(os.path.join(p, '*.*'))) # dir + files = sorted(glob.glob(os.path.join(p, "*.*"))) # dir elif os.path.isfile(p): files = [p] # files else: - raise Exception(f'ERROR: {p} does not exist') + raise Exception(f"ERROR: {p} does not exist") - images = [x for x in files if x.split('.')[-1].lower() in IMG_FORMATS] - videos = [x for x in files if x.split('.')[-1].lower() in VID_FORMATS] + images = [x for x in files if x.split(".")[-1].lower() in IMG_FORMATS] + videos = [x for x in files if x.split(".")[-1].lower() in VID_FORMATS] ni, nv = len(images), len(videos) self.img_size = img_size @@ -176,14 +208,16 @@ def __init__(self, path, img_size=640, stride=32, auto=True): self.files = images + videos self.nf = ni + nv # number of files self.video_flag = [False] * ni + [True] * nv - self.mode = 'image' + self.mode = "image" self.auto = auto if any(videos): self.new_video(videos[0]) # new video else: self.cap = None - assert self.nf > 0, f'No images or videos found in {p}. ' \ - f'Supported formats are:\nimages: {IMG_FORMATS}\nvideos: {VID_FORMATS}' + assert self.nf > 0, ( + f"No images or videos found in {p}. " + f"Supported formats are:\nimages: {IMG_FORMATS}\nvideos: {VID_FORMATS}" + ) def __iter__(self): self.count = 0 @@ -196,7 +230,7 @@ def __next__(self): if self.video_flag[self.count]: # Read video - self.mode = 'video' + self.mode = "video" ret_val, img0 = self.cap.read() if not ret_val: self.count += 1 @@ -209,14 +243,14 @@ def __next__(self): ret_val, img0 = self.cap.read() self.frame += 1 - s = f'video {self.count + 1}/{self.nf} ({self.frame}/{self.frames}) {path}: ' + s = f"video {self.count + 1}/{self.nf} ({self.frame}/{self.frames}) {path}: " else: # Read image self.count += 1 img0 = cv2.imread(path) # BGR - assert img0 is not None, f'Image Not Found {path}' - s = f'image {self.count}/{self.nf} {path}: ' + assert img0 is not None, f"Image Not Found {path}" + s = f"image {self.count}/{self.nf} {path}: " # Padded resize img = letterbox(img0, self.img_size, stride=self.stride, auto=self.auto)[0] @@ -238,7 +272,7 @@ def __len__(self): class LoadWebcam: # for inference # YOLOv5 local webcam dataloader, i.e. `python detect.py --source 0` - def __init__(self, pipe='0', img_size=640, stride=32): + def __init__(self, pipe="0", img_size=640, stride=32): self.img_size = img_size self.stride = stride self.pipe = eval(pipe) if pipe.isnumeric() else pipe @@ -251,7 +285,7 @@ def __iter__(self): def __next__(self): self.count += 1 - if cv2.waitKey(1) == ord('q'): # q to quit + if cv2.waitKey(1) == ord("q"): # q to quit self.cap.release() cv2.destroyAllWindows() raise StopIteration @@ -261,9 +295,9 @@ def __next__(self): img0 = cv2.flip(img0, 1) # flip left-right # Print - assert ret_val, f'Camera Error {self.pipe}' - img_path = 'webcam.jpg' - s = f'webcam {self.count}: ' + assert ret_val, f"Camera Error {self.pipe}" + img_path = "webcam.jpg" + s = f"webcam {self.count}: " # Padded resize img = letterbox(img0, self.img_size, stride=self.stride)[0] @@ -280,8 +314,8 @@ def __len__(self): class LoadStreams: # YOLOv5 streamloader, i.e. `python detect.py --source 'rtsp://example.com/media.mp4' # RTSP, RTMP, HTTP streams` - def __init__(self, sources='streams.txt', img_size=640, stride=32, auto=True): - self.mode = 'stream' + def __init__(self, sources="streams.txt", img_size=640, stride=32, auto=True): + self.mode = "stream" self.img_size = img_size self.stride = stride @@ -297,30 +331,31 @@ def __init__(self, sources='streams.txt', img_size=640, stride=32, auto=True): self.auto = auto for i, s in enumerate(sources): # index, source # Start thread to read frames from video stream - st = f'{i + 1}/{n}: {s}... ' - if 'youtube.com/' in s or 'youtu.be/' in s: # if source is YouTube video - check_requirements(('pafy', 'youtube_dl')) + st = f"{i + 1}/{n}: {s}... " + if "youtube.com/" in s or "youtu.be/" in s: # if source is YouTube video + check_requirements(("pafy", "youtube_dl")) import pafy + s = pafy.new(s).getbest(preftype="mp4").url # YouTube URL s = eval(s) if s.isnumeric() else s # i.e. s = '0' local webcam cap = cv2.VideoCapture(s) - assert cap.isOpened(), f'{st}Failed to open {s}' + assert cap.isOpened(), f"{st}Failed to open {s}" w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) self.fps[i] = max(cap.get(cv2.CAP_PROP_FPS) % 100, 0) or 30.0 # 30 FPS fallback - self.frames[i] = max(int(cap.get(cv2.CAP_PROP_FRAME_COUNT)), 0) or float('inf') # infinite stream fallback + self.frames[i] = max(int(cap.get(cv2.CAP_PROP_FRAME_COUNT)), 0) or float("inf") # infinite stream fallback _, self.imgs[i] = cap.read() # guarantee first frame self.threads[i] = Thread(target=self.update, args=([i, cap, s]), daemon=True) LOGGER.info(f"{st} Success ({self.frames[i]} frames {w}x{h} at {self.fps[i]:.2f} FPS)") self.threads[i].start() - LOGGER.info('') # newline + LOGGER.info("") # newline # check for common shapes s = np.stack([letterbox(x, self.img_size, stride=self.stride, auto=self.auto)[0].shape for x in self.imgs]) self.rect = np.unique(s, axis=0).shape[0] == 1 # rect inference if all shapes equal if not self.rect: - LOGGER.warning('WARNING: Stream shapes differ. For optimal performance supply similarly-shaped streams.') + LOGGER.warning("WARNING: Stream shapes differ. For optimal performance supply similarly-shaped streams.") def update(self, i, cap, stream): # Read stream `i` frames in daemon thread @@ -334,7 +369,7 @@ def update(self, i, cap, stream): if success: self.imgs[i] = im else: - LOGGER.warning('WARNING: Video stream unresponsive, please check your IP camera connection.') + LOGGER.warning("WARNING: Video stream unresponsive, please check your IP camera connection.") self.imgs[i] *= 0 cap.open(stream) # re-open stream if signal was lost time.sleep(1 / self.fps[i]) # wait time @@ -345,7 +380,7 @@ def __iter__(self): def __next__(self): self.count += 1 - if not all(x.is_alive() for x in self.threads) or cv2.waitKey(1) == ord('q'): # q to quit + if not all(x.is_alive() for x in self.threads) or cv2.waitKey(1) == ord("q"): # q to quit cv2.destroyAllWindows() raise StopIteration @@ -360,7 +395,7 @@ def __next__(self): img = img[..., ::-1].transpose((0, 3, 1, 2)) # BGR to RGB, BHWC to BCHW img = np.ascontiguousarray(img) - return self.sources, img, img0, None, '' + return self.sources, img, img0, None, "" def __len__(self): return len(self.sources) # 1E12 frames = 32 streams at 30 FPS for 30 years @@ -368,16 +403,29 @@ def __len__(self): def img2label_paths(img_paths): # Define label paths as a function of image paths - sa, sb = os.sep + 'images' + os.sep, os.sep + 'labels' + os.sep # /images/, /labels/ substrings - return [sb.join(x.rsplit(sa, 1)).rsplit('.', 1)[0] + '.txt' for x in img_paths] + sa, sb = os.sep + "images" + os.sep, os.sep + "labels" + os.sep # /images/, /labels/ substrings + return [sb.join(x.rsplit(sa, 1)).rsplit(".", 1)[0] + ".txt" for x in img_paths] class LoadImagesAndLabels(Dataset): # YOLOv5 train_loader/val_loader, loads images and labels for training and validation cache_version = 0.6 # dataset labels *.cache version - def __init__(self, path, img_size=640, batch_size=16, augment=False, hyp=None, rect=False, image_weights=False, - cache_images=False, single_cls=False, stride=32, pad=0.0, prefix=''): + def __init__( + self, + path, + img_size=640, + batch_size=16, + augment=False, + hyp=None, + rect=False, + image_weights=False, + cache_images=False, + single_cls=False, + stride=32, + pad=0.0, + prefix="", + ): self.img_size = img_size self.augment = augment self.hyp = hyp @@ -394,43 +442,43 @@ def __init__(self, path, img_size=640, batch_size=16, augment=False, hyp=None, r for p in path if isinstance(path, list) else [path]: p = Path(p) # os-agnostic if p.is_dir(): # dir - f += glob.glob(str(p / '**' / '*.*'), recursive=True) + f += glob.glob(str(p / "**" / "*.*"), recursive=True) # f = list(p.rglob('*.*')) # pathlib elif p.is_file(): # file with open(p) as t: t = t.read().strip().splitlines() parent = str(p.parent) + os.sep - f += [x.replace('./', parent) if x.startswith('./') else x for x in t] # local to global path + f += [x.replace("./", parent) if x.startswith("./") else x for x in t] # local to global path # f += [p.parent / x.lstrip(os.sep) for x in t] # local to global path (pathlib) else: - raise Exception(f'{prefix}{p} does not exist') - self.img_files = sorted(x.replace('/', os.sep) for x in f if x.split('.')[-1].lower() in IMG_FORMATS) + raise Exception(f"{prefix}{p} does not exist") + self.img_files = sorted(x.replace("/", os.sep) for x in f if x.split(".")[-1].lower() in IMG_FORMATS) # self.img_files = sorted([x for x in f if x.suffix[1:].lower() in IMG_FORMATS]) # pathlib - assert self.img_files, f'{prefix}No images found' + assert self.img_files, f"{prefix}No images found" except Exception as e: - raise Exception(f'{prefix}Error loading data from {path}: {e}\nSee {HELP_URL}') + raise Exception(f"{prefix}Error loading data from {path}: {e}\nSee {HELP_URL}") # Check cache self.label_files = img2label_paths(self.img_files) # labels - cache_path = (p if p.is_file() else Path(self.label_files[0]).parent).with_suffix('.cache') + cache_path = (p if p.is_file() else Path(self.label_files[0]).parent).with_suffix(".cache") try: cache, exists = np.load(cache_path, allow_pickle=True).item(), True # load dict - assert cache['version'] == self.cache_version # same version - assert cache['hash'] == get_hash(self.label_files + self.img_files) # same hash + assert cache["version"] == self.cache_version # same version + assert cache["hash"] == get_hash(self.label_files + self.img_files) # same hash except: cache, exists = self.cache_labels(cache_path, prefix), False # cache # Display cache - nf, nm, ne, nc, n = cache.pop('results') # found, missing, empty, corrupted, total + nf, nm, ne, nc, n = cache.pop("results") # found, missing, empty, corrupted, total if exists: d = f"Scanning '{cache_path}' images and labels... {nf} found, {nm} missing, {ne} empty, {nc} corrupted" tqdm(None, desc=prefix + d, total=n, initial=n) # display cache results - if cache['msgs']: - LOGGER.info('\n'.join(cache['msgs'])) # display warnings - assert nf > 0 or not augment, f'{prefix}No labels in {cache_path}. Can not train without labels. See {HELP_URL}' + if cache["msgs"]: + LOGGER.info("\n".join(cache["msgs"])) # display warnings + assert nf > 0 or not augment, f"{prefix}No labels in {cache_path}. Can not train without labels. See {HELP_URL}" # Read cache - [cache.pop(k) for k in ('hash', 'version', 'msgs')] # remove items + [cache.pop(k) for k in ("hash", "version", "msgs")] # remove items labels, shapes, self.segments = zip(*cache.values()) self.labels = list(labels) self.shapes = np.array(shapes, dtype=np.float64) @@ -484,33 +532,36 @@ def __init__(self, path, img_size=640, batch_size=16, augment=False, hyp=None, r # Cache images into memory for faster training (WARNING: large datasets may exceed system RAM) self.imgs, self.img_npy = [None] * n, [None] * n if cache_images: - if cache_images == 'disk': - self.im_cache_dir = Path(Path(self.img_files[0]).parent.as_posix() + '_npy') - self.img_npy = [self.im_cache_dir / Path(f).with_suffix('.npy').name for f in self.img_files] + if cache_images == "disk": + self.im_cache_dir = Path(Path(self.img_files[0]).parent.as_posix() + "_npy") + self.img_npy = [self.im_cache_dir / Path(f).with_suffix(".npy").name for f in self.img_files] self.im_cache_dir.mkdir(parents=True, exist_ok=True) gb = 0 # Gigabytes of cached images self.img_hw0, self.img_hw = [None] * n, [None] * n results = ThreadPool(NUM_THREADS).imap(lambda x: load_image(*x), zip(repeat(self), range(n))) pbar = tqdm(enumerate(results), total=n) for i, x in pbar: - if cache_images == 'disk': + if cache_images == "disk": if not self.img_npy[i].exists(): np.save(self.img_npy[i].as_posix(), x[0]) gb += self.img_npy[i].stat().st_size else: self.imgs[i], self.img_hw0[i], self.img_hw[i] = x # im, hw_orig, hw_resized = load_image(self, i) gb += self.imgs[i].nbytes - pbar.desc = f'{prefix}Caching images ({gb / 1E9:.1f}GB {cache_images})' + pbar.desc = f"{prefix}Caching images ({gb / 1E9:.1f}GB {cache_images})" pbar.close() - def cache_labels(self, path=Path('./labels.cache'), prefix=''): + def cache_labels(self, path=Path("./labels.cache"), prefix=""): # Cache dataset labels, check images and read shapes x = {} # dict nm, nf, ne, nc, msgs = 0, 0, 0, 0, [] # number missing, found, empty, corrupt, messages desc = f"{prefix}Scanning '{path.parent / path.stem}' images and labels..." with Pool(NUM_THREADS) as pool: - pbar = tqdm(pool.imap(verify_image_label, zip(self.img_files, self.label_files, repeat(prefix))), - desc=desc, total=len(self.img_files)) + pbar = tqdm( + pool.imap(verify_image_label, zip(self.img_files, self.label_files, repeat(prefix))), + desc=desc, + total=len(self.img_files), + ) for im_file, l, shape, segments, nm_f, nf_f, ne_f, nc_f, msg in pbar: nm += nm_f nf += nf_f @@ -524,19 +575,19 @@ def cache_labels(self, path=Path('./labels.cache'), prefix=''): pbar.close() if msgs: - LOGGER.info('\n'.join(msgs)) + LOGGER.info("\n".join(msgs)) if nf == 0: - LOGGER.warning(f'{prefix}WARNING: No labels found in {path}. See {HELP_URL}') - x['hash'] = get_hash(self.label_files + self.img_files) - x['results'] = nf, nm, ne, nc, len(self.img_files) - x['msgs'] = msgs # warnings - x['version'] = self.cache_version # cache version + LOGGER.warning(f"{prefix}WARNING: No labels found in {path}. See {HELP_URL}") + x["hash"] = get_hash(self.label_files + self.img_files) + x["results"] = nf, nm, ne, nc, len(self.img_files) + x["msgs"] = msgs # warnings + x["version"] = self.cache_version # cache version try: np.save(path, x) # save cache for next time - path.with_suffix('.cache.npy').rename(path) # remove .npy suffix - LOGGER.info(f'{prefix}New cache created: {path}') + path.with_suffix(".cache.npy").rename(path) # remove .npy suffix + LOGGER.info(f"{prefix}New cache created: {path}") except Exception as e: - LOGGER.warning(f'{prefix}WARNING: Cache directory {path.parent} is not writeable: {e}') # not writeable + LOGGER.warning(f"{prefix}WARNING: Cache directory {path.parent} is not writeable: {e}") # not writeable return x def __len__(self): @@ -552,14 +603,14 @@ def __getitem__(self, index): index = self.indices[index] # linear, shuffled, or image_weights hyp = self.hyp - mosaic = self.mosaic and random.random() < hyp['mosaic'] + mosaic = self.mosaic and random.random() < hyp["mosaic"] if mosaic: # Load mosaic img, labels = load_mosaic(self, index) shapes = None # MixUp augmentation - if random.random() < hyp['mixup']: + if random.random() < hyp["mixup"]: img, labels = mixup(img, labels, *load_mosaic(self, random.randint(0, self.n - 1))) else: @@ -576,16 +627,19 @@ def __getitem__(self, index): labels[:, 1:] = xywhn2xyxy(labels[:, 1:], ratio[0] * w, ratio[1] * h, padw=pad[0], padh=pad[1]) if self.augment: - img, labels = random_perspective(img, labels, - degrees=hyp['degrees'], - translate=hyp['translate'], - scale=hyp['scale'], - shear=hyp['shear'], - perspective=hyp['perspective']) + img, labels = random_perspective( + img, + labels, + degrees=hyp["degrees"], + translate=hyp["translate"], + scale=hyp["scale"], + shear=hyp["shear"], + perspective=hyp["perspective"], + ) nl = len(labels) # number of labels if nl: - labels[:, 1:5] = xyxy2xywhn(labels[:, 1:5], w=img.shape[1], h=img.shape[0], clip=True, eps=1E-3) + labels[:, 1:5] = xyxy2xywhn(labels[:, 1:5], w=img.shape[1], h=img.shape[0], clip=True, eps=1e-3) if self.augment: # Albumentations @@ -593,16 +647,16 @@ def __getitem__(self, index): nl = len(labels) # update after albumentations # HSV color-space - augment_hsv(img, hgain=hyp['hsv_h'], sgain=hyp['hsv_s'], vgain=hyp['hsv_v']) + augment_hsv(img, hgain=hyp["hsv_h"], sgain=hyp["hsv_s"], vgain=hyp["hsv_v"]) # Flip up-down - if random.random() < hyp['flipud']: + if random.random() < hyp["flipud"]: img = np.flipud(img) if nl: labels[:, 2] = 1 - labels[:, 2] # Flip left-right - if random.random() < hyp['fliplr']: + if random.random() < hyp["fliplr"]: img = np.fliplr(img) if nl: labels[:, 1] = 1 - labels[:, 1] @@ -639,8 +693,9 @@ def collate_fn4(batch): for i in range(n): # zidane torch.zeros(16,3,720,1280) # BCHW i *= 4 if random.random() < 0.5: - im = F.interpolate(img[i].unsqueeze(0).float(), scale_factor=2.0, mode='bilinear', align_corners=False)[ - 0].type(img[i].type()) + im = F.interpolate(img[i].unsqueeze(0).float(), scale_factor=2.0, mode="bilinear", align_corners=False)[ + 0 + ].type(img[i].type()) l = label[i] else: im = torch.cat((torch.cat((img[i], img[i + 1]), 1), torch.cat((img[i + 2], img[i + 3]), 1)), 2) @@ -665,12 +720,15 @@ def load_image(self, i): else: # read image path = self.img_files[i] im = cv2.imread(path) # BGR - assert im is not None, f'Image Not Found {path}' + assert im is not None, f"Image Not Found {path}" h0, w0 = im.shape[:2] # orig hw r = self.img_size / max(h0, w0) # ratio if r != 1: # if sizes are not equal - im = cv2.resize(im, (int(w0 * r), int(h0 * r)), - interpolation=cv2.INTER_AREA if r < 1 and not self.augment else cv2.INTER_LINEAR) + im = cv2.resize( + im, + (int(w0 * r), int(h0 * r)), + interpolation=cv2.INTER_AREA if r < 1 and not self.augment else cv2.INTER_LINEAR, + ) return im, (h0, w0), im.shape[:2] # im, hw_original, hw_resized else: return self.imgs[i], self.img_hw0[i], self.img_hw[i] # im, hw_original, hw_resized @@ -721,14 +779,18 @@ def load_mosaic(self, index): # img4, labels4 = replicate(img4, labels4) # replicate # Augment - img4, labels4, segments4 = copy_paste(img4, labels4, segments4, p=self.hyp['copy_paste']) - img4, labels4 = random_perspective(img4, labels4, segments4, - degrees=self.hyp['degrees'], - translate=self.hyp['translate'], - scale=self.hyp['scale'], - shear=self.hyp['shear'], - perspective=self.hyp['perspective'], - border=self.mosaic_border) # border to remove + img4, labels4, segments4 = copy_paste(img4, labels4, segments4, p=self.hyp["copy_paste"]) + img4, labels4 = random_perspective( + img4, + labels4, + segments4, + degrees=self.hyp["degrees"], + translate=self.hyp["translate"], + scale=self.hyp["scale"], + shear=self.hyp["shear"], + perspective=self.hyp["perspective"], + border=self.mosaic_border, + ) # border to remove return img4, labels4 @@ -777,12 +839,12 @@ def load_mosaic9(self, index): segments9.extend(segments) # Image - img9[y1:y2, x1:x2] = img[y1 - pady:, x1 - padx:] # img9[ymin:ymax, xmin:xmax] + img9[y1:y2, x1:x2] = img[y1 - pady :, x1 - padx :] # img9[ymin:ymax, xmin:xmax] hp, wp = h, w # height, width previous # Offset yc, xc = (int(random.uniform(0, s)) for _ in self.mosaic_border) # mosaic center x, y - img9 = img9[yc:yc + 2 * s, xc:xc + 2 * s] + img9 = img9[yc : yc + 2 * s, xc : xc + 2 * s] # Concat/clip labels labels9 = np.concatenate(labels9, 0) @@ -796,37 +858,41 @@ def load_mosaic9(self, index): # img9, labels9 = replicate(img9, labels9) # replicate # Augment - img9, labels9 = random_perspective(img9, labels9, segments9, - degrees=self.hyp['degrees'], - translate=self.hyp['translate'], - scale=self.hyp['scale'], - shear=self.hyp['shear'], - perspective=self.hyp['perspective'], - border=self.mosaic_border) # border to remove + img9, labels9 = random_perspective( + img9, + labels9, + segments9, + degrees=self.hyp["degrees"], + translate=self.hyp["translate"], + scale=self.hyp["scale"], + shear=self.hyp["shear"], + perspective=self.hyp["perspective"], + border=self.mosaic_border, + ) # border to remove return img9, labels9 -def create_folder(path='./new'): +def create_folder(path="./new"): # Create folder if os.path.exists(path): shutil.rmtree(path) # delete output folder os.makedirs(path) # make new output folder -def flatten_recursive(path='../datasets/coco128'): +def flatten_recursive(path="../datasets/coco128"): # Flatten a recursive directory by bringing all files to top level - new_path = Path(path + '_flat') + new_path = Path(path + "_flat") create_folder(new_path) - for file in tqdm(glob.glob(str(Path(path)) + '/**/*.*', recursive=True)): + for file in tqdm(glob.glob(str(Path(path)) + "/**/*.*", recursive=True)): shutil.copyfile(file, new_path / Path(file).name) -def extract_boxes(path='../datasets/coco128'): # from utils.datasets import *; extract_boxes() +def extract_boxes(path="../datasets/coco128"): # from utils.datasets import *; extract_boxes() # Convert detection dataset into classification dataset, with one directory per class path = Path(path) # images dir - shutil.rmtree(path / 'classifier') if (path / 'classifier').is_dir() else None # remove existing - files = list(path.rglob('*.*')) + shutil.rmtree(path / "classifier") if (path / "classifier").is_dir() else None # remove existing + files = list(path.rglob("*.*")) n = len(files) # number of files for im_file in tqdm(files, total=n): if im_file.suffix[1:] in IMG_FORMATS: @@ -842,7 +908,7 @@ def extract_boxes(path='../datasets/coco128'): # from utils.datasets import *; for j, x in enumerate(lb): c = int(x[0]) # class - f = (path / 'classifier') / f'{c}' / f'{path.stem}_{im_file.stem}_{j}.jpg' # new filename + f = (path / "classifier") / f"{c}" / f"{path.stem}_{im_file.stem}_{j}.jpg" # new filename if not f.parent.is_dir(): f.parent.mkdir(parents=True) @@ -853,11 +919,11 @@ def extract_boxes(path='../datasets/coco128'): # from utils.datasets import *; b[[0, 2]] = np.clip(b[[0, 2]], 0, w) # clip boxes outside of image b[[1, 3]] = np.clip(b[[1, 3]], 0, h) - assert cv2.imwrite(str(f), im[b[1]:b[3], b[0]:b[2]]), f'box failure in {f}' + assert cv2.imwrite(str(f), im[b[1] : b[3], b[0] : b[2]]), f"box failure in {f}" -def autosplit(path='../datasets/coco128/images', weights=(0.9, 0.1, 0.0), annotated_only=False): - """ Autosplit a dataset into train/val/test splits and save path/autosplit_*.txt files +def autosplit(path="../datasets/coco128/images", weights=(0.9, 0.1, 0.0), annotated_only=False): + """Autosplit a dataset into train/val/test splits and save path/autosplit_*.txt files Usage: from utils.datasets import *; autosplit() Arguments path: Path to images directory @@ -865,38 +931,38 @@ def autosplit(path='../datasets/coco128/images', weights=(0.9, 0.1, 0.0), annota annotated_only: Only use images with an annotated txt file """ path = Path(path) # images dir - files = sorted(x for x in path.rglob('*.*') if x.suffix[1:].lower() in IMG_FORMATS) # image files only + files = sorted(x for x in path.rglob("*.*") if x.suffix[1:].lower() in IMG_FORMATS) # image files only n = len(files) # number of files random.seed(0) # for reproducibility indices = random.choices([0, 1, 2], weights=weights, k=n) # assign each image to a split - txt = ['autosplit_train.txt', 'autosplit_val.txt', 'autosplit_test.txt'] # 3 txt files + txt = ["autosplit_train.txt", "autosplit_val.txt", "autosplit_test.txt"] # 3 txt files [(path.parent / x).unlink(missing_ok=True) for x in txt] # remove existing - print(f'Autosplitting images from {path}' + ', using *.txt labeled images only' * annotated_only) + print(f"Autosplitting images from {path}" + ", using *.txt labeled images only" * annotated_only) for i, img in tqdm(zip(indices, files), total=n): if not annotated_only or Path(img2label_paths([str(img)])[0]).exists(): # check label - with open(path.parent / txt[i], 'a') as f: - f.write('./' + img.relative_to(path.parent).as_posix() + '\n') # add image to txt file + with open(path.parent / txt[i], "a") as f: + f.write("./" + img.relative_to(path.parent).as_posix() + "\n") # add image to txt file def verify_image_label(args): # Verify one image-label pair im_file, lb_file, prefix = args - nm, nf, ne, nc, msg, segments = 0, 0, 0, 0, '', [] # number (missing, found, empty, corrupt), message, segments + nm, nf, ne, nc, msg, segments = 0, 0, 0, 0, "", [] # number (missing, found, empty, corrupt), message, segments try: # verify images im = Image.open(im_file) im.verify() # PIL verify shape = exif_size(im) # image size - assert (shape[0] > 9) & (shape[1] > 9), f'image size {shape} <10 pixels' - assert im.format.lower() in IMG_FORMATS, f'invalid image format {im.format}' - if im.format.lower() in ('jpg', 'jpeg'): - with open(im_file, 'rb') as f: + assert (shape[0] > 9) & (shape[1] > 9), f"image size {shape} <10 pixels" + assert im.format.lower() in IMG_FORMATS, f"invalid image format {im.format}" + if im.format.lower() in ("jpg", "jpeg"): + with open(im_file, "rb") as f: f.seek(-2, 2) - if f.read() != b'\xff\xd9': # corrupt JPEG - ImageOps.exif_transpose(Image.open(im_file)).save(im_file, 'JPEG', subsampling=0, quality=100) - msg = f'{prefix}WARNING: {im_file}: corrupt JPEG restored and saved' + if f.read() != b"\xff\xd9": # corrupt JPEG + ImageOps.exif_transpose(Image.open(im_file)).save(im_file, "JPEG", subsampling=0, quality=100) + msg = f"{prefix}WARNING: {im_file}: corrupt JPEG restored and saved" # verify labels if os.path.isfile(lb_file): @@ -910,13 +976,13 @@ def verify_image_label(args): l = np.array(l, dtype=np.float32) nl = len(l) if nl: - assert l.shape[1] == 5, f'labels require 5 columns, {l.shape[1]} columns detected' - assert (l >= 0).all(), f'negative label values {l[l < 0]}' - assert (l[:, 1:] <= 1).all(), f'non-normalized or out of bounds coordinates {l[:, 1:][l[:, 1:] > 1]}' + assert l.shape[1] == 5, f"labels require 5 columns, {l.shape[1]} columns detected" + assert (l >= 0).all(), f"negative label values {l[l < 0]}" + assert (l[:, 1:] <= 1).all(), f"non-normalized or out of bounds coordinates {l[:, 1:][l[:, 1:] > 1]}" l = np.unique(l, axis=0) # remove duplicate rows if len(l) < nl: segments = np.unique(segments, axis=0) - msg = f'{prefix}WARNING: {im_file}: {nl - len(l)} duplicate labels removed' + msg = f"{prefix}WARNING: {im_file}: {nl - len(l)} duplicate labels removed" else: ne = 1 # label empty l = np.zeros((0, 5), dtype=np.float32) @@ -926,12 +992,12 @@ def verify_image_label(args): return im_file, l, shape, segments, nm, nf, ne, nc, msg except Exception as e: nc = 1 - msg = f'{prefix}WARNING: {im_file}: ignoring corrupt image/label: {e}' + msg = f"{prefix}WARNING: {im_file}: ignoring corrupt image/label: {e}" return [None, None, None, None, nm, nf, ne, nc, msg] -def dataset_stats(path='coco128.yaml', autodownload=False, verbose=False, profile=False, hub=False): - """ Return dataset statistics dictionary with images and instances counts per split per class +def dataset_stats(path="coco128.yaml", autodownload=False, verbose=False, profile=False, hub=False): + """Return dataset statistics dictionary with images and instances counts per split per class To run in parent directory: export PYTHONPATH="$PWD/yolov5" Usage1: from utils.datasets import *; dataset_stats('coco128.yaml', autodownload=True) Usage2: from utils.datasets import *; dataset_stats('../datasets/coco128_with_yaml.zip') @@ -947,11 +1013,11 @@ def round_labels(labels): def unzip(path): # Unzip data.zip TODO: CONSTRAINT: path/to/abc.zip MUST unzip to 'path/to/abc/' - if str(path).endswith('.zip'): # path is data.zip - assert Path(path).is_file(), f'Error unzipping {path}, file not found' + if str(path).endswith(".zip"): # path is data.zip + assert Path(path).is_file(), f"Error unzipping {path}, file not found" ZipFile(path).extractall(path=path.parent) # unzip - dir = path.with_suffix('') # dataset directory == zip name - return True, str(dir), next(dir.rglob('*.yaml')) # zipped, data_dir, yaml_path + dir = path.with_suffix("") # dataset directory == zip name + return True, str(dir), next(dir.rglob("*.yaml")) # zipped, data_dir, yaml_path else: # path is data.yaml return False, None, path @@ -965,7 +1031,7 @@ def hub_ops(f, max_dim=1920): im = im.resize((int(im.width * r), int(im.height * r))) im.save(f_new, quality=75) # save except Exception as e: # use OpenCV - print(f'WARNING: HUB ops PIL failure {f}: {e}') + print(f"WARNING: HUB ops PIL failure {f}: {e}") im = cv2.imread(f) im_height, im_width = im.shape[:2] r = max_dim / max(im_height, im_width) # ratio @@ -974,58 +1040,64 @@ def hub_ops(f, max_dim=1920): cv2.imwrite(str(f_new), im) zipped, data_dir, yaml_path = unzip(Path(path)) - with open(check_yaml(yaml_path), errors='ignore') as f: + with open(check_yaml(yaml_path), errors="ignore") as f: data = yaml.safe_load(f) # data dict if zipped: - data['path'] = data_dir # TODO: should this be dir.resolve()? + data["path"] = data_dir # TODO: should this be dir.resolve()? check_dataset(data, autodownload) # download dataset if missing - hub_dir = Path(data['path'] + ('-hub' if hub else '')) - stats = {'nc': data['nc'], 'names': data['names']} # statistics dictionary - for split in 'train', 'val', 'test': + hub_dir = Path(data["path"] + ("-hub" if hub else "")) + stats = {"nc": data["nc"], "names": data["names"]} # statistics dictionary + for split in "train", "val", "test": if data.get(split) is None: stats[split] = None # i.e. no test set continue x = [] dataset = LoadImagesAndLabels(data[split]) # load dataset - for label in tqdm(dataset.labels, total=dataset.n, desc='Statistics'): - x.append(np.bincount(label[:, 0].astype(int), minlength=data['nc'])) + for label in tqdm(dataset.labels, total=dataset.n, desc="Statistics"): + x.append(np.bincount(label[:, 0].astype(int), minlength=data["nc"])) x = np.array(x) # shape(128x80) - stats[split] = {'instance_stats': {'total': int(x.sum()), 'per_class': x.sum(0).tolist()}, - 'image_stats': {'total': dataset.n, 'unlabelled': int(np.all(x == 0, 1).sum()), - 'per_class': (x > 0).sum(0).tolist()}, - 'labels': [{str(Path(k).name): round_labels(v.tolist())} for k, v in - zip(dataset.img_files, dataset.labels)]} + stats[split] = { + "instance_stats": {"total": int(x.sum()), "per_class": x.sum(0).tolist()}, + "image_stats": { + "total": dataset.n, + "unlabelled": int(np.all(x == 0, 1).sum()), + "per_class": (x > 0).sum(0).tolist(), + }, + "labels": [ + {str(Path(k).name): round_labels(v.tolist())} for k, v in zip(dataset.img_files, dataset.labels) + ], + } if hub: - im_dir = hub_dir / 'images' + im_dir = hub_dir / "images" im_dir.mkdir(parents=True, exist_ok=True) - for _ in tqdm(ThreadPool(NUM_THREADS).imap(hub_ops, dataset.img_files), total=dataset.n, desc='HUB Ops'): + for _ in tqdm(ThreadPool(NUM_THREADS).imap(hub_ops, dataset.img_files), total=dataset.n, desc="HUB Ops"): pass # Profile - stats_path = hub_dir / 'stats.json' + stats_path = hub_dir / "stats.json" if profile: for _ in range(1): - file = stats_path.with_suffix('.npy') + file = stats_path.with_suffix(".npy") t1 = time.time() np.save(file, stats) t2 = time.time() x = np.load(file, allow_pickle=True) - print(f'stats.npy times: {time.time() - t2:.3f}s read, {t2 - t1:.3f}s write') + print(f"stats.npy times: {time.time() - t2:.3f}s read, {t2 - t1:.3f}s write") - file = stats_path.with_suffix('.json') + file = stats_path.with_suffix(".json") t1 = time.time() - with open(file, 'w') as f: + with open(file, "w") as f: json.dump(stats, f) # save stats *.json t2 = time.time() with open(file) as f: x = json.load(f) # load hyps dict - print(f'stats.json times: {time.time() - t2:.3f}s read, {t2 - t1:.3f}s write') + print(f"stats.json times: {time.time() - t2:.3f}s read, {t2 - t1:.3f}s write") # Save, print and return if hub: - print(f'Saving {stats_path.resolve()}...') - with open(stats_path, 'w') as f: + print(f"Saving {stats_path.resolve()}...") + with open(stats_path, "w") as f: json.dump(stats, f) # save stats.json if verbose: print(json.dumps(stats, indent=2, sort_keys=False)) diff --git a/utils/downloads.py b/utils/downloads.py index 998a7a582a33..afa9efe28d6d 100644 --- a/utils/downloads.py +++ b/utils/downloads.py @@ -15,81 +15,93 @@ import torch -def gsutil_getsize(url=''): +def gsutil_getsize(url=""): # gs://bucket/file size https://cloud.google.com/storage/docs/gsutil/commands/du - s = subprocess.check_output(f'gsutil du {url}', shell=True).decode('utf-8') - return eval(s.split(' ')[0]) if len(s) else 0 # bytes + s = subprocess.check_output(f"gsutil du {url}", shell=True).decode("utf-8") + return eval(s.split(" ")[0]) if len(s) else 0 # bytes -def safe_download(file, url, url2=None, min_bytes=1E0, error_msg=''): +def safe_download(file, url, url2=None, min_bytes=1e0, error_msg=""): # Attempts to download file from url or url2, checks and removes incomplete downloads < min_bytes file = Path(file) assert_msg = f"Downloaded file '{file}' does not exist or size is < min_bytes={min_bytes}" try: # url1 - print(f'Downloading {url} to {file}...') + print(f"Downloading {url} to {file}...") torch.hub.download_url_to_file(url, str(file)) assert file.exists() and file.stat().st_size > min_bytes, assert_msg # check except Exception as e: # url2 file.unlink(missing_ok=True) # remove partial downloads - print(f'ERROR: {e}\nRe-attempting {url2 or url} to {file}...') + print(f"ERROR: {e}\nRe-attempting {url2 or url} to {file}...") os.system(f"curl -L '{url2 or url}' -o '{file}' --retry 3 -C -") # curl download, retry and resume on fail finally: if not file.exists() or file.stat().st_size < min_bytes: # check file.unlink(missing_ok=True) # remove partial downloads print(f"ERROR: {assert_msg}\n{error_msg}") - print('') + print("") -def attempt_download(file, repo='ultralytics/yolov5'): # from utils.downloads import *; attempt_download() +def attempt_download(file, repo="ultralytics/yolov5"): # from utils.downloads import *; attempt_download() # Attempt file download if does not exist - file = Path(str(file).strip().replace("'", '')) + file = Path(str(file).strip().replace("'", "")) if not file.exists(): # URL specified name = Path(urllib.parse.unquote(str(file))).name # decode '%2F' to '/' etc. - if str(file).startswith(('http:/', 'https:/')): # download - url = str(file).replace(':/', '://') # Pathlib turns :// -> :/ - name = name.split('?')[0] # parse authentication https://url.com/file.txt?auth... - safe_download(file=name, url=url, min_bytes=1E5) + if str(file).startswith(("http:/", "https:/")): # download + url = str(file).replace(":/", "://") # Pathlib turns :// -> :/ + name = name.split("?")[0] # parse authentication https://url.com/file.txt?auth... + safe_download(file=name, url=url, min_bytes=1e5) return name # GitHub assets file.parent.mkdir(parents=True, exist_ok=True) # make parent dir (if required) try: - response = requests.get(f'https://api.github.com/repos/{repo}/releases/latest').json() # github api - assets = [x['name'] for x in response['assets']] # release assets, i.e. ['yolov5s.pt', 'yolov5m.pt', ...] - tag = response['tag_name'] # i.e. 'v1.0' + response = requests.get(f"https://api.github.com/repos/{repo}/releases/latest").json() # github api + assets = [x["name"] for x in response["assets"]] # release assets, i.e. ['yolov5s.pt', 'yolov5m.pt', ...] + tag = response["tag_name"] # i.e. 'v1.0' except: # fallback plan - assets = ['yolov5n.pt', 'yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt', - 'yolov5n6.pt', 'yolov5s6.pt', 'yolov5m6.pt', 'yolov5l6.pt', 'yolov5x6.pt'] + assets = [ + "yolov5n.pt", + "yolov5s.pt", + "yolov5m.pt", + "yolov5l.pt", + "yolov5x.pt", + "yolov5n6.pt", + "yolov5s6.pt", + "yolov5m6.pt", + "yolov5l6.pt", + "yolov5x6.pt", + ] try: - tag = subprocess.check_output('git tag', shell=True, stderr=subprocess.STDOUT).decode().split()[-1] + tag = subprocess.check_output("git tag", shell=True, stderr=subprocess.STDOUT).decode().split()[-1] except: - tag = 'v6.0' # current release + tag = "v6.0" # current release if name in assets: - safe_download(file, - url=f'https://github.com/{repo}/releases/download/{tag}/{name}', - # url2=f'https://storage.googleapis.com/{repo}/ckpt/{name}', # backup url (optional) - min_bytes=1E5, - error_msg=f'{file} missing, try downloading from https://github.com/{repo}/releases/') + safe_download( + file, + url=f"https://github.com/{repo}/releases/download/{tag}/{name}", + # url2=f'https://storage.googleapis.com/{repo}/ckpt/{name}', # backup url (optional) + min_bytes=1e5, + error_msg=f"{file} missing, try downloading from https://github.com/{repo}/releases/", + ) return str(file) -def gdrive_download(id='16TiPfZj7htmTyhntwcZyEEAejOUxuT6m', file='tmp.zip'): +def gdrive_download(id="16TiPfZj7htmTyhntwcZyEEAejOUxuT6m", file="tmp.zip"): # Downloads a file from Google Drive. from yolov5.utils.downloads import *; gdrive_download() t = time.time() file = Path(file) - cookie = Path('cookie') # gdrive cookie - print(f'Downloading https://drive.google.com/uc?export=download&id={id} as {file}... ', end='') + cookie = Path("cookie") # gdrive cookie + print(f"Downloading https://drive.google.com/uc?export=download&id={id} as {file}... ", end="") file.unlink(missing_ok=True) # remove existing file cookie.unlink(missing_ok=True) # remove existing cookie # Attempt file download out = "NUL" if platform.system() == "Windows" else "/dev/null" os.system(f'curl -c ./cookie -s -L "drive.google.com/uc?export=download&id={id}" > {out}') - if os.path.exists('cookie'): # large file + if os.path.exists("cookie"): # large file s = f'curl -Lb ./cookie "drive.google.com/uc?export=download&confirm={get_token()}&id={id}" -o {file}' else: # small file s = f'curl -s -L -o {file} "drive.google.com/uc?export=download&id={id}"' @@ -99,16 +111,16 @@ def gdrive_download(id='16TiPfZj7htmTyhntwcZyEEAejOUxuT6m', file='tmp.zip'): # Error check if r != 0: file.unlink(missing_ok=True) # remove partial - print('Download error ') # raise Exception('Download error') + print("Download error ") # raise Exception('Download error') return r # Unzip if archive - if file.suffix == '.zip': - print('unzipping... ', end='') + if file.suffix == ".zip": + print("unzipping... ", end="") ZipFile(file).extractall(path=file.parent) # unzip file.unlink() # remove zip - print(f'Done ({time.time() - t:.1f}s)') + print(f"Done ({time.time() - t:.1f}s)") return r @@ -119,6 +131,7 @@ def get_token(cookie="./cookie"): return line.split()[-1] return "" + # Google utils: https://cloud.google.com/storage/docs/reference/libraries ---------------------------------------------- # # diff --git a/utils/general.py b/utils/general.py index b0ea1527129a..b2ff214d9028 100755 --- a/utils/general.py +++ b/utils/general.py @@ -32,11 +32,11 @@ from utils.metrics import box_iou, fitness # Settings -torch.set_printoptions(linewidth=320, precision=5, profile='long') -np.set_printoptions(linewidth=320, formatter={'float_kind': '{:11.5g}'.format}) # format short g, %precision=5 +torch.set_printoptions(linewidth=320, precision=5, profile="long") +np.set_printoptions(linewidth=320, formatter={"float_kind": "{:11.5g}".format}) # format short g, %precision=5 pd.options.display.max_columns = 10 cv2.setNumThreads(0) # prevent OpenCV from multithreading (incompatible with PyTorch DataLoader) -os.environ['NUMEXPR_MAX_THREADS'] = str(min(os.cpu_count(), 8)) # NumExpr max threads +os.environ["NUMEXPR_MAX_THREADS"] = str(min(os.cpu_count(), 8)) # NumExpr max threads FILE = Path(__file__).resolve() ROOT = FILE.parents[1] # YOLOv5 root directory @@ -44,7 +44,7 @@ def set_logging(name=None, verbose=True): # Sets level and returns logger - rank = int(os.getenv('RANK', -1)) # rank in world for Multi-GPU trainings + rank = int(os.getenv("RANK", -1)) # rank in world for Multi-GPU trainings logging.basicConfig(format="%(message)s", level=logging.INFO if (verbose and rank in (-1, 0)) else logging.WARNING) return logging.getLogger(name) @@ -58,12 +58,12 @@ def __enter__(self): self.start = time.time() def __exit__(self, type, value, traceback): - print(f'Profile results: {time.time() - self.start:.5f}s') + print(f"Profile results: {time.time() - self.start:.5f}s") class Timeout(contextlib.ContextDecorator): # Usage: @Timeout(seconds) decorator or 'with Timeout(seconds):' context manager - def __init__(self, seconds, *, timeout_msg='', suppress_timeout_errors=True): + def __init__(self, seconds, *, timeout_msg="", suppress_timeout_errors=True): self.seconds = int(seconds) self.timeout_message = timeout_msg self.suppress = bool(suppress_timeout_errors) @@ -112,13 +112,14 @@ def methods(instance): def print_args(name, opt): # Print argparser arguments - LOGGER.info(colorstr(f'{name}: ') + ', '.join(f'{k}={v}' for k, v in vars(opt).items())) + LOGGER.info(colorstr(f"{name}: ") + ", ".join(f"{k}={v}" for k, v in vars(opt).items())) def init_seeds(seed=0): # Initialize random number generator (RNG) seeds https://pytorch.org/docs/stable/notes/randomness.html # cudnn seed 0 settings are slower and more reproducible, else faster and less reproducible import torch.backends.cudnn as cudnn + random.seed(seed) np.random.seed(seed) torch.manual_seed(seed) @@ -130,21 +131,21 @@ def intersect_dicts(da, db, exclude=()): return {k: v for k, v in da.items() if k in db and not any(x in k for x in exclude) and v.shape == db[k].shape} -def get_latest_run(search_dir='.'): +def get_latest_run(search_dir="."): # Return path to most recent 'last.pt' in /runs (i.e. to --resume from) - last_list = glob.glob(f'{search_dir}/**/last*.pt', recursive=True) - return max(last_list, key=os.path.getctime) if last_list else '' + last_list = glob.glob(f"{search_dir}/**/last*.pt", recursive=True) + return max(last_list, key=os.path.getctime) if last_list else "" -def user_config_dir(dir='Ultralytics', env_var='YOLOV5_CONFIG_DIR'): +def user_config_dir(dir="Ultralytics", env_var="YOLOV5_CONFIG_DIR"): # Return path of user configuration directory. Prefer environment variable if exists. Make dir if required. env = os.getenv(env_var) if env: path = Path(env) # use environment variable else: - cfg = {'Windows': 'AppData/Roaming', 'Linux': '.config', 'Darwin': 'Library/Application Support'} # 3 OS dirs - path = Path.home() / cfg.get(platform.system(), '') # OS-specific config dir - path = (path if is_writeable(path) else Path('/tmp')) / dir # GCP and AWS lambda fix, only /tmp is writeable + cfg = {"Windows": "AppData/Roaming", "Linux": ".config", "Darwin": "Library/Application Support"} # 3 OS dirs + path = Path.home() / cfg.get(platform.system(), "") # OS-specific config dir + path = (path if is_writeable(path) else Path("/tmp")) / dir # GCP and AWS lambda fix, only /tmp is writeable path.mkdir(exist_ok=True) # make if required return path @@ -152,9 +153,9 @@ def user_config_dir(dir='Ultralytics', env_var='YOLOV5_CONFIG_DIR'): def is_writeable(dir, test=False): # Return True if directory has write permissions, test opening a file with write permissions if test=True if test: # method 1 - file = Path(dir) / 'tmp.txt' + file = Path(dir) / "tmp.txt" try: - with open(file, 'w'): # open file with write permissions + with open(file, "w"): # open file with write permissions pass file.unlink() # remove file return True @@ -166,13 +167,14 @@ def is_writeable(dir, test=False): def is_docker(): # Is environment a Docker container? - return Path('/workspace').exists() # or Path('/.dockerenv').exists() + return Path("/workspace").exists() # or Path('/.dockerenv').exists() def is_colab(): # Is environment a Google Colab instance? try: import google.colab + return True except ImportError: return False @@ -180,32 +182,32 @@ def is_colab(): def is_pip(): # Is file in a pip package? - return 'site-packages' in Path(__file__).resolve().parts + return "site-packages" in Path(__file__).resolve().parts -def is_ascii(s=''): +def is_ascii(s=""): # Is string composed of all ASCII (no UTF) characters? (note str().isascii() introduced in python 3.7) s = str(s) # convert list, tuple, None, etc. to str - return len(s.encode().decode('ascii', 'ignore')) == len(s) + return len(s.encode().decode("ascii", "ignore")) == len(s) -def is_chinese(s='人工智能'): +def is_chinese(s="人工智能"): # Is string composed of any Chinese characters? - return re.search('[\u4e00-\u9fff]', s) + return re.search("[\u4e00-\u9fff]", s) -def emojis(str=''): +def emojis(str=""): # Return platform-dependent emoji-safe version of string - return str.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else str + return str.encode().decode("ascii", "ignore") if platform.system() == "Windows" else str def file_size(path): # Return file/dir size (MB) path = Path(path) if path.is_file(): - return path.stat().st_size / 1E6 + return path.stat().st_size / 1e6 elif path.is_dir(): - return sum(f.stat().st_size for f in path.glob('**/*') if f.is_file()) / 1E6 + return sum(f.stat().st_size for f in path.glob("**/*") if f.is_file()) / 1e6 else: return 0.0 @@ -213,6 +215,7 @@ def file_size(path): def check_online(): # Check internet connectivity import socket + try: socket.create_connection(("1.1.1.1", 443), 5) # check host accessibility return True @@ -224,47 +227,47 @@ def check_online(): @WorkingDirectory(ROOT) def check_git_status(): # Recommend 'git pull' if code is out of date - msg = ', for updates see https://github.com/ultralytics/yolov5' - print(colorstr('github: '), end='') - assert Path('.git').exists(), 'skipping check (not a git repository)' + msg - assert not is_docker(), 'skipping check (Docker image)' + msg - assert check_online(), 'skipping check (offline)' + msg - - cmd = 'git fetch && git config --get remote.origin.url' - url = check_output(cmd, shell=True, timeout=5).decode().strip().rstrip('.git') # git fetch - branch = check_output('git rev-parse --abbrev-ref HEAD', shell=True).decode().strip() # checked out - n = int(check_output(f'git rev-list {branch}..origin/master --count', shell=True)) # commits behind + msg = ", for updates see https://github.com/ultralytics/yolov5" + print(colorstr("github: "), end="") + assert Path(".git").exists(), "skipping check (not a git repository)" + msg + assert not is_docker(), "skipping check (Docker image)" + msg + assert check_online(), "skipping check (offline)" + msg + + cmd = "git fetch && git config --get remote.origin.url" + url = check_output(cmd, shell=True, timeout=5).decode().strip().rstrip(".git") # git fetch + branch = check_output("git rev-parse --abbrev-ref HEAD", shell=True).decode().strip() # checked out + n = int(check_output(f"git rev-list {branch}..origin/master --count", shell=True)) # commits behind if n > 0: s = f"⚠️ YOLOv5 is out of date by {n} commit{'s' * (n > 1)}. Use `git pull` or `git clone {url}` to update." else: - s = f'up to date with {url} ✅' + s = f"up to date with {url} ✅" print(emojis(s)) # emoji-safe -def check_python(minimum='3.6.2'): +def check_python(minimum="3.6.2"): # Check current python version vs. required python version - check_version(platform.python_version(), minimum, name='Python ', hard=True) + check_version(platform.python_version(), minimum, name="Python ", hard=True) -def check_version(current='0.0.0', minimum='0.0.0', name='version ', pinned=False, hard=False): +def check_version(current="0.0.0", minimum="0.0.0", name="version ", pinned=False, hard=False): # Check version vs. required version current, minimum = (pkg.parse_version(x) for x in (current, minimum)) result = (current == minimum) if pinned else (current >= minimum) # bool if hard: # assert min requirements met - assert result, f'{name}{minimum} required by YOLOv5, but {name}{current} is currently installed' + assert result, f"{name}{minimum} required by YOLOv5, but {name}{current} is currently installed" else: return result @try_except -def check_requirements(requirements=ROOT / 'requirements.txt', exclude=(), install=True): +def check_requirements(requirements=ROOT / "requirements.txt", exclude=(), install=True): # Check installed dependencies meet requirements (pass *.txt file or list of packages) - prefix = colorstr('red', 'bold', 'requirements:') + prefix = colorstr("red", "bold", "requirements:") check_python() # check python version if isinstance(requirements, (str, Path)): # requirements.txt file file = Path(requirements) assert file.exists(), f"{prefix} {file.resolve()} not found, check failed." - requirements = [f'{x.name}{x.specifier}' for x in pkg.parse_requirements(file.open()) if x.name not in exclude] + requirements = [f"{x.name}{x.specifier}" for x in pkg.parse_requirements(file.open()) if x.name not in exclude] else: # list or tuple of packages requirements = [x for x in requirements if x not in exclude] @@ -281,14 +284,16 @@ def check_requirements(requirements=ROOT / 'requirements.txt', exclude=(), insta print(check_output(f"pip install '{r}'", shell=True).decode()) n += 1 except Exception as e: - print(f'{prefix} {e}') + print(f"{prefix} {e}") else: - print(f'{s}. Please install and rerun your command.') + print(f"{s}. Please install and rerun your command.") if n: # if packages updated - source = file.resolve() if 'file' in locals() else requirements - s = f"{prefix} {n} package{'s' * (n > 1)} updated per {source}\n" \ + source = file.resolve() if "file" in locals() else requirements + s = ( + f"{prefix} {n} package{'s' * (n > 1)} updated per {source}\n" f"{prefix} ⚠️ {colorstr('bold', 'Restart runtime or rerun command for updates to take effect')}\n" + ) print(emojis(s)) @@ -299,26 +304,26 @@ def check_img_size(imgsz, s=32, floor=0): else: # list i.e. img_size=[640, 480] new_size = [max(make_divisible(x, int(s)), floor) for x in imgsz] if new_size != imgsz: - print(f'WARNING: --img-size {imgsz} must be multiple of max stride {s}, updating to {new_size}') + print(f"WARNING: --img-size {imgsz} must be multiple of max stride {s}, updating to {new_size}") return new_size def check_imshow(): # Check if environment supports image displays try: - assert not is_docker(), 'cv2.imshow() is disabled in Docker environments' - assert not is_colab(), 'cv2.imshow() is disabled in Google Colab environments' - cv2.imshow('test', np.zeros((1, 1, 3))) + assert not is_docker(), "cv2.imshow() is disabled in Docker environments" + assert not is_colab(), "cv2.imshow() is disabled in Google Colab environments" + cv2.imshow("test", np.zeros((1, 1, 3))) cv2.waitKey(1) cv2.destroyAllWindows() cv2.waitKey(1) return True except Exception as e: - print(f'WARNING: Environment does not support cv2.imshow() or PIL Image.show() image displays\n{e}') + print(f"WARNING: Environment does not support cv2.imshow() or PIL Image.show() image displays\n{e}") return False -def check_suffix(file='yolov5s.pt', suffix=('.pt',), msg=''): +def check_suffix(file="yolov5s.pt", suffix=(".pt",), msg=""): # Check file(s) for acceptable suffix if file and suffix: if isinstance(suffix, str): @@ -329,32 +334,32 @@ def check_suffix(file='yolov5s.pt', suffix=('.pt',), msg=''): assert s in suffix, f"{msg}{f} acceptable suffix is {suffix}" -def check_yaml(file, suffix=('.yaml', '.yml')): +def check_yaml(file, suffix=(".yaml", ".yml")): # Search/download YAML file (if necessary) and return path, checking suffix return check_file(file, suffix) -def check_file(file, suffix=''): +def check_file(file, suffix=""): # Search/download file (if necessary) and return path check_suffix(file, suffix) # optional file = str(file) # convert to str() - if Path(file).is_file() or file == '': # exists + if Path(file).is_file() or file == "": # exists return file - elif file.startswith(('http:/', 'https:/')): # download - url = str(Path(file)).replace(':/', '://') # Pathlib turns :// -> :/ - file = Path(urllib.parse.unquote(file).split('?')[0]).name # '%2F' to '/', split https://url.com/file.txt?auth + elif file.startswith(("http:/", "https:/")): # download + url = str(Path(file)).replace(":/", "://") # Pathlib turns :// -> :/ + file = Path(urllib.parse.unquote(file).split("?")[0]).name # '%2F' to '/', split https://url.com/file.txt?auth if Path(file).is_file(): - print(f'Found {url} locally at {file}') # file already exists + print(f"Found {url} locally at {file}") # file already exists else: - print(f'Downloading {url} to {file}...') + print(f"Downloading {url} to {file}...") torch.hub.download_url_to_file(url, file) - assert Path(file).exists() and Path(file).stat().st_size > 0, f'File download failed: {url}' # check + assert Path(file).exists() and Path(file).stat().st_size > 0, f"File download failed: {url}" # check return file else: # search files = [] - for d in 'data', 'models', 'utils': # search directories - files.extend(glob.glob(str(ROOT / d / '**' / file), recursive=True)) # find file - assert len(files), f'File not found: {file}' # assert file was found + for d in "data", "models", "utils": # search directories + files.extend(glob.glob(str(ROOT / d / "**" / file), recursive=True)) # find file + assert len(files), f"File not found: {file}" # assert file was found assert len(files) == 1, f"Multiple files match '{file}', specify exact path: {files}" # assert unique return files[0] # return file @@ -364,61 +369,61 @@ def check_dataset(data, autodownload=True): # Usage: https://github.com/ultralytics/yolov5/releases/download/v1.0/coco128_with_yaml.zip # Download (optional) - extract_dir = '' - if isinstance(data, (str, Path)) and str(data).endswith('.zip'): # i.e. gs://bucket/dir/coco128.zip - download(data, dir='../datasets', unzip=True, delete=False, curl=False, threads=1) - data = next((Path('../datasets') / Path(data).stem).rglob('*.yaml')) + extract_dir = "" + if isinstance(data, (str, Path)) and str(data).endswith(".zip"): # i.e. gs://bucket/dir/coco128.zip + download(data, dir="../datasets", unzip=True, delete=False, curl=False, threads=1) + data = next((Path("../datasets") / Path(data).stem).rglob("*.yaml")) extract_dir, autodownload = data.parent, False # Read yaml (optional) if isinstance(data, (str, Path)): - with open(data, errors='ignore') as f: + with open(data, errors="ignore") as f: data = yaml.safe_load(f) # dictionary # Parse yaml - path = extract_dir or Path(data.get('path') or '') # optional 'path' default to '.' - for k in 'train', 'val', 'test': + path = extract_dir or Path(data.get("path") or "") # optional 'path' default to '.' + for k in "train", "val", "test": if data.get(k): # prepend path data[k] = str(path / data[k]) if isinstance(data[k], str) else [str(path / x) for x in data[k]] - assert 'nc' in data, "Dataset 'nc' key missing." - if 'names' not in data: - data['names'] = [f'class{i}' for i in range(data['nc'])] # assign class names if missing - train, val, test, s = (data.get(x) for x in ('train', 'val', 'test', 'download')) + assert "nc" in data, "Dataset 'nc' key missing." + if "names" not in data: + data["names"] = [f"class{i}" for i in range(data["nc"])] # assign class names if missing + train, val, test, s = (data.get(x) for x in ("train", "val", "test", "download")) if val: val = [Path(x).resolve() for x in (val if isinstance(val, list) else [val])] # val path if not all(x.exists() for x in val): - print('\nWARNING: Dataset not found, nonexistent paths: %s' % [str(x) for x in val if not x.exists()]) + print("\nWARNING: Dataset not found, nonexistent paths: %s" % [str(x) for x in val if not x.exists()]) if s and autodownload: # download script - root = path.parent if 'path' in data else '..' # unzip directory i.e. '../' - if s.startswith('http') and s.endswith('.zip'): # URL + root = path.parent if "path" in data else ".." # unzip directory i.e. '../' + if s.startswith("http") and s.endswith(".zip"): # URL f = Path(s).name # filename - print(f'Downloading {s} to {f}...') + print(f"Downloading {s} to {f}...") torch.hub.download_url_to_file(s, f) Path(root).mkdir(parents=True, exist_ok=True) # create root ZipFile(f).extractall(path=root) # unzip Path(f).unlink() # remove zip r = None # success - elif s.startswith('bash '): # bash script - print(f'Running {s} ...') + elif s.startswith("bash "): # bash script + print(f"Running {s} ...") r = os.system(s) else: # python script - r = exec(s, {'yaml': data}) # return None + r = exec(s, {"yaml": data}) # return None print(f"Dataset autodownload {f'success, saved to {root}' if r in (0, None) else 'failure'}\n") else: - raise Exception('Dataset not found.') + raise Exception("Dataset not found.") return data # dictionary def url2file(url): # Convert URL to filename, i.e. https://url.com/file.txt?auth -> file.txt - url = str(Path(url)).replace(':/', '://') # Pathlib turns :// -> :/ - file = Path(urllib.parse.unquote(url)).name.split('?')[0] # '%2F' to '/', split https://url.com/file.txt?auth + url = str(Path(url)).replace(":/", "://") # Pathlib turns :// -> :/ + file = Path(urllib.parse.unquote(url)).name.split("?")[0] # '%2F' to '/', split https://url.com/file.txt?auth return file -def download(url, dir='.', unzip=True, delete=True, curl=False, threads=1): +def download(url, dir=".", unzip=True, delete=True, curl=False, threads=1): # Multi-threaded file download and unzip function, used in data.yaml for autodownload def download_one(url, dir): # Download 1 file @@ -426,17 +431,17 @@ def download_one(url, dir): if Path(url).is_file(): # exists in current path Path(url).rename(f) # move to dir elif not f.exists(): - print(f'Downloading {url} to {f}...') + print(f"Downloading {url} to {f}...") if curl: os.system(f"curl -L '{url}' -o '{f}' --retry 9 -C -") # curl download, retry and resume on fail else: torch.hub.download_url_to_file(url, f, progress=True) # torch download - if unzip and f.suffix in ('.zip', '.gz'): - print(f'Unzipping {f}...') - if f.suffix == '.zip': + if unzip and f.suffix in (".zip", ".gz"): + print(f"Unzipping {f}...") + if f.suffix == ".zip": ZipFile(f).extractall(path=dir) # unzip - elif f.suffix == '.gz': - os.system(f'tar xfz {f} --directory {f.parent}') # unzip + elif f.suffix == ".gz": + os.system(f"tar xfz {f} --directory {f.parent}") # unzip if delete: f.unlink() # remove zip @@ -469,27 +474,29 @@ def one_cycle(y1=0.0, y2=1.0, steps=100): def colorstr(*input): # Colors a string https://en.wikipedia.org/wiki/ANSI_escape_code, i.e. colorstr('blue', 'hello world') - *args, string = input if len(input) > 1 else ('blue', 'bold', input[0]) # color arguments, string - colors = {'black': '\033[30m', # basic colors - 'red': '\033[31m', - 'green': '\033[32m', - 'yellow': '\033[33m', - 'blue': '\033[34m', - 'magenta': '\033[35m', - 'cyan': '\033[36m', - 'white': '\033[37m', - 'bright_black': '\033[90m', # bright colors - 'bright_red': '\033[91m', - 'bright_green': '\033[92m', - 'bright_yellow': '\033[93m', - 'bright_blue': '\033[94m', - 'bright_magenta': '\033[95m', - 'bright_cyan': '\033[96m', - 'bright_white': '\033[97m', - 'end': '\033[0m', # misc - 'bold': '\033[1m', - 'underline': '\033[4m'} - return ''.join(colors[x] for x in args) + f'{string}' + colors['end'] + *args, string = input if len(input) > 1 else ("blue", "bold", input[0]) # color arguments, string + colors = { + "black": "\033[30m", # basic colors + "red": "\033[31m", + "green": "\033[32m", + "yellow": "\033[33m", + "blue": "\033[34m", + "magenta": "\033[35m", + "cyan": "\033[36m", + "white": "\033[37m", + "bright_black": "\033[90m", # bright colors + "bright_red": "\033[91m", + "bright_green": "\033[92m", + "bright_yellow": "\033[93m", + "bright_blue": "\033[94m", + "bright_magenta": "\033[95m", + "bright_cyan": "\033[96m", + "bright_white": "\033[97m", + "end": "\033[0m", # misc + "bold": "\033[1m", + "underline": "\033[4m", + } + return "".join(colors[x] for x in args) + f"{string}" + colors["end"] def labels_to_class_weights(labels, nc=80): @@ -525,9 +532,88 @@ def coco80_to_coco91_class(): # converts 80-index (val2014) to 91-index (paper) # b = np.loadtxt('data/coco_paper.names', dtype='str', delimiter='\n') # x1 = [list(a[i] == b).index(True) + 1 for i in range(80)] # darknet to coco # x2 = [list(b[i] == a).index(True) if any(b[i] == a) else None for i in range(91)] # coco to darknet - x = [1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 27, 28, 31, 32, 33, 34, - 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 46, 47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, - 64, 65, 67, 70, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 84, 85, 86, 87, 88, 89, 90] + x = [ + 1, + 2, + 3, + 4, + 5, + 6, + 7, + 8, + 9, + 10, + 11, + 13, + 14, + 15, + 16, + 17, + 18, + 19, + 20, + 21, + 22, + 23, + 24, + 25, + 27, + 28, + 31, + 32, + 33, + 34, + 35, + 36, + 37, + 38, + 39, + 40, + 41, + 42, + 43, + 44, + 46, + 47, + 48, + 49, + 50, + 51, + 52, + 53, + 54, + 55, + 56, + 57, + 58, + 59, + 60, + 61, + 62, + 63, + 64, + 65, + 67, + 70, + 72, + 73, + 74, + 75, + 76, + 77, + 78, + 79, + 80, + 81, + 82, + 84, + 85, + 86, + 87, + 88, + 89, + 90, + ] return x @@ -585,7 +671,10 @@ def segment2box(segment, width=640, height=640): # Convert 1 segment label to 1 box label, applying inside-image constraint, i.e. (xy1, xy2, ...) to (xyxy) x, y = segment.T # segment xy inside = (x >= 0) & (y >= 0) & (x <= width) & (y <= height) - x, y, = x[inside], y[inside] + x, y, = ( + x[inside], + y[inside], + ) return np.array([x.min(), y.min(), x.max(), y.max()]) if any(x) else np.zeros((1, 4)) # xyxy @@ -635,8 +724,9 @@ def clip_coords(boxes, shape): boxes[:, [1, 3]] = boxes[:, [1, 3]].clip(0, shape[0]) # y1, y2 -def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, multi_label=False, - labels=(), max_det=300): +def non_max_suppression( + prediction, conf_thres=0.25, iou_thres=0.45, classes=None, agnostic=False, multi_label=False, labels=(), max_det=300 +): """Runs Non-Maximum Suppression (NMS) on inference results Returns: @@ -647,8 +737,8 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non xc = prediction[..., 4] > conf_thres # candidates # Checks - assert 0 <= conf_thres <= 1, f'Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0' - assert 0 <= iou_thres <= 1, f'Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0' + assert 0 <= conf_thres <= 1, f"Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0" + assert 0 <= iou_thres <= 1, f"Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0" # Settings min_wh, max_wh = 2, 4096 # (pixels) minimum and maximum box width and height @@ -713,7 +803,7 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non i = torchvision.ops.nms(boxes, scores, iou_thres) # NMS if i.shape[0] > max_det: # limit detections i = i[:max_det] - if merge and (1 < n < 3E3): # Merge NMS (boxes merged using weighted mean) + if merge and (1 < n < 3e3): # Merge NMS (boxes merged using weighted mean) # update boxes as boxes(i,4) = weights(i,n) * boxes(n,4) iou = box_iou(boxes[i], boxes) > iou_thres # iou matrix weights = iou * scores[None] # box weights @@ -723,65 +813,84 @@ def non_max_suppression(prediction, conf_thres=0.25, iou_thres=0.45, classes=Non output[xi] = x[i] if (time.time() - t) > time_limit: - print(f'WARNING: NMS time limit {time_limit}s exceeded') + print(f"WARNING: NMS time limit {time_limit}s exceeded") break # time limit exceeded return output -def strip_optimizer(f='best.pt', s=''): # from utils.general import *; strip_optimizer() +def strip_optimizer(f="best.pt", s=""): # from utils.general import *; strip_optimizer() # Strip optimizer from 'f' to finalize training, optionally save as 's' - x = torch.load(f, map_location=torch.device('cpu')) - if x.get('ema'): - x['model'] = x['ema'] # replace model with ema - for k in 'optimizer', 'training_results', 'wandb_id', 'ema', 'updates': # keys + x = torch.load(f, map_location=torch.device("cpu")) + if x.get("ema"): + x["model"] = x["ema"] # replace model with ema + for k in "optimizer", "training_results", "wandb_id", "ema", "updates": # keys x[k] = None - x['epoch'] = -1 - x['model'].half() # to FP16 - for p in x['model'].parameters(): + x["epoch"] = -1 + x["model"].half() # to FP16 + for p in x["model"].parameters(): p.requires_grad = False torch.save(x, s or f) - mb = os.path.getsize(s or f) / 1E6 # filesize + mb = os.path.getsize(s or f) / 1e6 # filesize print(f"Optimizer stripped from {f},{(' saved as %s,' % s) if s else ''} {mb:.1f}MB") def print_mutation(results, hyp, save_dir, bucket): - evolve_csv, results_csv, evolve_yaml = save_dir / 'evolve.csv', save_dir / 'results.csv', save_dir / 'hyp_evolve.yaml' - keys = ('metrics/precision', 'metrics/recall', 'metrics/mAP_0.5', 'metrics/mAP_0.5:0.95', - 'val/box_loss', 'val/obj_loss', 'val/cls_loss') + tuple(hyp.keys()) # [results + hyps] + evolve_csv, results_csv, evolve_yaml = ( + save_dir / "evolve.csv", + save_dir / "results.csv", + save_dir / "hyp_evolve.yaml", + ) + keys = ( + "metrics/precision", + "metrics/recall", + "metrics/mAP_0.5", + "metrics/mAP_0.5:0.95", + "val/box_loss", + "val/obj_loss", + "val/cls_loss", + ) + tuple( + hyp.keys() + ) # [results + hyps] keys = tuple(x.strip() for x in keys) vals = results + tuple(hyp.values()) n = len(keys) # Download (optional) if bucket: - url = f'gs://{bucket}/evolve.csv' + url = f"gs://{bucket}/evolve.csv" if gsutil_getsize(url) > (os.path.getsize(evolve_csv) if os.path.exists(evolve_csv) else 0): - os.system(f'gsutil cp {url} {save_dir}') # download evolve.csv if larger than local + os.system(f"gsutil cp {url} {save_dir}") # download evolve.csv if larger than local # Log to evolve.csv - s = '' if evolve_csv.exists() else (('%20s,' * n % keys).rstrip(',') + '\n') # add header - with open(evolve_csv, 'a') as f: - f.write(s + ('%20.5g,' * n % vals).rstrip(',') + '\n') + s = "" if evolve_csv.exists() else (("%20s," * n % keys).rstrip(",") + "\n") # add header + with open(evolve_csv, "a") as f: + f.write(s + ("%20.5g," * n % vals).rstrip(",") + "\n") # Print to screen - print(colorstr('evolve: ') + ', '.join(f'{x.strip():>20s}' for x in keys)) - print(colorstr('evolve: ') + ', '.join(f'{x:20.5g}' for x in vals), end='\n\n\n') + print(colorstr("evolve: ") + ", ".join(f"{x.strip():>20s}" for x in keys)) + print(colorstr("evolve: ") + ", ".join(f"{x:20.5g}" for x in vals), end="\n\n\n") # Save yaml - with open(evolve_yaml, 'w') as f: + with open(evolve_yaml, "w") as f: data = pd.read_csv(evolve_csv) data = data.rename(columns=lambda x: x.strip()) # strip keys i = np.argmax(fitness(data.values[:, :7])) # - f.write('# YOLOv5 Hyperparameter Evolution Results\n' + - f'# Best generation: {i}\n' + - f'# Last generation: {len(data)}\n' + - '# ' + ', '.join(f'{x.strip():>20s}' for x in keys[:7]) + '\n' + - '# ' + ', '.join(f'{x:>20.5g}' for x in data.values[i, :7]) + '\n\n') + f.write( + "# YOLOv5 Hyperparameter Evolution Results\n" + + f"# Best generation: {i}\n" + + f"# Last generation: {len(data)}\n" + + "# " + + ", ".join(f"{x.strip():>20s}" for x in keys[:7]) + + "\n" + + "# " + + ", ".join(f"{x:>20.5g}" for x in data.values[i, :7]) + + "\n\n" + ) yaml.safe_dump(hyp, f, sort_keys=False) if bucket: - os.system(f'gsutil cp {evolve_csv} {evolve_yaml} gs://{bucket}') # upload + os.system(f"gsutil cp {evolve_csv} {evolve_yaml} gs://{bucket}") # upload def apply_classifier(x, model, img, im0): @@ -804,7 +913,7 @@ def apply_classifier(x, model, img, im0): pred_cls1 = d[:, 5].long() ims = [] for j, a in enumerate(d): # per item - cutout = im0[i][int(a[1]):int(a[3]), int(a[0]):int(a[2])] + cutout = im0[i][int(a[1]) : int(a[3]), int(a[0]) : int(a[2])] im = cv2.resize(cutout, (224, 224)) # BGR # cv2.imwrite('example%i.jpg' % j, cutout) @@ -819,11 +928,11 @@ def apply_classifier(x, model, img, im0): return x -def increment_path(path, exist_ok=False, sep='', mkdir=False): +def increment_path(path, exist_ok=False, sep="", mkdir=False): # Increment file or directory path, i.e. runs/exp --> runs/exp{sep}2, runs/exp{sep}3, ... etc. path = Path(path) # os-agnostic if path.exists() and not exist_ok: - path, suffix = (path.with_suffix(''), path.suffix) if path.is_file() else (path, '') + path, suffix = (path.with_suffix(""), path.suffix) if path.is_file() else (path, "") dirs = glob.glob(f"{path}{sep}*") # similar paths matches = [re.search(rf"%s{sep}(\d+)" % path.stem, d) for d in dirs] i = [int(m.groups()[0]) for m in matches if m] # indices diff --git a/utils/loggers/__init__.py b/utils/loggers/__init__.py index ae2d98bdc36d..3de741ccae78 100644 --- a/utils/loggers/__init__.py +++ b/utils/loggers/__init__.py @@ -16,14 +16,14 @@ from utils.plots import plot_images, plot_results from utils.torch_utils import de_parallel -LOGGERS = ('csv', 'tb', 'wandb') # text-file, TensorBoard, Weights & Biases -RANK = int(os.getenv('RANK', -1)) +LOGGERS = ("csv", "tb", "wandb") # text-file, TensorBoard, Weights & Biases +RANK = int(os.getenv("RANK", -1)) try: import wandb - assert hasattr(wandb, '__version__') # verify package import not local dir - if pkg.parse_version(wandb.__version__) >= pkg.parse_version('0.12.2') and RANK in [0, -1]: + assert hasattr(wandb, "__version__") # verify package import not local dir + if pkg.parse_version(wandb.__version__) >= pkg.parse_version("0.12.2") and RANK in [0, -1]: wandb_login_success = wandb.login(timeout=30) if not wandb_login_success: wandb = None @@ -31,7 +31,7 @@ wandb = None -class Loggers(): +class Loggers: # YOLOv5 Loggers class def __init__(self, save_dir=None, weights=None, opt=None, hyp=None, logger=None, include=LOGGERS): self.save_dir = save_dir @@ -40,31 +40,42 @@ def __init__(self, save_dir=None, weights=None, opt=None, hyp=None, logger=None, self.hyp = hyp self.logger = logger # for printing results to console self.include = include - self.keys = ['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', # metrics - 'val/box_loss', 'val/obj_loss', 'val/cls_loss', # val loss - 'x/lr0', 'x/lr1', 'x/lr2'] # params + self.keys = [ + "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", # metrics + "val/box_loss", + "val/obj_loss", + "val/cls_loss", # val loss + "x/lr0", + "x/lr1", + "x/lr2", + ] # params for k in LOGGERS: setattr(self, k, None) # init empty logger dictionary self.csv = True # always log to csv # Message if not wandb: - prefix = colorstr('Weights & Biases: ') + prefix = colorstr("Weights & Biases: ") s = f"{prefix}run 'pip install wandb' to automatically track and visualize YOLOv5 🚀 runs (RECOMMENDED)" print(emojis(s)) # TensorBoard s = self.save_dir - if 'tb' in self.include and not self.opt.evolve: - prefix = colorstr('TensorBoard: ') + if "tb" in self.include and not self.opt.evolve: + prefix = colorstr("TensorBoard: ") self.logger.info(f"{prefix}Start with 'tensorboard --logdir {s.parent}', view at http://localhost:6006/") self.tb = SummaryWriter(str(s)) # W&B - if wandb and 'wandb' in self.include: - wandb_artifact_resume = isinstance(self.opt.resume, str) and self.opt.resume.startswith('wandb-artifact://') - run_id = torch.load(self.weights).get('wandb_id') if self.opt.resume and not wandb_artifact_resume else None + if wandb and "wandb" in self.include: + wandb_artifact_resume = isinstance(self.opt.resume, str) and self.opt.resume.startswith("wandb-artifact://") + run_id = torch.load(self.weights).get("wandb_id") if self.opt.resume and not wandb_artifact_resume else None self.opt.hyp = self.hyp # add hyperparameters self.wandb = WandbLogger(self.opt, run_id) else: @@ -72,7 +83,7 @@ def __init__(self, save_dir=None, weights=None, opt=None, hyp=None, logger=None, def on_pretrain_routine_end(self): # Callback runs on pre-train routine end - paths = self.save_dir.glob('*labels*.jpg') # training labels + paths = self.save_dir.glob("*labels*.jpg") # training labels if self.wandb: self.wandb.log({"Labels": [wandb.Image(str(x), caption=x.name) for x in paths]}) @@ -82,14 +93,14 @@ def on_train_batch_end(self, ni, model, imgs, targets, paths, plots, sync_bn): if ni == 0: if not sync_bn: # tb.add_graph() --sync known issue https://github.com/ultralytics/yolov5/issues/3754 with warnings.catch_warnings(): - warnings.simplefilter('ignore') # suppress jit trace warning + warnings.simplefilter("ignore") # suppress jit trace warning self.tb.add_graph(torch.jit.trace(de_parallel(model), imgs[0:1], strict=False), []) if ni < 3: - f = self.save_dir / f'train_batch{ni}.jpg' # filename + f = self.save_dir / f"train_batch{ni}.jpg" # filename Thread(target=plot_images, args=(imgs, targets, paths, f), daemon=True).start() if self.wandb and ni == 10: - files = sorted(self.save_dir.glob('train*.jpg')) - self.wandb.log({'Mosaics': [wandb.Image(str(f), caption=f.name) for f in files if f.exists()]}) + files = sorted(self.save_dir.glob("train*.jpg")) + self.wandb.log({"Mosaics": [wandb.Image(str(f), caption=f.name) for f in files if f.exists()]}) def on_train_epoch_end(self, epoch): # Callback runs on train epoch end @@ -104,18 +115,18 @@ def on_val_image_end(self, pred, predn, path, names, im): def on_val_end(self): # Callback runs on val end if self.wandb: - files = sorted(self.save_dir.glob('val*.jpg')) + files = sorted(self.save_dir.glob("val*.jpg")) self.wandb.log({"Validation": [wandb.Image(str(f), caption=f.name) for f in files]}) def on_fit_epoch_end(self, vals, epoch, best_fitness, fi): # Callback runs at the end of each fit (train+val) epoch x = {k: v for k, v in zip(self.keys, vals)} # dict if self.csv: - file = self.save_dir / 'results.csv' + file = self.save_dir / "results.csv" n = len(x) + 1 # number of cols - s = '' if file.exists() else (('%20s,' * n % tuple(['epoch'] + self.keys)).rstrip(',') + '\n') # add header - with open(file, 'a') as f: - f.write(s + ('%20.5g,' * n % tuple([epoch] + vals)).rstrip(',') + '\n') + s = "" if file.exists() else (("%20s," * n % tuple(["epoch"] + self.keys)).rstrip(",") + "\n") # add header + with open(file, "a") as f: + f.write(s + ("%20.5g," * n % tuple([epoch] + vals)).rstrip(",") + "\n") if self.tb: for k, v in x.items(): @@ -134,22 +145,26 @@ def on_model_save(self, last, epoch, final_epoch, best_fitness, fi): def on_train_end(self, last, best, plots, epoch, results): # Callback runs on training end if plots: - plot_results(file=self.save_dir / 'results.csv') # save results.png - files = ['results.png', 'confusion_matrix.png', *(f'{x}_curve.png' for x in ('F1', 'PR', 'P', 'R'))] + plot_results(file=self.save_dir / "results.csv") # save results.png + files = ["results.png", "confusion_matrix.png", *(f"{x}_curve.png" for x in ("F1", "PR", "P", "R"))] files = [(self.save_dir / f) for f in files if (self.save_dir / f).exists()] # filter if self.tb: import cv2 + for f in files: - self.tb.add_image(f.stem, cv2.imread(str(f))[..., ::-1], epoch, dataformats='HWC') + self.tb.add_image(f.stem, cv2.imread(str(f))[..., ::-1], epoch, dataformats="HWC") if self.wandb: self.wandb.log({"Results": [wandb.Image(str(f), caption=f.name) for f in files]}) # Calling wandb.log. TODO: Refactor this into WandbLogger.log_model if not self.opt.evolve: - wandb.log_artifact(str(best if best.exists() else last), type='model', - name='run_' + self.wandb.wandb_run.id + '_model', - aliases=['latest', 'best', 'stripped']) + wandb.log_artifact( + str(best if best.exists() else last), + type="model", + name="run_" + self.wandb.wandb_run.id + "_model", + aliases=["latest", "best", "stripped"], + ) self.wandb.finish_run() else: self.wandb.finish_run() diff --git a/utils/loggers/wandb/log_dataset.py b/utils/loggers/wandb/log_dataset.py index 8447272cdb48..9a77de1de386 100644 --- a/utils/loggers/wandb/log_dataset.py +++ b/utils/loggers/wandb/log_dataset.py @@ -2,20 +2,20 @@ from wandb_utils import WandbLogger -WANDB_ARTIFACT_PREFIX = 'wandb-artifact://' +WANDB_ARTIFACT_PREFIX = "wandb-artifact://" def create_dataset_artifact(opt): - logger = WandbLogger(opt, None, job_type='Dataset Creation') # TODO: return value unused + logger = WandbLogger(opt, None, job_type="Dataset Creation") # TODO: return value unused -if __name__ == '__main__': +if __name__ == "__main__": parser = argparse.ArgumentParser() - parser.add_argument('--data', type=str, default='data/coco128.yaml', help='data.yaml path') - parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset') - parser.add_argument('--project', type=str, default='YOLOv5', help='name of W&B Project') - parser.add_argument('--entity', default=None, help='W&B entity') - parser.add_argument('--name', type=str, default='log dataset', help='name of W&B run') + parser.add_argument("--data", type=str, default="data/coco128.yaml", help="data.yaml path") + parser.add_argument("--single-cls", action="store_true", help="train as single-class dataset") + parser.add_argument("--project", type=str, default="YOLOv5", help="name of W&B Project") + parser.add_argument("--entity", default=None, help="W&B entity") + parser.add_argument("--name", type=str, default="log dataset", help="name of W&B run") opt = parser.parse_args() opt.resume = False # Explicitly disallow resume check for dataset upload job diff --git a/utils/loggers/wandb/wandb_utils.py b/utils/loggers/wandb/wandb_utils.py index a71bc6ce96d2..99c2651aa366 100644 --- a/utils/loggers/wandb/wandb_utils.py +++ b/utils/loggers/wandb/wandb_utils.py @@ -22,20 +22,20 @@ try: import wandb - assert hasattr(wandb, '__version__') # verify package import not local dir + assert hasattr(wandb, "__version__") # verify package import not local dir except (ImportError, AssertionError): wandb = None -RANK = int(os.getenv('RANK', -1)) -WANDB_ARTIFACT_PREFIX = 'wandb-artifact://' +RANK = int(os.getenv("RANK", -1)) +WANDB_ARTIFACT_PREFIX = "wandb-artifact://" def remove_prefix(from_string, prefix=WANDB_ARTIFACT_PREFIX): - return from_string[len(prefix):] + return from_string[len(prefix) :] def check_wandb_config_file(data_config_file): - wandb_config = '_wandb.'.join(data_config_file.rsplit('.', 1)) # updated data.yaml path + wandb_config = "_wandb.".join(data_config_file.rsplit(".", 1)) # updated data.yaml path if Path(wandb_config).is_file(): return wandb_config return data_config_file @@ -44,13 +44,15 @@ def check_wandb_config_file(data_config_file): def check_wandb_dataset(data_file): is_trainset_wandb_artifact = False is_valset_wandb_artifact = False - if check_file(data_file) and data_file.endswith('.yaml'): - with open(data_file, errors='ignore') as f: + if check_file(data_file) and data_file.endswith(".yaml"): + with open(data_file, errors="ignore") as f: data_dict = yaml.safe_load(f) - is_trainset_wandb_artifact = (isinstance(data_dict['train'], str) and - data_dict['train'].startswith(WANDB_ARTIFACT_PREFIX)) - is_valset_wandb_artifact = (isinstance(data_dict['val'], str) and - data_dict['val'].startswith(WANDB_ARTIFACT_PREFIX)) + is_trainset_wandb_artifact = isinstance(data_dict["train"], str) and data_dict["train"].startswith( + WANDB_ARTIFACT_PREFIX + ) + is_valset_wandb_artifact = isinstance(data_dict["val"], str) and data_dict["val"].startswith( + WANDB_ARTIFACT_PREFIX + ) if is_trainset_wandb_artifact or is_valset_wandb_artifact: return data_dict else: @@ -62,7 +64,7 @@ def get_run_info(run_path): run_id = run_path.stem project = run_path.parent.stem entity = run_path.parent.parent.stem - model_artifact_name = 'run_' + run_id + '_model' + model_artifact_name = "run_" + run_id + "_model" return entity, project, run_id, model_artifact_name @@ -73,7 +75,7 @@ def check_wandb_resume(opt): if RANK not in [-1, 0]: # For resuming DDP runs entity, project, run_id, model_artifact_name = get_run_info(opt.resume) api = wandb.Api() - artifact = api.artifact(entity + '/' + project + '/' + model_artifact_name + ':latest') + artifact = api.artifact(entity + "/" + project + "/" + model_artifact_name + ":latest") modeldir = artifact.download() opt.weights = str(Path(modeldir) / "last.pt") return True @@ -81,30 +83,30 @@ def check_wandb_resume(opt): def process_wandb_config_ddp_mode(opt): - with open(check_file(opt.data), errors='ignore') as f: + with open(check_file(opt.data), errors="ignore") as f: data_dict = yaml.safe_load(f) # data dict train_dir, val_dir = None, None - if isinstance(data_dict['train'], str) and data_dict['train'].startswith(WANDB_ARTIFACT_PREFIX): + if isinstance(data_dict["train"], str) and data_dict["train"].startswith(WANDB_ARTIFACT_PREFIX): api = wandb.Api() - train_artifact = api.artifact(remove_prefix(data_dict['train']) + ':' + opt.artifact_alias) + train_artifact = api.artifact(remove_prefix(data_dict["train"]) + ":" + opt.artifact_alias) train_dir = train_artifact.download() - train_path = Path(train_dir) / 'data/images/' - data_dict['train'] = str(train_path) + train_path = Path(train_dir) / "data/images/" + data_dict["train"] = str(train_path) - if isinstance(data_dict['val'], str) and data_dict['val'].startswith(WANDB_ARTIFACT_PREFIX): + if isinstance(data_dict["val"], str) and data_dict["val"].startswith(WANDB_ARTIFACT_PREFIX): api = wandb.Api() - val_artifact = api.artifact(remove_prefix(data_dict['val']) + ':' + opt.artifact_alias) + val_artifact = api.artifact(remove_prefix(data_dict["val"]) + ":" + opt.artifact_alias) val_dir = val_artifact.download() - val_path = Path(val_dir) / 'data/images/' - data_dict['val'] = str(val_path) + val_path = Path(val_dir) / "data/images/" + data_dict["val"] = str(val_path) if train_dir or val_dir: - ddp_data_path = str(Path(val_dir) / 'wandb_local_data.yaml') - with open(ddp_data_path, 'w') as f: + ddp_data_path = str(Path(val_dir) / "wandb_local_data.yaml") + with open(ddp_data_path, "w") as f: yaml.safe_dump(data_dict, f) opt.data = ddp_data_path -class WandbLogger(): +class WandbLogger: """Log training runs, datasets, models, and predictions to Weights & Biases. This logger sends information to W&B at wandb.ai. By default, this information @@ -118,7 +120,7 @@ class WandbLogger(): https://docs.wandb.com/guides/integrations/yolov5 """ - def __init__(self, opt, run_id=None, job_type='Training'): + def __init__(self, opt, run_id=None, job_type="Training"): """ - Initialize WandbLogger instance - Upload dataset if opt.upload_dataset is True @@ -129,7 +131,7 @@ def __init__(self, opt, run_id=None, job_type='Training'): run_id (str) -- Run ID of W&B run to be resumed job_type (str) -- To set the job_type for this run - """ + """ # Pre-training routine -- self.job_type = job_type self.wandb, self.wandb_run = wandb, None if not wandb else wandb.run @@ -148,25 +150,29 @@ def __init__(self, opt, run_id=None, job_type='Training'): if opt.resume.startswith(WANDB_ARTIFACT_PREFIX): entity, project, run_id, model_artifact_name = get_run_info(opt.resume) model_artifact_name = WANDB_ARTIFACT_PREFIX + model_artifact_name - assert wandb, 'install wandb to resume wandb runs' + assert wandb, "install wandb to resume wandb runs" # Resume wandb-artifact:// runs here| workaround for not overwriting wandb.config - self.wandb_run = wandb.init(id=run_id, - project=project, - entity=entity, - resume='allow', - allow_val_change=True) + self.wandb_run = wandb.init( + id=run_id, project=project, entity=entity, resume="allow", allow_val_change=True + ) opt.resume = model_artifact_name elif self.wandb: - self.wandb_run = wandb.init(config=opt, - resume="allow", - project='YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem, - entity=opt.entity, - name=opt.name if opt.name != 'exp' else None, - job_type=job_type, - id=run_id, - allow_val_change=True) if not wandb.run else wandb.run + self.wandb_run = ( + wandb.init( + config=opt, + resume="allow", + project="YOLOv5" if opt.project == "runs/train" else Path(opt.project).stem, + entity=opt.entity, + name=opt.name if opt.name != "exp" else None, + job_type=job_type, + id=run_id, + allow_val_change=True, + ) + if not wandb.run + else wandb.run + ) if self.wandb_run: - if self.job_type == 'Training': + if self.job_type == "Training": if opt.upload_dataset: if not opt.resume: self.wandb_artifact_data_dict = self.check_and_upload_dataset(opt) @@ -182,11 +188,10 @@ def __init__(self, opt, run_id=None, job_type='Training'): self.wandb_artifact_data_dict = self.wandb_artifact_data_dict or self.data_dict # write data_dict to config. useful for resuming from artifacts. Do this only when not resuming. - self.wandb_run.config.update({'data_dict': self.wandb_artifact_data_dict}, - allow_val_change=True) + self.wandb_run.config.update({"data_dict": self.wandb_artifact_data_dict}, allow_val_change=True) self.setup_training(opt) - if self.job_type == 'Dataset Creation': + if self.job_type == "Dataset Creation": self.data_dict = self.check_and_upload_dataset(opt) def check_and_upload_dataset(self, opt): @@ -199,12 +204,12 @@ def check_and_upload_dataset(self, opt): returns: Updated dataset info dictionary where local dataset paths are replaced by WAND_ARFACT_PREFIX links. """ - assert wandb, 'Install wandb to upload dataset' - config_path = self.log_dataset_artifact(opt.data, - opt.single_cls, - 'YOLOv5' if opt.project == 'runs/train' else Path(opt.project).stem) + assert wandb, "Install wandb to upload dataset" + config_path = self.log_dataset_artifact( + opt.data, opt.single_cls, "YOLOv5" if opt.project == "runs/train" else Path(opt.project).stem + ) print("Created dataset config file ", config_path) - with open(config_path, errors='ignore') as f: + with open(config_path, errors="ignore") as f: wandb_data_dict = yaml.safe_load(f) return wandb_data_dict @@ -226,22 +231,29 @@ def setup_training(self, opt): if modeldir: self.weights = Path(modeldir) / "last.pt" config = self.wandb_run.config - opt.weights, opt.save_period, opt.batch_size, opt.bbox_interval, opt.epochs, opt.hyp = str( - self.weights), config.save_period, config.batch_size, config.bbox_interval, config.epochs, \ - config.hyp + opt.weights, opt.save_period, opt.batch_size, opt.bbox_interval, opt.epochs, opt.hyp = ( + str(self.weights), + config.save_period, + config.batch_size, + config.bbox_interval, + config.epochs, + config.hyp, + ) data_dict = self.data_dict if self.val_artifact is None: # If --upload_dataset is set, use the existing artifact, don't download - self.train_artifact_path, self.train_artifact = self.download_dataset_artifact(data_dict.get('train'), - opt.artifact_alias) - self.val_artifact_path, self.val_artifact = self.download_dataset_artifact(data_dict.get('val'), - opt.artifact_alias) + self.train_artifact_path, self.train_artifact = self.download_dataset_artifact( + data_dict.get("train"), opt.artifact_alias + ) + self.val_artifact_path, self.val_artifact = self.download_dataset_artifact( + data_dict.get("val"), opt.artifact_alias + ) if self.train_artifact_path is not None: - train_path = Path(self.train_artifact_path) / 'data/images/' - data_dict['train'] = str(train_path) + train_path = Path(self.train_artifact_path) / "data/images/" + data_dict["train"] = str(train_path) if self.val_artifact_path is not None: - val_path = Path(self.val_artifact_path) / 'data/images/' - data_dict['val'] = str(val_path) + val_path = Path(self.val_artifact_path) / "data/images/" + data_dict["val"] = str(val_path) if self.val_artifact is not None: self.result_artifact = wandb.Artifact("run_" + wandb.run.id + "_progress", "evaluation") @@ -271,7 +283,7 @@ def download_dataset_artifact(self, path, alias): if isinstance(path, str) and path.startswith(WANDB_ARTIFACT_PREFIX): artifact_path = Path(remove_prefix(path, WANDB_ARTIFACT_PREFIX) + ":" + alias) dataset_artifact = wandb.use_artifact(artifact_path.as_posix().replace("\\", "/")) - assert dataset_artifact is not None, "'Error: W&B dataset artifact doesn\'t exist'" + assert dataset_artifact is not None, "'Error: W&B dataset artifact doesn't exist'" datadir = dataset_artifact.download() return datadir, dataset_artifact return None, None @@ -285,12 +297,12 @@ def download_model_artifact(self, opt): """ if opt.resume.startswith(WANDB_ARTIFACT_PREFIX): model_artifact = wandb.use_artifact(remove_prefix(opt.resume, WANDB_ARTIFACT_PREFIX) + ":latest") - assert model_artifact is not None, 'Error: W&B model artifact doesn\'t exist' + assert model_artifact is not None, "Error: W&B model artifact doesn't exist" modeldir = model_artifact.download() - epochs_trained = model_artifact.metadata.get('epochs_trained') - total_epochs = model_artifact.metadata.get('total_epochs') + epochs_trained = model_artifact.metadata.get("epochs_trained") + total_epochs = model_artifact.metadata.get("total_epochs") is_finished = total_epochs is None - assert not is_finished, 'training is finished, can only resume incomplete runs.' + assert not is_finished, "training is finished, can only resume incomplete runs." return modeldir, model_artifact return None, None @@ -305,17 +317,22 @@ def log_model(self, path, opt, epoch, fitness_score, best_model=False): fitness_score (float) -- fitness score for current epoch best_model (boolean) -- Boolean representing if the current checkpoint is the best yet. """ - model_artifact = wandb.Artifact('run_' + wandb.run.id + '_model', type='model', metadata={ - 'original_url': str(path), - 'epochs_trained': epoch + 1, - 'save period': opt.save_period, - 'project': opt.project, - 'total_epochs': opt.epochs, - 'fitness_score': fitness_score - }) - model_artifact.add_file(str(path / 'last.pt'), name='last.pt') - wandb.log_artifact(model_artifact, - aliases=['latest', 'last', 'epoch ' + str(self.current_epoch), 'best' if best_model else '']) + model_artifact = wandb.Artifact( + "run_" + wandb.run.id + "_model", + type="model", + metadata={ + "original_url": str(path), + "epochs_trained": epoch + 1, + "save period": opt.save_period, + "project": opt.project, + "total_epochs": opt.epochs, + "fitness_score": fitness_score, + }, + ) + model_artifact.add_file(str(path / "last.pt"), name="last.pt") + wandb.log_artifact( + model_artifact, aliases=["latest", "last", "epoch " + str(self.current_epoch), "best" if best_model else ""] + ) print("Saving model artifact on epoch ", epoch + 1) def log_dataset_artifact(self, data_file, single_cls, project, overwrite_config=False): @@ -334,28 +351,34 @@ def log_dataset_artifact(self, data_file, single_cls, project, overwrite_config= """ self.data_dict = check_dataset(data_file) # parse and check data = dict(self.data_dict) - nc, names = (1, ['item']) if single_cls else (int(data['nc']), data['names']) + nc, names = (1, ["item"]) if single_cls else (int(data["nc"]), data["names"]) names = {k: v for k, v in enumerate(names)} # to index dictionary - self.train_artifact = self.create_dataset_table(LoadImagesAndLabels( - data['train'], rect=True, batch_size=1), names, name='train') if data.get('train') else None - self.val_artifact = self.create_dataset_table(LoadImagesAndLabels( - data['val'], rect=True, batch_size=1), names, name='val') if data.get('val') else None - if data.get('train'): - data['train'] = WANDB_ARTIFACT_PREFIX + str(Path(project) / 'train') - if data.get('val'): - data['val'] = WANDB_ARTIFACT_PREFIX + str(Path(project) / 'val') + self.train_artifact = ( + self.create_dataset_table(LoadImagesAndLabels(data["train"], rect=True, batch_size=1), names, name="train") + if data.get("train") + else None + ) + self.val_artifact = ( + self.create_dataset_table(LoadImagesAndLabels(data["val"], rect=True, batch_size=1), names, name="val") + if data.get("val") + else None + ) + if data.get("train"): + data["train"] = WANDB_ARTIFACT_PREFIX + str(Path(project) / "train") + if data.get("val"): + data["val"] = WANDB_ARTIFACT_PREFIX + str(Path(project) / "val") path = Path(data_file).stem - path = (path if overwrite_config else path + '_wandb') + '.yaml' # updated data.yaml path - data.pop('download', None) - data.pop('path', None) - with open(path, 'w') as f: + path = (path if overwrite_config else path + "_wandb") + ".yaml" # updated data.yaml path + data.pop("download", None) + data.pop("path", None) + with open(path, "w") as f: yaml.safe_dump(data, f) - if self.job_type == 'Training': # builds correct artifact pipeline graph + if self.job_type == "Training": # builds correct artifact pipeline graph self.wandb_run.use_artifact(self.val_artifact) self.wandb_run.use_artifact(self.train_artifact) self.val_artifact.wait() - self.val_table = self.val_artifact.get('val') + self.val_table = self.val_artifact.get("val") self.map_val_table_path() else: self.wandb_run.log_artifact(self.train_artifact) @@ -372,7 +395,7 @@ def map_val_table_path(self): for i, data in enumerate(tqdm(self.val_table.data)): self.val_table_path_map[data[3]] = data[0] - def create_dataset_table(self, dataset: LoadImagesAndLabels, class_to_id: Dict[int,str], name: str = 'dataset'): + def create_dataset_table(self, dataset: LoadImagesAndLabels, class_to_id: Dict[int, str], name: str = "dataset"): """ Create and return W&B artifact containing W&B Table of the dataset. @@ -390,27 +413,33 @@ def create_dataset_table(self, dataset: LoadImagesAndLabels, class_to_id: Dict[i img_files = tqdm(dataset.img_files) if not img_files else img_files for img_file in img_files: if Path(img_file).is_dir(): - artifact.add_dir(img_file, name='data/images') - labels_path = 'labels'.join(dataset.path.rsplit('images', 1)) - artifact.add_dir(labels_path, name='data/labels') + artifact.add_dir(img_file, name="data/images") + labels_path = "labels".join(dataset.path.rsplit("images", 1)) + artifact.add_dir(labels_path, name="data/labels") else: - artifact.add_file(img_file, name='data/images/' + Path(img_file).name) + artifact.add_file(img_file, name="data/images/" + Path(img_file).name) label_file = Path(img2label_paths([img_file])[0]) - artifact.add_file(str(label_file), - name='data/labels/' + label_file.name) if label_file.exists() else None + artifact.add_file( + str(label_file), name="data/labels/" + label_file.name + ) if label_file.exists() else None table = wandb.Table(columns=["id", "train_image", "Classes", "name"]) - class_set = wandb.Classes([{'id': id, 'name': name} for id, name in class_to_id.items()]) + class_set = wandb.Classes([{"id": id, "name": name} for id, name in class_to_id.items()]) for si, (img, labels, paths, shapes) in enumerate(tqdm(dataset)): box_data, img_classes = [], {} for cls, *xywh in labels[:, 1:].tolist(): cls = int(cls) - box_data.append({"position": {"middle": [xywh[0], xywh[1]], "width": xywh[2], "height": xywh[3]}, - "class_id": cls, - "box_caption": "%s" % (class_to_id[cls])}) + box_data.append( + { + "position": {"middle": [xywh[0], xywh[1]], "width": xywh[2], "height": xywh[3]}, + "class_id": cls, + "box_caption": "%s" % (class_to_id[cls]), + } + ) img_classes[cls] = class_to_id[cls] boxes = {"ground_truth": {"box_data": box_data, "class_labels": class_to_id}} # inference-space - table.add_data(si, wandb.Image(paths, classes=class_set, boxes=boxes), list(img_classes.values()), - Path(paths).name) + table.add_data( + si, wandb.Image(paths, classes=class_set, boxes=boxes), list(img_classes.values()), Path(paths).name + ) artifact.add(table, name) return artifact @@ -423,26 +452,30 @@ def log_training_progress(self, predn, path, names): path (str): local path of the current evaluation image names (dict(int, str)): hash map that maps class ids to labels """ - class_set = wandb.Classes([{'id': id, 'name': name} for id, name in names.items()]) + class_set = wandb.Classes([{"id": id, "name": name} for id, name in names.items()]) box_data = [] total_conf = 0 for *xyxy, conf, cls in predn.tolist(): if conf >= 0.25: box_data.append( - {"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]}, - "class_id": int(cls), - "box_caption": f"{names[cls]} {conf:.3f}", - "scores": {"class_score": conf}, - "domain": "pixel"}) + { + "position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]}, + "class_id": int(cls), + "box_caption": f"{names[cls]} {conf:.3f}", + "scores": {"class_score": conf}, + "domain": "pixel", + } + ) total_conf += conf boxes = {"predictions": {"box_data": box_data, "class_labels": names}} # inference-space id = self.val_table_path_map[Path(path).name] - self.result_table.add_data(self.current_epoch, - id, - self.val_table.data[id][1], - wandb.Image(self.val_table.data[id][1], boxes=boxes, classes=class_set), - total_conf / max(1, len(box_data)) - ) + self.result_table.add_data( + self.current_epoch, + id, + self.val_table.data[id][1], + wandb.Image(self.val_table.data[id][1], boxes=boxes, classes=class_set), + total_conf / max(1, len(box_data)), + ) def val_one_image(self, pred, predn, path, names, im): """ @@ -458,11 +491,16 @@ def val_one_image(self, pred, predn, path, names, im): if len(self.bbox_media_panel_images) < self.max_imgs_to_log and self.current_epoch > 0: if self.current_epoch % self.bbox_interval == 0: - box_data = [{"position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]}, - "class_id": int(cls), - "box_caption": f"{names[cls]} {conf:.3f}", - "scores": {"class_score": conf}, - "domain": "pixel"} for *xyxy, conf, cls in pred.tolist()] + box_data = [ + { + "position": {"minX": xyxy[0], "minY": xyxy[1], "maxX": xyxy[2], "maxY": xyxy[3]}, + "class_id": int(cls), + "box_caption": f"{names[cls]} {conf:.3f}", + "scores": {"class_score": conf}, + "domain": "pixel", + } + for *xyxy, conf, cls in pred.tolist() + ] boxes = {"predictions": {"box_data": box_data, "class_labels": names}} # inference-space self.bbox_media_panel_images.append(wandb.Image(im, boxes=boxes, caption=path.name)) @@ -492,9 +530,11 @@ def end_epoch(self, best_result=False): self.log_dict = {} self.bbox_media_panel_images = [] if self.result_artifact: - self.result_artifact.add(self.result_table, 'result') - wandb.log_artifact(self.result_artifact, aliases=['latest', 'last', 'epoch ' + str(self.current_epoch), - ('best' if best_result else '')]) + self.result_artifact.add(self.result_table, "result") + wandb.log_artifact( + self.result_artifact, + aliases=["latest", "last", "epoch " + str(self.current_epoch), ("best" if best_result else "")], + ) wandb.log({"evaluation": self.result_table}) self.result_table = wandb.Table(["epoch", "id", "ground truth", "prediction", "avg_confidence"]) @@ -513,7 +553,7 @@ def finish_run(self): @contextmanager def all_logging_disabled(highest_level=logging.CRITICAL): - """ source - https://gist.github.com/simon-weber/7853144 + """source - https://gist.github.com/simon-weber/7853144 A context manager that will prevent any logging messages triggered during the body from being processed. :param highest_level: the maximum logging level in use. This would only need to be changed if a custom level greater than CRITICAL is defined. diff --git a/utils/loss.py b/utils/loss.py index 194c8e503e0e..7a6f60a9b934 100644 --- a/utils/loss.py +++ b/utils/loss.py @@ -19,7 +19,7 @@ class BCEBlurWithLogitsLoss(nn.Module): # BCEwithLogitLoss() with reduced missing label effects. def __init__(self, alpha=0.05): super().__init__() - self.loss_fcn = nn.BCEWithLogitsLoss(reduction='none') # must be nn.BCEWithLogitsLoss() + self.loss_fcn = nn.BCEWithLogitsLoss(reduction="none") # must be nn.BCEWithLogitsLoss() self.alpha = alpha def forward(self, pred, true): @@ -40,7 +40,7 @@ def __init__(self, loss_fcn, gamma=1.5, alpha=0.25): self.gamma = gamma self.alpha = alpha self.reduction = loss_fcn.reduction - self.loss_fcn.reduction = 'none' # required to apply FL to each element + self.loss_fcn.reduction = "none" # required to apply FL to each element def forward(self, pred, true): loss = self.loss_fcn(pred, true) @@ -54,9 +54,9 @@ def forward(self, pred, true): modulating_factor = (1.0 - p_t) ** self.gamma loss *= alpha_factor * modulating_factor - if self.reduction == 'mean': + if self.reduction == "mean": return loss.mean() - elif self.reduction == 'sum': + elif self.reduction == "sum": return loss.sum() else: # 'none' return loss @@ -70,7 +70,7 @@ def __init__(self, loss_fcn, gamma=1.5, alpha=0.25): self.gamma = gamma self.alpha = alpha self.reduction = loss_fcn.reduction - self.loss_fcn.reduction = 'none' # required to apply FL to each element + self.loss_fcn.reduction = "none" # required to apply FL to each element def forward(self, pred, true): loss = self.loss_fcn(pred, true) @@ -80,9 +80,9 @@ def forward(self, pred, true): modulating_factor = torch.abs(true - pred_prob) ** self.gamma loss *= alpha_factor * modulating_factor - if self.reduction == 'mean': + if self.reduction == "mean": return loss.mean() - elif self.reduction == 'sum': + elif self.reduction == "sum": return loss.sum() else: # 'none' return loss @@ -96,14 +96,14 @@ def __init__(self, model, autobalance=False): h = model.hyp # hyperparameters # Define criteria - BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['cls_pw']], device=device)) - BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h['obj_pw']], device=device)) + BCEcls = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h["cls_pw"]], device=device)) + BCEobj = nn.BCEWithLogitsLoss(pos_weight=torch.tensor([h["obj_pw"]], device=device)) # Class label smoothing https://arxiv.org/pdf/1902.04103.pdf eqn 3 - self.cp, self.cn = smooth_BCE(eps=h.get('label_smoothing', 0.0)) # positive, negative BCE targets + self.cp, self.cn = smooth_BCE(eps=h.get("label_smoothing", 0.0)) # positive, negative BCE targets # Focal loss - g = h['fl_gamma'] # focal loss gamma + g = h["fl_gamma"] # focal loss gamma if g > 0: BCEcls, BCEobj = FocalLoss(BCEcls, g), FocalLoss(BCEobj, g) @@ -111,7 +111,7 @@ def __init__(self, model, autobalance=False): self.balance = {3: [4.0, 1.0, 0.4]}.get(det.nl, [4.0, 1.0, 0.25, 0.06, 0.02]) # P3-P7 self.ssi = list(det.stride).index(16) if autobalance else 0 # stride 16 index self.BCEcls, self.BCEobj, self.gr, self.hyp, self.autobalance = BCEcls, BCEobj, 1.0, h, autobalance - for k in 'na', 'nc', 'nl', 'anchors': + for k in "na", "nc", "nl", "anchors": setattr(self, k, getattr(det, k)) def __call__(self, p, targets): # predictions, targets, model @@ -159,9 +159,9 @@ def __call__(self, p, targets): # predictions, targets, model if self.autobalance: self.balance = [x / self.balance[self.ssi] for x in self.balance] - lbox *= self.hyp['box'] - lobj *= self.hyp['obj'] - lcls *= self.hyp['cls'] + lbox *= self.hyp["box"] + lobj *= self.hyp["obj"] + lcls *= self.hyp["cls"] bs = tobj.shape[0] # batch size return (lbox + lobj + lcls) * bs, torch.cat((lbox, lobj, lcls)).detach() @@ -175,10 +175,20 @@ def build_targets(self, p, targets): targets = torch.cat((targets.repeat(na, 1, 1), ai[:, :, None]), 2) # append anchor indices g = 0.5 # bias - off = torch.tensor([[0, 0], - [1, 0], [0, 1], [-1, 0], [0, -1], # j,k,l,m - # [1, 1], [1, -1], [-1, 1], [-1, -1], # jk,jm,lk,lm - ], device=targets.device).float() * g # offsets + off = ( + torch.tensor( + [ + [0, 0], + [1, 0], + [0, 1], + [-1, 0], + [0, -1], # j,k,l,m + # [1, 1], [1, -1], [-1, 1], [-1, -1], # jk,jm,lk,lm + ], + device=targets.device, + ).float() + * g + ) # offsets for i in range(self.nl): anchors = self.anchors[i] @@ -189,7 +199,7 @@ def build_targets(self, p, targets): if nt: # Matches r = t[:, :, 4:6] / anchors[:, None] # wh ratio - j = torch.max(r, 1 / r).max(2)[0] < self.hyp['anchor_t'] # compare + j = torch.max(r, 1 / r).max(2)[0] < self.hyp["anchor_t"] # compare # j = wh_iou(anchors, t[:, 4:6]) > model.hyp['iou_t'] # iou(3,n)=wh_iou(anchors(3,2), gwh(n,2)) t = t[j] # filter diff --git a/utils/metrics.py b/utils/metrics.py index 2e0e0c65e63d..2d8bc657919c 100644 --- a/utils/metrics.py +++ b/utils/metrics.py @@ -18,8 +18,8 @@ def fitness(x): return (x[:, :4] * w).sum(1) -def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir='.', names=()): - """ Compute the average precision, given the recall and precision curves. +def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir=".", names=()): + """Compute the average precision, given the recall and precision curves. Source: https://github.com/rafaelpadilla/Object-Detection-Metrics. # Arguments tp: True positives (nparray, nx1 or nx10). @@ -74,17 +74,17 @@ def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir='.', names names = [v for k, v in names.items() if k in unique_classes] # list: only classes that have data names = {i: v for i, v in enumerate(names)} # to dict if plot: - plot_pr_curve(px, py, ap, Path(save_dir) / 'PR_curve.png', names) - plot_mc_curve(px, f1, Path(save_dir) / 'F1_curve.png', names, ylabel='F1') - plot_mc_curve(px, p, Path(save_dir) / 'P_curve.png', names, ylabel='Precision') - plot_mc_curve(px, r, Path(save_dir) / 'R_curve.png', names, ylabel='Recall') + plot_pr_curve(px, py, ap, Path(save_dir) / "PR_curve.png", names) + plot_mc_curve(px, f1, Path(save_dir) / "F1_curve.png", names, ylabel="F1") + plot_mc_curve(px, p, Path(save_dir) / "P_curve.png", names, ylabel="Precision") + plot_mc_curve(px, r, Path(save_dir) / "R_curve.png", names, ylabel="Recall") i = f1.mean(0).argmax() # max F1 index - return p[:, i], r[:, i], ap, f1[:, i], unique_classes.astype('int32') + return p[:, i], r[:, i], ap, f1[:, i], unique_classes.astype("int32") def compute_ap(recall, precision): - """ Compute the average precision, given the recall and precision curves + """Compute the average precision, given the recall and precision curves # Arguments recall: The recall curve (list) precision: The precision curve (list) @@ -100,8 +100,8 @@ def compute_ap(recall, precision): mpre = np.flip(np.maximum.accumulate(np.flip(mpre))) # Integrate area under curve - method = 'interp' # methods: 'continuous', 'interp' - if method == 'interp': + method = "interp" # methods: 'continuous', 'interp' + if method == "interp": x = np.linspace(0, 1, 101) # 101-point interp (COCO) ap = np.trapz(np.interp(x, mrec, mpre), x) # integrate else: # 'continuous' @@ -162,31 +162,38 @@ def process_batch(self, detections, labels): def matrix(self): return self.matrix - def plot(self, normalize=True, save_dir='', names=()): + def plot(self, normalize=True, save_dir="", names=()): try: import seaborn as sn - array = self.matrix / ((self.matrix.sum(0).reshape(1, -1) + 1E-6) if normalize else 1) # normalize columns + array = self.matrix / ((self.matrix.sum(0).reshape(1, -1) + 1e-6) if normalize else 1) # normalize columns array[array < 0.005] = np.nan # don't annotate (would appear as 0.00) fig = plt.figure(figsize=(12, 9), tight_layout=True) sn.set(font_scale=1.0 if self.nc < 50 else 0.8) # for label size labels = (0 < len(names) < 99) and len(names) == self.nc # apply names to ticklabels with warnings.catch_warnings(): - warnings.simplefilter('ignore') # suppress empty matrix RuntimeWarning: All-NaN slice encountered - sn.heatmap(array, annot=self.nc < 30, annot_kws={"size": 8}, cmap='Blues', fmt='.2f', square=True, - xticklabels=names + ['background FP'] if labels else "auto", - yticklabels=names + ['background FN'] if labels else "auto").set_facecolor((1, 1, 1)) - fig.axes[0].set_xlabel('True') - fig.axes[0].set_ylabel('Predicted') - fig.savefig(Path(save_dir) / 'confusion_matrix.png', dpi=250) + warnings.simplefilter("ignore") # suppress empty matrix RuntimeWarning: All-NaN slice encountered + sn.heatmap( + array, + annot=self.nc < 30, + annot_kws={"size": 8}, + cmap="Blues", + fmt=".2f", + square=True, + xticklabels=names + ["background FP"] if labels else "auto", + yticklabels=names + ["background FN"] if labels else "auto", + ).set_facecolor((1, 1, 1)) + fig.axes[0].set_xlabel("True") + fig.axes[0].set_ylabel("Predicted") + fig.savefig(Path(save_dir) / "confusion_matrix.png", dpi=250) plt.close() except Exception as e: - print(f'WARNING: ConfusionMatrix plot failure: {e}') + print(f"WARNING: ConfusionMatrix plot failure: {e}") def print(self): for i in range(self.nc + 1): - print(' '.join(map(str, self.matrix[i]))) + print(" ".join(map(str, self.matrix[i]))) def bbox_iou(box1, box2, x1y1x2y2=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-7): @@ -204,8 +211,9 @@ def bbox_iou(box1, box2, x1y1x2y2=True, GIoU=False, DIoU=False, CIoU=False, eps= b2_y1, b2_y2 = box2[1] - box2[3] / 2, box2[1] + box2[3] / 2 # Intersection area - inter = (torch.min(b1_x2, b2_x2) - torch.max(b1_x1, b2_x1)).clamp(0) * \ - (torch.min(b1_y2, b2_y2) - torch.max(b1_y1, b2_y1)).clamp(0) + inter = (torch.min(b1_x2, b2_x2) - torch.max(b1_x1, b2_x1)).clamp(0) * ( + torch.min(b1_y2, b2_y2) - torch.max(b1_y1, b2_y1) + ).clamp(0) # Union Area w1, h1 = b1_x2 - b1_x1, b1_y2 - b1_y1 + eps @@ -218,8 +226,9 @@ def bbox_iou(box1, box2, x1y1x2y2=True, GIoU=False, DIoU=False, CIoU=False, eps= ch = torch.max(b1_y2, b2_y2) - torch.min(b1_y1, b2_y1) # convex height if CIoU or DIoU: # Distance or Complete IoU https://arxiv.org/abs/1911.08287v1 c2 = cw ** 2 + ch ** 2 + eps # convex diagonal squared - rho2 = ((b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2 + - (b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2) / 4 # center distance squared + rho2 = ( + (b2_x1 + b2_x2 - b1_x1 - b1_x2) ** 2 + (b2_y1 + b2_y2 - b1_y1 - b1_y2) ** 2 + ) / 4 # center distance squared if DIoU: return iou - rho2 / c2 # DIoU elif CIoU: # https://github.com/Zzh-tju/DIoU-SSD-pytorch/blob/master/utils/box/box_utils.py#L47 @@ -259,8 +268,8 @@ def box_area(box): return inter / (area1[:, None] + area2 - inter) # iou = inter / (area1 + area2 - inter) -def bbox_ioa(box1, box2, eps=1E-7): - """ Returns the intersection over box2 area given box1, box2. Boxes are x1y1x2y2 +def bbox_ioa(box1, box2, eps=1e-7): + """Returns the intersection over box2 area given box1, box2. Boxes are x1y1x2y2 box1: np.array of shape(4) box2: np.array of shape(nx4) returns: np.array of shape(n) @@ -273,8 +282,9 @@ def bbox_ioa(box1, box2, eps=1E-7): b2_x1, b2_y1, b2_x2, b2_y2 = box2[0], box2[1], box2[2], box2[3] # Intersection area - inter_area = (np.minimum(b1_x2, b2_x2) - np.maximum(b1_x1, b2_x1)).clip(0) * \ - (np.minimum(b1_y2, b2_y2) - np.maximum(b1_y1, b2_y1)).clip(0) + inter_area = (np.minimum(b1_x2, b2_x2) - np.maximum(b1_x1, b2_x1)).clip(0) * ( + np.minimum(b1_y2, b2_y2) - np.maximum(b1_y1, b2_y1) + ).clip(0) # box2 area box2_area = (b2_x2 - b2_x1) * (b2_y2 - b2_y1) + eps @@ -293,20 +303,21 @@ def wh_iou(wh1, wh2): # Plots ---------------------------------------------------------------------------------------------------------------- -def plot_pr_curve(px, py, ap, save_dir='pr_curve.png', names=()): + +def plot_pr_curve(px, py, ap, save_dir="pr_curve.png", names=()): # Precision-recall curve fig, ax = plt.subplots(1, 1, figsize=(9, 6), tight_layout=True) py = np.stack(py, axis=1) if 0 < len(names) < 21: # display per-class legend if < 21 classes for i, y in enumerate(py.T): - ax.plot(px, y, linewidth=1, label=f'{names[i]} {ap[i, 0]:.3f}') # plot(recall, precision) + ax.plot(px, y, linewidth=1, label=f"{names[i]} {ap[i, 0]:.3f}") # plot(recall, precision) else: - ax.plot(px, py, linewidth=1, color='grey') # plot(recall, precision) + ax.plot(px, py, linewidth=1, color="grey") # plot(recall, precision) - ax.plot(px, py.mean(1), linewidth=3, color='blue', label='all classes %.3f mAP@0.5' % ap[:, 0].mean()) - ax.set_xlabel('Recall') - ax.set_ylabel('Precision') + ax.plot(px, py.mean(1), linewidth=3, color="blue", label="all classes %.3f mAP@0.5" % ap[:, 0].mean()) + ax.set_xlabel("Recall") + ax.set_ylabel("Precision") ax.set_xlim(0, 1) ax.set_ylim(0, 1) plt.legend(bbox_to_anchor=(1.04, 1), loc="upper left") @@ -314,18 +325,18 @@ def plot_pr_curve(px, py, ap, save_dir='pr_curve.png', names=()): plt.close() -def plot_mc_curve(px, py, save_dir='mc_curve.png', names=(), xlabel='Confidence', ylabel='Metric'): +def plot_mc_curve(px, py, save_dir="mc_curve.png", names=(), xlabel="Confidence", ylabel="Metric"): # Metric-confidence curve fig, ax = plt.subplots(1, 1, figsize=(9, 6), tight_layout=True) if 0 < len(names) < 21: # display per-class legend if < 21 classes for i, y in enumerate(py): - ax.plot(px, y, linewidth=1, label=f'{names[i]}') # plot(confidence, metric) + ax.plot(px, y, linewidth=1, label=f"{names[i]}") # plot(confidence, metric) else: - ax.plot(px, py.T, linewidth=1, color='grey') # plot(confidence, metric) + ax.plot(px, py.T, linewidth=1, color="grey") # plot(confidence, metric) y = py.mean(0) - ax.plot(px, y, linewidth=3, color='blue', label=f'all classes {y.max():.2f} at {px[y.argmax()]:.3f}') + ax.plot(px, y, linewidth=3, color="blue", label=f"all classes {y.max():.2f} at {px[y.argmax()]:.3f}") ax.set_xlabel(xlabel) ax.set_ylabel(ylabel) ax.set_xlim(0, 1) diff --git a/utils/plots.py b/utils/plots.py index b5e25d668d22..34b873e08750 100644 --- a/utils/plots.py +++ b/utils/plots.py @@ -22,18 +22,38 @@ # Settings CONFIG_DIR = user_config_dir() # Ultralytics settings dir -RANK = int(os.getenv('RANK', -1)) -matplotlib.rc('font', **{'size': 11}) -matplotlib.use('Agg') # for writing to files only +RANK = int(os.getenv("RANK", -1)) +matplotlib.rc("font", **{"size": 11}) +matplotlib.use("Agg") # for writing to files only class Colors: # Ultralytics color palette https://ultralytics.com/ def __init__(self): # hex = matplotlib.colors.TABLEAU_COLORS.values() - hex = ('FF3838', 'FF9D97', 'FF701F', 'FFB21D', 'CFD231', '48F90A', '92CC17', '3DDB86', '1A9334', '00D4BB', - '2C99A8', '00C2FF', '344593', '6473FF', '0018EC', '8438FF', '520085', 'CB38FF', 'FF95C8', 'FF37C7') - self.palette = [self.hex2rgb('#' + c) for c in hex] + hex = ( + "FF3838", + "FF9D97", + "FF701F", + "FFB21D", + "CFD231", + "48F90A", + "92CC17", + "3DDB86", + "1A9334", + "00D4BB", + "2C99A8", + "00C2FF", + "344593", + "6473FF", + "0018EC", + "8438FF", + "520085", + "CB38FF", + "FF95C8", + "FF37C7", + ) + self.palette = [self.hex2rgb("#" + c) for c in hex] self.n = len(self.palette) def __call__(self, i, bgr=False): @@ -42,13 +62,13 @@ def __call__(self, i, bgr=False): @staticmethod def hex2rgb(h): # rgb order (PIL) - return tuple(int(h[1 + i:1 + i + 2], 16) for i in (0, 2, 4)) + return tuple(int(h[1 + i : 1 + i + 2], 16) for i in (0, 2, 4)) colors = Colors() # create instance for 'from utils.plots import colors' -def check_font(font='Arial.ttf', size=10): +def check_font(font="Arial.ttf", size=10): # Return a PIL TrueType Font, downloading to CONFIG_DIR if necessary font = Path(font) font = font if font.exists() else (CONFIG_DIR / font.name) @@ -56,7 +76,7 @@ def check_font(font='Arial.ttf', size=10): return ImageFont.truetype(str(font) if font.exists() else font.name, size) except Exception as e: # download if missing url = "https://ultralytics.com/assets/" + font.name - print(f'Downloading {url} to {font}...') + print(f"Downloading {url} to {font}...") torch.hub.download_url_to_file(url, str(font), progress=False) return ImageFont.truetype(str(font), size) @@ -66,29 +86,36 @@ class Annotator: check_font() # download TTF if necessary # YOLOv5 Annotator for train/val mosaics and jpgs and detect/hub inference annotations - def __init__(self, im, line_width=None, font_size=None, font='Arial.ttf', pil=False, example='abc'): - assert im.data.contiguous, 'Image not contiguous. Apply np.ascontiguousarray(im) to Annotator() input images.' + def __init__(self, im, line_width=None, font_size=None, font="Arial.ttf", pil=False, example="abc"): + assert im.data.contiguous, "Image not contiguous. Apply np.ascontiguousarray(im) to Annotator() input images." self.pil = pil or not is_ascii(example) or is_chinese(example) if self.pil: # use PIL self.im = im if isinstance(im, Image.Image) else Image.fromarray(im) self.draw = ImageDraw.Draw(self.im) - self.font = check_font(font='Arial.Unicode.ttf' if is_chinese(example) else font, - size=font_size or max(round(sum(self.im.size) / 2 * 0.035), 12)) + self.font = check_font( + font="Arial.Unicode.ttf" if is_chinese(example) else font, + size=font_size or max(round(sum(self.im.size) / 2 * 0.035), 12), + ) else: # use cv2 self.im = im self.lw = line_width or max(round(sum(im.shape) / 2 * 0.003), 2) # line width - def box_label(self, box, label='', color=(128, 128, 128), txt_color=(255, 255, 255)): + def box_label(self, box, label="", color=(128, 128, 128), txt_color=(255, 255, 255)): # Add one xyxy box to image with label if self.pil or not is_ascii(label): self.draw.rectangle(box, width=self.lw, outline=color) # box if label: w, h = self.font.getsize(label) # text width, height outside = box[1] - h >= 0 # label fits outside box - self.draw.rectangle([box[0], - box[1] - h if outside else box[1], - box[0] + w + 1, - box[1] + 1 if outside else box[1] + h + 1], fill=color) + self.draw.rectangle( + [ + box[0], + box[1] - h if outside else box[1], + box[0] + w + 1, + box[1] + 1 if outside else box[1] + h + 1, + ], + fill=color, + ) # self.draw.text((box[0], box[1]), label, fill=txt_color, font=self.font, anchor='ls') # for PIL>8.0 self.draw.text((box[0], box[1] - h if outside else box[1]), label, fill=txt_color, font=self.font) else: # cv2 @@ -100,8 +127,16 @@ def box_label(self, box, label='', color=(128, 128, 128), txt_color=(255, 255, 2 outside = p1[1] - h - 3 >= 0 # label fits outside box p2 = p1[0] + w, p1[1] - h - 3 if outside else p1[1] + h + 3 cv2.rectangle(self.im, p1, p2, color, -1, cv2.LINE_AA) # filled - cv2.putText(self.im, label, (p1[0], p1[1] - 2 if outside else p1[1] + h + 2), 0, self.lw / 3, txt_color, - thickness=tf, lineType=cv2.LINE_AA) + cv2.putText( + self.im, + label, + (p1[0], p1[1] - 2 if outside else p1[1] + h + 2), + 0, + self.lw / 3, + txt_color, + thickness=tf, + lineType=cv2.LINE_AA, + ) def rectangle(self, xy, fill=None, outline=None, width=1): # Add rectangle to image (PIL-only) @@ -117,7 +152,7 @@ def result(self): return np.asarray(self.im) -def feature_visualization(x, module_type, stage, n=32, save_dir=Path('runs/detect/exp')): +def feature_visualization(x, module_type, stage, n=32, save_dir=Path("runs/detect/exp")): """ x: Features to be visualized module_type: Module type @@ -125,7 +160,7 @@ def feature_visualization(x, module_type, stage, n=32, save_dir=Path('runs/detec n: Maximum number of feature maps to plot save_dir: Directory to save results """ - if 'Detect' not in module_type: + if "Detect" not in module_type: batch, channels, height, width = x.shape # batch, channels, height, width if height > 1 and width > 1: f = f"stage{stage}_{module_type.split('.')[-1]}_features.png" # filename @@ -137,10 +172,10 @@ def feature_visualization(x, module_type, stage, n=32, save_dir=Path('runs/detec plt.subplots_adjust(wspace=0.05, hspace=0.05) for i in range(n): ax[i].imshow(blocks[i].squeeze()) # cmap='gray' - ax[i].axis('off') + ax[i].axis("off") - print(f'Saving {save_dir / f}... ({n}/{channels})') - plt.savefig(save_dir / f, dpi=300, bbox_inches='tight') + print(f"Saving {save_dir / f}... ({n}/{channels})") + plt.savefig(save_dir / f, dpi=300, bbox_inches="tight") plt.close() @@ -160,7 +195,7 @@ def butter_lowpass_filtfilt(data, cutoff=1500, fs=50000, order=5): def butter_lowpass(cutoff, fs, order): nyq = 0.5 * fs normal_cutoff = cutoff / nyq - return butter(order, normal_cutoff, btype='low', analog=False) + return butter(order, normal_cutoff, btype="low", analog=False) b, a = butter_lowpass(cutoff, fs, order=order) return filtfilt(b, a, data) # forward-backward filter @@ -175,7 +210,7 @@ def output_to_target(output): return np.array(targets) -def plot_images(images, targets, paths=None, fname='images.jpg', names=None, max_size=1920, max_subplots=16): +def plot_images(images, targets, paths=None, fname="images.jpg", names=None, max_size=1920, max_subplots=16): # Plot image grid with labels if isinstance(images, torch.Tensor): images = images.cpu().float().numpy() @@ -194,7 +229,7 @@ def plot_images(images, targets, paths=None, fname='images.jpg', names=None, max break x, y = int(w * (i // ns)), int(h * (i % ns)) # block origin im = im.transpose(1, 2, 0) - mosaic[y:y + h, x:x + w, :] = im + mosaic[y : y + h, x : x + w, :] = im # Resize (optional) scale = max_size / ns / max(h, w) @@ -214,7 +249,7 @@ def plot_images(images, targets, paths=None, fname='images.jpg', names=None, max if len(targets) > 0: ti = targets[targets[:, 0] == i] # image targets boxes = xywh2xyxy(ti[:, 2:6]).T - classes = ti[:, 1].astype('int') + classes = ti[:, 1].astype("int") labels = ti.shape[1] == 6 # labels if no conf column conf = None if labels else ti[:, 6] # check for confidence presence (label vs pred) @@ -231,59 +266,59 @@ def plot_images(images, targets, paths=None, fname='images.jpg', names=None, max color = colors(cls) cls = names[cls] if names else cls if labels or conf[j] > 0.25: # 0.25 conf thresh - label = f'{cls}' if labels else f'{cls} {conf[j]:.1f}' + label = f"{cls}" if labels else f"{cls} {conf[j]:.1f}" annotator.box_label(box, label, color=color) annotator.im.save(fname) # save -def plot_lr_scheduler(optimizer, scheduler, epochs=300, save_dir=''): +def plot_lr_scheduler(optimizer, scheduler, epochs=300, save_dir=""): # Plot LR simulating training for full epochs optimizer, scheduler = copy(optimizer), copy(scheduler) # do not modify originals y = [] for _ in range(epochs): scheduler.step() - y.append(optimizer.param_groups[0]['lr']) - plt.plot(y, '.-', label='LR') - plt.xlabel('epoch') - plt.ylabel('LR') + y.append(optimizer.param_groups[0]["lr"]) + plt.plot(y, ".-", label="LR") + plt.xlabel("epoch") + plt.ylabel("LR") plt.grid() plt.xlim(0, epochs) plt.ylim(0) - plt.savefig(Path(save_dir) / 'LR.png', dpi=200) + plt.savefig(Path(save_dir) / "LR.png", dpi=200) plt.close() def plot_val_txt(): # from utils.plots import *; plot_val() # Plot val.txt histograms - x = np.loadtxt('val.txt', dtype=np.float32) + x = np.loadtxt("val.txt", dtype=np.float32) box = xyxy2xywh(x[:, :4]) cx, cy = box[:, 0], box[:, 1] fig, ax = plt.subplots(1, 1, figsize=(6, 6), tight_layout=True) ax.hist2d(cx, cy, bins=600, cmax=10, cmin=0) - ax.set_aspect('equal') - plt.savefig('hist2d.png', dpi=300) + ax.set_aspect("equal") + plt.savefig("hist2d.png", dpi=300) fig, ax = plt.subplots(1, 2, figsize=(12, 6), tight_layout=True) ax[0].hist(cx, bins=600) ax[1].hist(cy, bins=600) - plt.savefig('hist1d.png', dpi=200) + plt.savefig("hist1d.png", dpi=200) def plot_targets_txt(): # from utils.plots import *; plot_targets_txt() # Plot targets.txt histograms - x = np.loadtxt('targets.txt', dtype=np.float32).T - s = ['x targets', 'y targets', 'width targets', 'height targets'] + x = np.loadtxt("targets.txt", dtype=np.float32).T + s = ["x targets", "y targets", "width targets", "height targets"] fig, ax = plt.subplots(2, 2, figsize=(8, 8), tight_layout=True) ax = ax.ravel() for i in range(4): - ax[i].hist(x[i], bins=100, label=f'{x[i].mean():.3g} +/- {x[i].std():.3g}') + ax[i].hist(x[i], bins=100, label=f"{x[i].mean():.3g} +/- {x[i].std():.3g}") ax[i].legend() ax[i].set_title(s[i]) - plt.savefig('targets.jpg', dpi=200) + plt.savefig("targets.jpg", dpi=200) -def plot_val_study(file='', dir='', x=None): # from utils.plots import *; plot_val_study() +def plot_val_study(file="", dir="", x=None): # from utils.plots import *; plot_val_study() # Plot file=study.txt generated by val.py (or plot all study*.txt in dir) save_dir = Path(file).parent if file else Path(dir) plot2 = False # plot additional results @@ -292,59 +327,72 @@ def plot_val_study(file='', dir='', x=None): # from utils.plots import *; plot_ fig2, ax2 = plt.subplots(1, 1, figsize=(8, 4), tight_layout=True) # for f in [save_dir / f'study_coco_{x}.txt' for x in ['yolov5n6', 'yolov5s6', 'yolov5m6', 'yolov5l6', 'yolov5x6']]: - for f in sorted(save_dir.glob('study*.txt')): + for f in sorted(save_dir.glob("study*.txt")): y = np.loadtxt(f, dtype=np.float32, usecols=[0, 1, 2, 3, 7, 8, 9], ndmin=2).T x = np.arange(y.shape[1]) if x is None else np.array(x) if plot2: - s = ['P', 'R', 'mAP@.5', 'mAP@.5:.95', 't_preprocess (ms/img)', 't_inference (ms/img)', 't_NMS (ms/img)'] + s = ["P", "R", "mAP@.5", "mAP@.5:.95", "t_preprocess (ms/img)", "t_inference (ms/img)", "t_NMS (ms/img)"] for i in range(7): - ax[i].plot(x, y[i], '.-', linewidth=2, markersize=8) + ax[i].plot(x, y[i], ".-", linewidth=2, markersize=8) ax[i].set_title(s[i]) j = y[3].argmax() + 1 - ax2.plot(y[5, 1:j], y[3, 1:j] * 1E2, '.-', linewidth=2, markersize=8, - label=f.stem.replace('study_coco_', '').replace('yolo', 'YOLO')) - - ax2.plot(1E3 / np.array([209, 140, 97, 58, 35, 18]), [34.6, 40.5, 43.0, 47.5, 49.7, 51.5], - 'k.-', linewidth=2, markersize=8, alpha=.25, label='EfficientDet') + ax2.plot( + y[5, 1:j], + y[3, 1:j] * 1e2, + ".-", + linewidth=2, + markersize=8, + label=f.stem.replace("study_coco_", "").replace("yolo", "YOLO"), + ) + + ax2.plot( + 1e3 / np.array([209, 140, 97, 58, 35, 18]), + [34.6, 40.5, 43.0, 47.5, 49.7, 51.5], + "k.-", + linewidth=2, + markersize=8, + alpha=0.25, + label="EfficientDet", + ) ax2.grid(alpha=0.2) ax2.set_yticks(np.arange(20, 60, 5)) ax2.set_xlim(0, 57) ax2.set_ylim(25, 55) - ax2.set_xlabel('GPU Speed (ms/img)') - ax2.set_ylabel('COCO AP val') - ax2.legend(loc='lower right') - f = save_dir / 'study.png' - print(f'Saving {f}...') + ax2.set_xlabel("GPU Speed (ms/img)") + ax2.set_ylabel("COCO AP val") + ax2.legend(loc="lower right") + f = save_dir / "study.png" + print(f"Saving {f}...") plt.savefig(f, dpi=300) -def plot_labels(labels, names=(), save_dir=Path('')): +def plot_labels(labels, names=(), save_dir=Path("")): # plot dataset labels - print('Plotting labels... ') + print("Plotting labels... ") c, b = labels[:, 0], labels[:, 1:].transpose() # classes, boxes nc = int(c.max() + 1) # number of classes - x = pd.DataFrame(b.transpose(), columns=['x', 'y', 'width', 'height']) + x = pd.DataFrame(b.transpose(), columns=["x", "y", "width", "height"]) # seaborn correlogram - sn.pairplot(x, corner=True, diag_kind='auto', kind='hist', diag_kws=dict(bins=50), plot_kws=dict(pmax=0.9)) - plt.savefig(save_dir / 'labels_correlogram.jpg', dpi=200) + sn.pairplot(x, corner=True, diag_kind="auto", kind="hist", diag_kws=dict(bins=50), plot_kws=dict(pmax=0.9)) + plt.savefig(save_dir / "labels_correlogram.jpg", dpi=200) plt.close() # matplotlib labels - matplotlib.use('svg') # faster + matplotlib.use("svg") # faster ax = plt.subplots(2, 2, figsize=(8, 8), tight_layout=True)[1].ravel() y = ax[0].hist(c, bins=np.linspace(0, nc, nc + 1) - 0.5, rwidth=0.8) # [y[2].patches[i].set_color([x / 255 for x in colors(i)]) for i in range(nc)] # update colors bug #3195 - ax[0].set_ylabel('instances') + ax[0].set_ylabel("instances") if 0 < len(names) < 30: ax[0].set_xticks(range(len(names))) ax[0].set_xticklabels(names, rotation=90, fontsize=10) else: - ax[0].set_xlabel('classes') - sn.histplot(x, x='x', y='y', ax=ax[2], bins=50, pmax=0.9) - sn.histplot(x, x='width', y='height', ax=ax[3], bins=50, pmax=0.9) + ax[0].set_xlabel("classes") + sn.histplot(x, x="x", y="y", ax=ax[2], bins=50, pmax=0.9) + sn.histplot(x, x="width", y="height", ax=ax[3], bins=50, pmax=0.9) # rectangles labels[:, 1:3] = 0.5 # center @@ -353,18 +401,18 @@ def plot_labels(labels, names=(), save_dir=Path('')): for cls, *box in labels[:1000]: ImageDraw.Draw(img).rectangle(box, width=1, outline=colors(cls)) # plot ax[1].imshow(img) - ax[1].axis('off') + ax[1].axis("off") for a in [0, 1, 2, 3]: - for s in ['top', 'right', 'left', 'bottom']: + for s in ["top", "right", "left", "bottom"]: ax[a].spines[s].set_visible(False) - plt.savefig(save_dir / 'labels.jpg', dpi=200) - matplotlib.use('Agg') + plt.savefig(save_dir / "labels.jpg", dpi=200) + matplotlib.use("Agg") plt.close() -def plot_evolve(evolve_csv='path/to/evolve.csv'): # from utils.plots import *; plot_evolve() +def plot_evolve(evolve_csv="path/to/evolve.csv"): # from utils.plots import *; plot_evolve() # Plot evolve.csv hyp evolution results evolve_csv = Path(evolve_csv) data = pd.read_csv(evolve_csv) @@ -373,30 +421,30 @@ def plot_evolve(evolve_csv='path/to/evolve.csv'): # from utils.plots import *; f = fitness(x) j = np.argmax(f) # max fitness index plt.figure(figsize=(10, 12), tight_layout=True) - matplotlib.rc('font', **{'size': 8}) + matplotlib.rc("font", **{"size": 8}) for i, k in enumerate(keys[7:]): v = x[:, 7 + i] mu = v[j] # best single result plt.subplot(6, 5, i + 1) - plt.scatter(v, f, c=hist2d(v, f, 20), cmap='viridis', alpha=.8, edgecolors='none') - plt.plot(mu, f.max(), 'k+', markersize=15) - plt.title(f'{k} = {mu:.3g}', fontdict={'size': 9}) # limit to 40 characters + plt.scatter(v, f, c=hist2d(v, f, 20), cmap="viridis", alpha=0.8, edgecolors="none") + plt.plot(mu, f.max(), "k+", markersize=15) + plt.title(f"{k} = {mu:.3g}", fontdict={"size": 9}) # limit to 40 characters if i % 5 != 0: plt.yticks([]) - print(f'{k:>15}: {mu:.3g}') - f = evolve_csv.with_suffix('.png') # filename + print(f"{k:>15}: {mu:.3g}") + f = evolve_csv.with_suffix(".png") # filename plt.savefig(f, dpi=200) plt.close() - print(f'Saved {f}') + print(f"Saved {f}") -def plot_results(file='path/to/results.csv', dir=''): +def plot_results(file="path/to/results.csv", dir=""): # Plot training results.csv. Usage: from utils.plots import *; plot_results('path/to/results.csv') save_dir = Path(file).parent if file else Path(dir) fig, ax = plt.subplots(2, 5, figsize=(12, 6), tight_layout=True) ax = ax.ravel() - files = list(save_dir.glob('results*.csv')) - assert len(files), f'No results.csv files found in {save_dir.resolve()}, nothing to plot.' + files = list(save_dir.glob("results*.csv")) + assert len(files), f"No results.csv files found in {save_dir.resolve()}, nothing to plot." for fi, f in enumerate(files): try: data = pd.read_csv(f) @@ -405,49 +453,49 @@ def plot_results(file='path/to/results.csv', dir=''): for i, j in enumerate([1, 2, 3, 4, 5, 8, 9, 10, 6, 7]): y = data.values[:, j] # y[y == 0] = np.nan # don't show zero values - ax[i].plot(x, y, marker='.', label=f.stem, linewidth=2, markersize=8) + ax[i].plot(x, y, marker=".", label=f.stem, linewidth=2, markersize=8) ax[i].set_title(s[j], fontsize=12) # if j in [8, 9, 10]: # share train and val loss y axes # ax[i].get_shared_y_axes().join(ax[i], ax[i - 5]) except Exception as e: - print(f'Warning: Plotting error for {f}: {e}') + print(f"Warning: Plotting error for {f}: {e}") ax[1].legend() - fig.savefig(save_dir / 'results.png', dpi=200) + fig.savefig(save_dir / "results.png", dpi=200) plt.close() -def profile_idetection(start=0, stop=0, labels=(), save_dir=''): +def profile_idetection(start=0, stop=0, labels=(), save_dir=""): # Plot iDetection '*.txt' per-image logs. from utils.plots import *; profile_idetection() ax = plt.subplots(2, 4, figsize=(12, 6), tight_layout=True)[1].ravel() - s = ['Images', 'Free Storage (GB)', 'RAM Usage (GB)', 'Battery', 'dt_raw (ms)', 'dt_smooth (ms)', 'real-world FPS'] - files = list(Path(save_dir).glob('frames*.txt')) + s = ["Images", "Free Storage (GB)", "RAM Usage (GB)", "Battery", "dt_raw (ms)", "dt_smooth (ms)", "real-world FPS"] + files = list(Path(save_dir).glob("frames*.txt")) for fi, f in enumerate(files): try: results = np.loadtxt(f, ndmin=2).T[:, 90:-30] # clip first and last rows n = results.shape[1] # number of rows x = np.arange(start, min(stop, n) if stop else n) results = results[:, x] - t = (results[0] - results[0].min()) # set t0=0s + t = results[0] - results[0].min() # set t0=0s results[0] = x for i, a in enumerate(ax): if i < len(results): - label = labels[fi] if len(labels) else f.stem.replace('frames_', '') - a.plot(t, results[i], marker='.', label=label, linewidth=1, markersize=5) + label = labels[fi] if len(labels) else f.stem.replace("frames_", "") + a.plot(t, results[i], marker=".", label=label, linewidth=1, markersize=5) a.set_title(s[i]) - a.set_xlabel('time (s)') + a.set_xlabel("time (s)") # if fi == len(files) - 1: # a.set_ylim(bottom=0) - for side in ['top', 'right']: + for side in ["top", "right"]: a.spines[side].set_visible(False) else: a.remove() except Exception as e: - print(f'Warning: Plotting error for {f}; {e}') + print(f"Warning: Plotting error for {f}; {e}") ax[1].legend() - plt.savefig(Path(save_dir) / 'idetection_profile.png', dpi=200) + plt.savefig(Path(save_dir) / "idetection_profile.png", dpi=200) -def save_one_box(xyxy, im, file='image.jpg', gain=1.02, pad=10, square=False, BGR=False, save=True): +def save_one_box(xyxy, im, file="image.jpg", gain=1.02, pad=10, square=False, BGR=False, save=True): # Save image crop as {file} with crop size multiple {gain} and {pad} pixels. Save and/or return crop xyxy = torch.tensor(xyxy).view(-1, 4) b = xyxy2xywh(xyxy) # boxes @@ -456,8 +504,8 @@ def save_one_box(xyxy, im, file='image.jpg', gain=1.02, pad=10, square=False, BG b[:, 2:] = b[:, 2:] * gain + pad # box wh * gain + pad xyxy = xywh2xyxy(b).long() clip_coords(xyxy, im.shape) - crop = im[int(xyxy[0, 1]):int(xyxy[0, 3]), int(xyxy[0, 0]):int(xyxy[0, 2]), ::(1 if BGR else -1)] + crop = im[int(xyxy[0, 1]) : int(xyxy[0, 3]), int(xyxy[0, 0]) : int(xyxy[0, 2]), :: (1 if BGR else -1)] if save: file.parent.mkdir(parents=True, exist_ok=True) # make directory - cv2.imwrite(str(increment_path(file).with_suffix('.jpg')), crop) + cv2.imwrite(str(increment_path(file).with_suffix(".jpg")), crop) return crop diff --git a/utils/torch_utils.py b/utils/torch_utils.py index 73acec8e819c..d06b73b2332b 100644 --- a/utils/torch_utils.py +++ b/utils/torch_utils.py @@ -42,44 +42,44 @@ def torch_distributed_zero_first(local_rank: int): def date_modified(path=__file__): # return human-readable file modification date, i.e. '2021-3-26' t = datetime.datetime.fromtimestamp(Path(path).stat().st_mtime) - return f'{t.year}-{t.month}-{t.day}' + return f"{t.year}-{t.month}-{t.day}" def git_describe(path=Path(__file__).parent): # path must be a directory # return human-readable git description, i.e. v5.0-5-g3e25f1e https://git-scm.com/docs/git-describe - s = f'git -C {path} describe --tags --long --always' + s = f"git -C {path} describe --tags --long --always" try: return subprocess.check_output(s, shell=True, stderr=subprocess.STDOUT).decode()[:-1] except subprocess.CalledProcessError as e: - return '' # not a git repository + return "" # not a git repository -def select_device(device='', batch_size=None): +def select_device(device="", batch_size=None): # device = 'cpu' or '0' or '0,1,2,3' - s = f'YOLOv5 🚀 {git_describe() or date_modified()} torch {torch.__version__} ' # string - device = str(device).strip().lower().replace('cuda:', '') # to string, 'cuda:0' to '0' - cpu = device == 'cpu' + s = f"YOLOv5 🚀 {git_describe() or date_modified()} torch {torch.__version__} " # string + device = str(device).strip().lower().replace("cuda:", "") # to string, 'cuda:0' to '0' + cpu = device == "cpu" if cpu: - os.environ['CUDA_VISIBLE_DEVICES'] = '-1' # force torch.cuda.is_available() = False + os.environ["CUDA_VISIBLE_DEVICES"] = "-1" # force torch.cuda.is_available() = False elif device: # non-cpu device requested - os.environ['CUDA_VISIBLE_DEVICES'] = device # set environment variable - assert torch.cuda.is_available(), f'CUDA unavailable, invalid device {device} requested' # check availability + os.environ["CUDA_VISIBLE_DEVICES"] = device # set environment variable + assert torch.cuda.is_available(), f"CUDA unavailable, invalid device {device} requested" # check availability cuda = not cpu and torch.cuda.is_available() if cuda: - devices = device.split(',') if device else '0' # range(torch.cuda.device_count()) # i.e. 0,1,6,7 + devices = device.split(",") if device else "0" # range(torch.cuda.device_count()) # i.e. 0,1,6,7 n = len(devices) # device count if n > 1 and batch_size: # check batch_size is divisible by device_count - assert batch_size % n == 0, f'batch-size {batch_size} not multiple of GPU count {n}' - space = ' ' * (len(s) + 1) + assert batch_size % n == 0, f"batch-size {batch_size} not multiple of GPU count {n}" + space = " " * (len(s) + 1) for i, d in enumerate(devices): p = torch.cuda.get_device_properties(i) s += f"{'' if i == 0 else space}CUDA:{d} ({p.name}, {p.total_memory / 1024 ** 2:.0f}MiB)\n" # bytes to MB else: - s += 'CPU\n' + s += "CPU\n" - LOGGER.info(s.encode().decode('ascii', 'ignore') if platform.system() == 'Windows' else s) # emoji-safe - return torch.device('cuda:0' if cuda else 'cpu') + LOGGER.info(s.encode().decode("ascii", "ignore") if platform.system() == "Windows" else s) # emoji-safe + return torch.device("cuda:0" if cuda else "cpu") def time_sync(): @@ -100,18 +100,20 @@ def profile(input, ops, n=10, device=None): results = [] device = device or select_device() - print(f"{'Params':>12s}{'GFLOPs':>12s}{'GPU_mem (GB)':>14s}{'forward (ms)':>14s}{'backward (ms)':>14s}" - f"{'input':>24s}{'output':>24s}") + print( + f"{'Params':>12s}{'GFLOPs':>12s}{'GPU_mem (GB)':>14s}{'forward (ms)':>14s}{'backward (ms)':>14s}" + f"{'input':>24s}{'output':>24s}" + ) for x in input if isinstance(input, list) else [input]: x = x.to(device) x.requires_grad = True for m in ops if isinstance(ops, list) else [ops]: - m = m.to(device) if hasattr(m, 'to') else m # device - m = m.half() if hasattr(m, 'half') and isinstance(x, torch.Tensor) and x.dtype is torch.float16 else m + m = m.to(device) if hasattr(m, "to") else m # device + m = m.half() if hasattr(m, "half") and isinstance(x, torch.Tensor) and x.dtype is torch.float16 else m tf, tb, t = 0, 0, [0, 0, 0] # dt forward, backward try: - flops = thop.profile(m, inputs=(x,), verbose=False)[0] / 1E9 * 2 # GFLOPs + flops = thop.profile(m, inputs=(x,), verbose=False)[0] / 1e9 * 2 # GFLOPs except: flops = 0 @@ -125,14 +127,14 @@ def profile(input, ops, n=10, device=None): t[2] = time_sync() except Exception as e: # no backward method # print(e) # for debug - t[2] = float('nan') + t[2] = float("nan") tf += (t[1] - t[0]) * 1000 / n # ms per op forward tb += (t[2] - t[1]) * 1000 / n # ms per op backward - mem = torch.cuda.memory_reserved() / 1E9 if torch.cuda.is_available() else 0 # (GB) - s_in = tuple(x.shape) if isinstance(x, torch.Tensor) else 'list' - s_out = tuple(y.shape) if isinstance(y, torch.Tensor) else 'list' + mem = torch.cuda.memory_reserved() / 1e9 if torch.cuda.is_available() else 0 # (GB) + s_in = tuple(x.shape) if isinstance(x, torch.Tensor) else "list" + s_out = tuple(y.shape) if isinstance(y, torch.Tensor) else "list" p = sum(list(x.numel() for x in m.parameters())) if isinstance(m, nn.Module) else 0 # parameters - print(f'{p:12}{flops:12.4g}{mem:>14.3f}{tf:14.4g}{tb:14.4g}{str(s_in):>24s}{str(s_out):>24s}') + print(f"{p:12}{flops:12.4g}{mem:>14.3f}{tf:14.4g}{tb:14.4g}{str(s_in):>24s}{str(s_out):>24s}") results.append([p, flops, mem, tf, tb, s_in, s_out]) except Exception as e: print(e) @@ -180,23 +182,30 @@ def sparsity(model): def prune(model, amount=0.3): # Prune model to requested global sparsity import torch.nn.utils.prune as prune - print('Pruning model... ', end='') + + print("Pruning model... ", end="") for name, m in model.named_modules(): if isinstance(m, nn.Conv2d): - prune.l1_unstructured(m, name='weight', amount=amount) # prune - prune.remove(m, 'weight') # make permanent - print(' %.3g global sparsity' % sparsity(model)) + prune.l1_unstructured(m, name="weight", amount=amount) # prune + prune.remove(m, "weight") # make permanent + print(" %.3g global sparsity" % sparsity(model)) def fuse_conv_and_bn(conv, bn): # Fuse convolution and batchnorm layers https://tehnokv.com/posts/fusing-batchnorm-and-conv/ - fusedconv = nn.Conv2d(conv.in_channels, - conv.out_channels, - kernel_size=conv.kernel_size, - stride=conv.stride, - padding=conv.padding, - groups=conv.groups, - bias=True).requires_grad_(False).to(conv.weight.device) + fusedconv = ( + nn.Conv2d( + conv.in_channels, + conv.out_channels, + kernel_size=conv.kernel_size, + stride=conv.stride, + padding=conv.padding, + groups=conv.groups, + bias=True, + ) + .requires_grad_(False) + .to(conv.weight.device) + ) # prepare filters w_conv = conv.weight.clone().view(conv.out_channels, -1) @@ -218,24 +227,27 @@ def model_info(model, verbose=False, img_size=640): if verbose: print(f"{'layer':>5} {'name':>40} {'gradient':>9} {'parameters':>12} {'shape':>20} {'mu':>10} {'sigma':>10}") for i, (name, p) in enumerate(model.named_parameters()): - name = name.replace('module_list.', '') - print('%5g %40s %9s %12g %20s %10.3g %10.3g' % - (i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std())) + name = name.replace("module_list.", "") + print( + "%5g %40s %9s %12g %20s %10.3g %10.3g" + % (i, name, p.requires_grad, p.numel(), list(p.shape), p.mean(), p.std()) + ) try: # FLOPs from thop import profile - stride = max(int(model.stride.max()), 32) if hasattr(model, 'stride') else 32 - img = torch.zeros((1, model.yaml.get('ch', 3), stride, stride), device=next(model.parameters()).device) # input - flops = profile(deepcopy(model), inputs=(img,), verbose=False)[0] / 1E9 * 2 # stride GFLOPs + + stride = max(int(model.stride.max()), 32) if hasattr(model, "stride") else 32 + img = torch.zeros((1, model.yaml.get("ch", 3), stride, stride), device=next(model.parameters()).device) # input + flops = profile(deepcopy(model), inputs=(img,), verbose=False)[0] / 1e9 * 2 # stride GFLOPs img_size = img_size if isinstance(img_size, list) else [img_size, img_size] # expand if int/float - fs = ', %.1f GFLOPs' % (flops * img_size[0] / stride * img_size[1] / stride) # 640x640 GFLOPs + fs = ", %.1f GFLOPs" % (flops * img_size[0] / stride * img_size[1] / stride) # 640x640 GFLOPs except (ImportError, Exception): - fs = '' + fs = "" LOGGER.info(f"Model Summary: {len(list(model.modules()))} layers, {n_p} parameters, {n_g} gradients{fs}") -def load_classifier(name='resnet101', n=2): +def load_classifier(name="resnet101", n=2): # Loads a pretrained model reshaped to n-class output model = torchvision.models.__dict__[name](pretrained=True) @@ -261,7 +273,7 @@ def scale_img(img, ratio=1.0, same_shape=False, gs=32): # img(16,3,256,416) else: h, w = img.shape[2:] s = (int(h * ratio), int(w * ratio)) # new size - img = F.interpolate(img, size=s, mode='bilinear', align_corners=False) # resize + img = F.interpolate(img, size=s, mode="bilinear", align_corners=False) # resize if not same_shape: # pad/crop img h, w = (math.ceil(x * ratio / gs) * gs for x in (h, w)) return F.pad(img, [0, w - s[1], 0, h - s[0]], value=0.447) # value = imagenet mean @@ -270,7 +282,7 @@ def scale_img(img, ratio=1.0, same_shape=False, gs=32): # img(16,3,256,416) def copy_attr(a, b, include=(), exclude=()): # Copy attributes from b to a, options to only include [...] and to exclude [...] for k, v in b.__dict__.items(): - if (len(include) and k not in include) or k.startswith('_') or k in exclude: + if (len(include) and k not in include) or k.startswith("_") or k in exclude: continue else: setattr(a, k, v) @@ -281,7 +293,7 @@ class EarlyStopping: def __init__(self, patience=30): self.best_fitness = 0.0 # i.e. mAP self.best_epoch = 0 - self.patience = patience or float('inf') # epochs to wait after fitness stops improving to stop + self.patience = patience or float("inf") # epochs to wait after fitness stops improving to stop self.possible_stop = False # possible stop may occur next epoch def __call__(self, epoch, fitness): @@ -292,15 +304,17 @@ def __call__(self, epoch, fitness): self.possible_stop = delta >= (self.patience - 1) # possible stop may occur next epoch stop = delta >= self.patience # stop training if patience exceeded if stop: - LOGGER.info(f'Stopping training early as no improvement observed in last {self.patience} epochs. ' - f'Best results observed at epoch {self.best_epoch}, best model saved as best.pt.\n' - f'To update EarlyStopping(patience={self.patience}) pass a new patience value, ' - f'i.e. `python train.py --patience 300` or use `--patience 0` to disable EarlyStopping.') + LOGGER.info( + f"Stopping training early as no improvement observed in last {self.patience} epochs. " + f"Best results observed at epoch {self.best_epoch}, best model saved as best.pt.\n" + f"To update EarlyStopping(patience={self.patience}) pass a new patience value, " + f"i.e. `python train.py --patience 300` or use `--patience 0` to disable EarlyStopping." + ) return stop class ModelEMA: - """ Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models + """Model Exponential Moving Average from https://github.com/rwightman/pytorch-image-models Keep a moving average of everything in the model state_dict (parameters and buffers). This is intended to allow functionality like https://www.tensorflow.org/api_docs/python/tf/train/ExponentialMovingAverage @@ -331,6 +345,6 @@ def update(self, model): v *= d v += (1 - d) * msd[k].detach() - def update_attr(self, model, include=(), exclude=('process_group', 'reducer')): + def update_attr(self, model, include=(), exclude=("process_group", "reducer")): # Update EMA attributes copy_attr(self.ema, model, include, exclude) diff --git a/val.py b/val.py index d2797f1189ec..f9c267e20ec2 100644 --- a/val.py +++ b/val.py @@ -26,9 +26,23 @@ from models.experimental import attempt_load from utils.callbacks import Callbacks from utils.datasets import create_dataloader -from utils.general import (LOGGER, box_iou, check_dataset, check_img_size, check_requirements, check_suffix, check_yaml, - coco80_to_coco91_class, colorstr, increment_path, non_max_suppression, print_args, - scale_coords, xywh2xyxy, xyxy2xywh) +from utils.general import ( + LOGGER, + box_iou, + check_dataset, + check_img_size, + check_requirements, + check_suffix, + check_yaml, + coco80_to_coco91_class, + colorstr, + increment_path, + non_max_suppression, + print_args, + scale_coords, + xywh2xyxy, + xyxy2xywh, +) from utils.metrics import ConfusionMatrix, ap_per_class from utils.plots import output_to_target, plot_images, plot_val_study from utils.torch_utils import select_device, time_sync @@ -40,8 +54,8 @@ def save_one_txt(predn, save_conf, shape, file): for *xyxy, conf, cls in predn.tolist(): xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh line = (cls, *xywh, conf) if save_conf else (cls, *xywh) # label format - with open(file, 'a') as f: - f.write(('%g ' * len(line)).rstrip() % line + '\n') + with open(file, "a") as f: + f.write(("%g " * len(line)).rstrip() % line + "\n") def save_one_json(predn, jdict, path, class_map): @@ -50,10 +64,14 @@ def save_one_json(predn, jdict, path, class_map): box = xyxy2xywh(predn[:, :4]) # xywh box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner for p, b in zip(predn.tolist(), box.tolist()): - jdict.append({'image_id': image_id, - 'category_id': class_map[int(p[5])], - 'bbox': [round(x, 3) for x in b], - 'score': round(p[4], 5)}) + jdict.append( + { + "image_id": image_id, + "category_id": class_map[int(p[5])], + "bbox": [round(x, 3) for x in b], + "score": round(p[4], 5), + } + ) def process_batch(detections, labels, iouv): @@ -81,32 +99,33 @@ def process_batch(detections, labels, iouv): @torch.no_grad() -def run(data, - weights=None, # model.pt path(s) - batch_size=32, # batch size - imgsz=640, # inference size (pixels) - conf_thres=0.001, # confidence threshold - iou_thres=0.6, # NMS IoU threshold - task='val', # train, val, test, speed or study - device='', # cuda device, i.e. 0 or 0,1,2,3 or cpu - single_cls=False, # treat as single-class dataset - augment=False, # augmented inference - verbose=False, # verbose output - save_txt=False, # save results to *.txt - save_hybrid=False, # save label+prediction hybrid results to *.txt - save_conf=False, # save confidences in --save-txt labels - save_json=False, # save a COCO-JSON results file - project=ROOT / 'runs/val', # save to project/name - name='exp', # save to project/name - exist_ok=False, # existing project/name ok, do not increment - half=True, # use FP16 half-precision inference - model=None, - dataloader=None, - save_dir=Path(''), - plots=True, - callbacks=Callbacks(), - compute_loss=None, - ): +def run( + data, + weights=None, # model.pt path(s) + batch_size=32, # batch size + imgsz=640, # inference size (pixels) + conf_thres=0.001, # confidence threshold + iou_thres=0.6, # NMS IoU threshold + task="val", # train, val, test, speed or study + device="", # cuda device, i.e. 0 or 0,1,2,3 or cpu + single_cls=False, # treat as single-class dataset + augment=False, # augmented inference + verbose=False, # verbose output + save_txt=False, # save results to *.txt + save_hybrid=False, # save label+prediction hybrid results to *.txt + save_conf=False, # save confidences in --save-txt labels + save_json=False, # save a COCO-JSON results file + project=ROOT / "runs/val", # save to project/name + name="exp", # save to project/name + exist_ok=False, # existing project/name ok, do not increment + half=True, # use FP16 half-precision inference + model=None, + dataloader=None, + save_dir=Path(""), + plots=True, + callbacks=Callbacks(), + compute_loss=None, +): # Initialize/load model and set device training = model is not None if training: # called by train.py @@ -117,10 +136,10 @@ def run(data, # Directories save_dir = increment_path(Path(project) / name, exist_ok=exist_ok) # increment run - (save_dir / 'labels' if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir + (save_dir / "labels" if save_txt else save_dir).mkdir(parents=True, exist_ok=True) # make dir # Load model - check_suffix(weights, '.pt') + check_suffix(weights, ".pt") model = attempt_load(weights, map_location=device) # load FP32 model gs = max(int(model.stride.max()), 32) # grid size (max stride) imgsz = check_img_size(imgsz, s=gs) # check image size @@ -133,30 +152,31 @@ def run(data, data = check_dataset(data) # check # Half - half &= device.type != 'cpu' # half precision only supported on CUDA + half &= device.type != "cpu" # half precision only supported on CUDA model.half() if half else model.float() # Configure model.eval() - is_coco = isinstance(data.get('val'), str) and data['val'].endswith('coco/val2017.txt') # COCO dataset - nc = 1 if single_cls else int(data['nc']) # number of classes + is_coco = isinstance(data.get("val"), str) and data["val"].endswith("coco/val2017.txt") # COCO dataset + 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 mAP@0.5:0.95 niou = iouv.numel() # Dataloader if not training: - if device.type != 'cpu': + if device.type != "cpu": model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters()))) # run once - pad = 0.0 if task == 'speed' else 0.5 - task = task if task in ('train', 'val', 'test') else 'val' # path to train/val/test images - dataloader = create_dataloader(data[task], imgsz, batch_size, gs, single_cls, pad=pad, rect=True, - prefix=colorstr(f'{task}: '))[0] + pad = 0.0 if task == "speed" else 0.5 + task = task if task in ("train", "val", "test") else "val" # path to train/val/test images + dataloader = create_dataloader( + data[task], imgsz, batch_size, gs, single_cls, pad=pad, rect=True, prefix=colorstr(f"{task}: ") + )[0] seen = 0 confusion_matrix = ConfusionMatrix(nc=nc) - names = {k: v for k, v in enumerate(model.names if hasattr(model, 'names') else model.module.names)} + names = {k: v for k, v in enumerate(model.names if hasattr(model, "names") else model.module.names)} class_map = coco80_to_coco91_class() if is_coco else list(range(1000)) - s = ('%20s' + '%11s' * 6) % ('Class', 'Images', 'Labels', 'P', 'R', 'mAP@.5', 'mAP@.5:.95') + s = ("%20s" + "%11s" * 6) % ("Class", "Images", "Labels", "P", "R", "mAP@.5", "mAP@.5:.95") dt, p, r, f1, mp, mr, map50, map = [0.0, 0.0, 0.0], 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 loss = torch.zeros(3, device=device) jdict, stats, ap, ap_class = [], [], [], [] @@ -218,16 +238,16 @@ def run(data, # Save/log if save_txt: - save_one_txt(predn, save_conf, shape, file=save_dir / 'labels' / (path.stem + '.txt')) + save_one_txt(predn, save_conf, shape, file=save_dir / "labels" / (path.stem + ".txt")) if save_json: save_one_json(predn, jdict, path, class_map) # append to COCO-JSON dictionary - callbacks.run('on_val_image_end', pred, predn, path, names, img[si]) + callbacks.run("on_val_image_end", pred, predn, path, names, img[si]) # Plot images if plots and batch_i < 3: - f = save_dir / f'val_batch{batch_i}_labels.jpg' # labels + f = save_dir / f"val_batch{batch_i}_labels.jpg" # labels Thread(target=plot_images, args=(img, targets, paths, f, names), daemon=True).start() - f = save_dir / f'val_batch{batch_i}_pred.jpg' # predictions + f = save_dir / f"val_batch{batch_i}_pred.jpg" # predictions Thread(target=plot_images, args=(img, output_to_target(out), paths, f, names), daemon=True).start() # Compute statistics @@ -241,8 +261,8 @@ def run(data, nt = torch.zeros(1) # Print results - pf = '%20s' + '%11i' * 2 + '%11.3g' * 4 # print format - LOGGER.info(pf % ('all', seen, nt.sum(), mp, mr, map50, map)) + pf = "%20s" + "%11i" * 2 + "%11.3g" * 4 # print format + LOGGER.info(pf % ("all", seen, nt.sum(), mp, mr, map50, map)) # Print results per class if (verbose or (nc < 50 and not training)) and nc > 1 and len(stats): @@ -250,33 +270,33 @@ def run(data, LOGGER.info(pf % (names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i])) # Print speeds - t = tuple(x / seen * 1E3 for x in dt) # speeds per image + t = tuple(x / seen * 1e3 for x in dt) # speeds per image if not training: shape = (batch_size, 3, imgsz, imgsz) - LOGGER.info(f'Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {shape}' % t) + LOGGER.info(f"Speed: %.1fms pre-process, %.1fms inference, %.1fms NMS per image at shape {shape}" % t) # Plots if plots: confusion_matrix.plot(save_dir=save_dir, names=list(names.values())) - callbacks.run('on_val_end') + callbacks.run("on_val_end") # Save JSON if save_json and len(jdict): - w = Path(weights[0] if isinstance(weights, list) else weights).stem if weights is not None else '' # weights - anno_json = str(Path(data.get('path', '../coco')) / 'annotations/instances_val2017.json') # annotations json + w = Path(weights[0] if isinstance(weights, list) else weights).stem if weights is not None else "" # weights + anno_json = str(Path(data.get("path", "../coco")) / "annotations/instances_val2017.json") # annotations json pred_json = str(save_dir / f"{w}_predictions.json") # predictions json - LOGGER.info(f'\nEvaluating pycocotools mAP... saving {pred_json}...') - with open(pred_json, 'w') as f: + LOGGER.info(f"\nEvaluating pycocotools mAP... saving {pred_json}...") + with open(pred_json, "w") as f: json.dump(jdict, f) try: # https://github.com/cocodataset/cocoapi/blob/master/PythonAPI/pycocoEvalDemo.ipynb - check_requirements(['pycocotools']) + check_requirements(["pycocotools"]) from pycocotools.coco import COCO from pycocotools.cocoeval import COCOeval anno = COCO(anno_json) # init annotations api pred = anno.loadRes(pred_json) # init predictions api - eval = COCOeval(anno, pred, 'bbox') + eval = COCOeval(anno, pred, "bbox") if is_coco: eval.params.imgIds = [int(Path(x).stem) for x in dataloader.dataset.img_files] # image IDs to evaluate eval.evaluate() @@ -284,12 +304,12 @@ def run(data, eval.summarize() map, map50 = eval.stats[:2] # update results (mAP@0.5:0.95, mAP@0.5) except Exception as e: - LOGGER.info(f'pycocotools unable to run: {e}') + LOGGER.info(f"pycocotools unable to run: {e}") # Return results model.float() # for training if not training: - s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else '' + s = f"\n{len(list(save_dir.glob('labels/*.txt')))} labels saved to {save_dir / 'labels'}" if save_txt else "" LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}") maps = np.zeros(nc) + map for i, c in enumerate(ap_class): @@ -299,60 +319,78 @@ def run(data, def parse_opt(): parser = argparse.ArgumentParser() - parser.add_argument('--data', type=str, default=ROOT / 'data/coco128.yaml', help='dataset.yaml path') - parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.pt', help='model.pt path(s)') - parser.add_argument('--batch-size', type=int, default=32, help='batch size') - parser.add_argument('--imgsz', '--img', '--img-size', type=int, default=640, help='inference size (pixels)') - parser.add_argument('--conf-thres', type=float, default=0.001, help='confidence threshold') - parser.add_argument('--iou-thres', type=float, default=0.6, help='NMS IoU threshold') - parser.add_argument('--task', default='val', help='train, val, test, speed or study') - parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') - parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset') - parser.add_argument('--augment', action='store_true', help='augmented inference') - parser.add_argument('--verbose', action='store_true', help='report mAP by class') - parser.add_argument('--save-txt', action='store_true', help='save results to *.txt') - parser.add_argument('--save-hybrid', action='store_true', help='save label+prediction hybrid results to *.txt') - parser.add_argument('--save-conf', action='store_true', help='save confidences in --save-txt labels') - parser.add_argument('--save-json', action='store_true', help='save a COCO-JSON results file') - parser.add_argument('--project', default=ROOT / 'runs/val', help='save to project/name') - parser.add_argument('--name', default='exp', help='save to project/name') - parser.add_argument('--exist-ok', action='store_true', help='existing project/name ok, do not increment') - parser.add_argument('--half', action='store_true', help='use FP16 half-precision inference') + parser.add_argument("--data", type=str, default=ROOT / "data/coco128.yaml", help="dataset.yaml path") + parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "yolov5s.pt", help="model.pt path(s)") + parser.add_argument("--batch-size", type=int, default=32, help="batch size") + parser.add_argument("--imgsz", "--img", "--img-size", type=int, default=640, help="inference size (pixels)") + parser.add_argument("--conf-thres", type=float, default=0.001, help="confidence threshold") + parser.add_argument("--iou-thres", type=float, default=0.6, help="NMS IoU threshold") + parser.add_argument("--task", default="val", help="train, val, test, speed or study") + parser.add_argument("--device", default="", help="cuda device, i.e. 0 or 0,1,2,3 or cpu") + parser.add_argument("--single-cls", action="store_true", help="treat as single-class dataset") + parser.add_argument("--augment", action="store_true", help="augmented inference") + parser.add_argument("--verbose", action="store_true", help="report mAP by class") + parser.add_argument("--save-txt", action="store_true", help="save results to *.txt") + parser.add_argument("--save-hybrid", action="store_true", help="save label+prediction hybrid results to *.txt") + parser.add_argument("--save-conf", action="store_true", help="save confidences in --save-txt labels") + parser.add_argument("--save-json", action="store_true", help="save a COCO-JSON results file") + parser.add_argument("--project", default=ROOT / "runs/val", help="save to project/name") + parser.add_argument("--name", default="exp", help="save to project/name") + parser.add_argument("--exist-ok", action="store_true", help="existing project/name ok, do not increment") + parser.add_argument("--half", action="store_true", help="use FP16 half-precision inference") opt = parser.parse_args() opt.data = check_yaml(opt.data) # check YAML - opt.save_json |= opt.data.endswith('coco.yaml') + opt.save_json |= opt.data.endswith("coco.yaml") opt.save_txt |= opt.save_hybrid print_args(FILE.stem, opt) return opt def main(opt): - check_requirements(requirements=ROOT / 'requirements.txt', exclude=('tensorboard', 'thop')) + check_requirements(requirements=ROOT / "requirements.txt", exclude=("tensorboard", "thop")) - if opt.task in ('train', 'val', 'test'): # run normally + if opt.task in ("train", "val", "test"): # run normally if opt.conf_thres > 0.001: # https://github.com/ultralytics/yolov5/issues/1466 - LOGGER.info(f'WARNING: confidence threshold {opt.conf_thres} >> 0.001 will produce invalid mAP values.') + LOGGER.info(f"WARNING: confidence threshold {opt.conf_thres} >> 0.001 will produce invalid mAP values.") run(**vars(opt)) - elif opt.task == 'speed': # speed benchmarks + elif opt.task == "speed": # speed benchmarks # python val.py --task speed --data coco.yaml --batch 1 --weights yolov5n.pt yolov5s.pt... for w in opt.weights if isinstance(opt.weights, list) else [opt.weights]: - run(opt.data, weights=w, batch_size=opt.batch_size, imgsz=opt.imgsz, conf_thres=.25, iou_thres=.45, - device=opt.device, save_json=False, plots=False) - - elif opt.task == 'study': # run over a range of settings and save/plot + run( + opt.data, + weights=w, + batch_size=opt.batch_size, + imgsz=opt.imgsz, + conf_thres=0.25, + iou_thres=0.45, + device=opt.device, + save_json=False, + plots=False, + ) + + elif opt.task == "study": # run over a range of settings and save/plot # python val.py --task study --data coco.yaml --iou 0.7 --weights yolov5n.pt yolov5s.pt... x = list(range(256, 1536 + 128, 128)) # x axis (image sizes) for w in opt.weights if isinstance(opt.weights, list) else [opt.weights]: - f = f'study_{Path(opt.data).stem}_{Path(w).stem}.txt' # filename to save to + f = f"study_{Path(opt.data).stem}_{Path(w).stem}.txt" # filename to save to y = [] # y axis for i in x: # img-size - LOGGER.info(f'\nRunning {f} point {i}...') - r, _, t = run(opt.data, weights=w, batch_size=opt.batch_size, imgsz=i, conf_thres=opt.conf_thres, - iou_thres=opt.iou_thres, device=opt.device, save_json=opt.save_json, plots=False) + LOGGER.info(f"\nRunning {f} point {i}...") + r, _, t = run( + opt.data, + weights=w, + batch_size=opt.batch_size, + imgsz=i, + conf_thres=opt.conf_thres, + iou_thres=opt.iou_thres, + device=opt.device, + save_json=opt.save_json, + plots=False, + ) y.append(r + t) # results and times - np.savetxt(f, y, fmt='%10.4g') # save - os.system('zip -r study.zip study_*.txt') + np.savetxt(f, y, fmt="%10.4g") # save + os.system("zip -r study.zip study_*.txt") plot_val_study(x=x) # plot