diff --git a/.gitignore b/.gitignore index 3a96807..8bf2fea 100644 --- a/.gitignore +++ b/.gitignore @@ -363,4 +363,4 @@ MigrationBackup/ FodyWeavers.xsd .idea -cmake-build-* \ No newline at end of file +cmake-build-* diff --git a/YOLO/yolov5-master/detect.py b/YOLO/yolov5-master/detect.py new file mode 100644 index 0000000..d1df240 --- /dev/null +++ b/YOLO/yolov5-master/detect.py @@ -0,0 +1,437 @@ +# Ultralytics YOLOv5 🚀, AGPL-3.0 license +""" +Run YOLOv5 detection inference on images, videos, directories, globs, YouTube, webcam, streams, etc. + +Usage - sources: + $ python detect.py --weights yolov5s.pt --source 0 # webcam + img.jpg # image + vid.mp4 # video + screen # screenshot + path/ # directory + list.txt # list of images + list.streams # list of streams + 'path/*.jpg' # glob + 'https://youtu.be/LNwODJXcvt4' # YouTube + 'rtsp://example.com/media.mp4' # RTSP, RTMP, HTTP stream + +Usage - formats: + $ python detect.py --weights yolov5s.pt # PyTorch + yolov5s.torchscript # TorchScript + yolov5s.onnx # ONNX Runtime or OpenCV DNN with --dnn + yolov5s_openvino_model # OpenVINO + yolov5s.engine # TensorRT + yolov5s.mlpackage # CoreML (macOS-only) + yolov5s_saved_model # TensorFlow SavedModel + yolov5s.pb # TensorFlow GraphDef + yolov5s.tflite # TensorFlow Lite + yolov5s_edgetpu.tflite # TensorFlow Edge TPU + yolov5s_paddle_model # PaddlePaddle +""" + +import argparse +import csv +import os +import platform +import sys +from pathlib import Path + +import torch + +FILE = Path(__file__).resolve() +ROOT = FILE.parents[0] # YOLOv5 root directory +if str(ROOT) not in sys.path: + sys.path.append(str(ROOT)) # add ROOT to PATH +ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative + +from ultralytics.utils.plotting import Annotator, colors, save_one_box + +from models.common import DetectMultiBackend +from utils.dataloaders import IMG_FORMATS, VID_FORMATS, LoadImages, LoadScreenshots, LoadStreams +from utils.general import ( + LOGGER, + Profile, + check_file, + check_img_size, + check_imshow, + check_requirements, + colorstr, + cv2, + increment_path, + non_max_suppression, + print_args, + scale_boxes, + strip_optimizer, + xyxy2xywh, +) +from utils.torch_utils import select_device, smart_inference_mode + + +@smart_inference_mode() +def run( + weights=ROOT / "yolov5s.pt", # model path or triton URL + source=ROOT / "data/images", # file/dir/URL/glob/screen/0(webcam) + data=ROOT / "data/coco128.yaml", # dataset.yaml path + imgsz=(640, 640), # inference size (height, width) + 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_format=0, # save boxes coordinates in YOLO format or Pascal-VOC format (0 for YOLO and 1 for Pascal-VOC) + save_csv=False, # save results in CSV format + 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 + vid_stride=1, # video frame-rate stride +): + """ + Runs YOLOv5 detection inference on various sources like images, videos, directories, streams, etc. + + Args: + weights (str | Path): Path to the model weights file or a Triton URL. Default is 'yolov5s.pt'. + source (str | Path): Input source, which can be a file, directory, URL, glob pattern, screen capture, or webcam + index. Default is 'data/images'. + data (str | Path): Path to the dataset YAML file. Default is 'data/coco128.yaml'. + imgsz (tuple[int, int]): Inference image size as a tuple (height, width). Default is (640, 640). + conf_thres (float): Confidence threshold for detections. Default is 0.25. + iou_thres (float): Intersection Over Union (IOU) threshold for non-max suppression. Default is 0.45. + max_det (int): Maximum number of detections per image. Default is 1000. + device (str): CUDA device identifier (e.g., '0' or '0,1,2,3') or 'cpu'. Default is an empty string, which uses the + best available device. + view_img (bool): If True, display inference results using OpenCV. Default is False. + save_txt (bool): If True, save results in a text file. Default is False. + save_csv (bool): If True, save results in a CSV file. Default is False. + save_conf (bool): If True, include confidence scores in the saved results. Default is False. + save_crop (bool): If True, save cropped prediction boxes. Default is False. + nosave (bool): If True, do not save inference images or videos. Default is False. + classes (list[int]): List of class indices to filter detections by. Default is None. + agnostic_nms (bool): If True, perform class-agnostic non-max suppression. Default is False. + augment (bool): If True, use augmented inference. Default is False. + visualize (bool): If True, visualize feature maps. Default is False. + update (bool): If True, update all models' weights. Default is False. + project (str | Path): Directory to save results. Default is 'runs/detect'. + name (str): Name of the current experiment; used to create a subdirectory within 'project'. Default is 'exp'. + exist_ok (bool): If True, existing directories with the same name are reused instead of being incremented. Default is + False. + line_thickness (int): Thickness of bounding box lines in pixels. Default is 3. + hide_labels (bool): If True, do not display labels on bounding boxes. Default is False. + hide_conf (bool): If True, do not display confidence scores on bounding boxes. Default is False. + half (bool): If True, use FP16 half-precision inference. Default is False. + dnn (bool): If True, use OpenCV DNN backend for ONNX inference. Default is False. + vid_stride (int): Stride for processing video frames, to skip frames between processing. Default is 1. + + Returns: + None + + Examples: + ```python + from ultralytics import run + + # Run inference on an image + run(source='data/images/example.jpg', weights='yolov5s.pt', device='0') + + # Run inference on a video with specific confidence threshold + run(source='data/videos/example.mp4', weights='yolov5s.pt', conf_thres=0.4, device='0') + ``` + """ + source = str(source) + 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(".streams") or (is_url and not is_file) + screenshot = source.lower().startswith("screen") + 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 + + # Load model + device = select_device(device) + model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half) + stride, names, pt = model.stride, model.names, model.pt + imgsz = check_img_size(imgsz, s=stride) # check image size + + # Dataloader + bs = 1 # batch_size + if webcam: + view_img = check_imshow(warn=True) + dataset = LoadStreams(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride) + bs = len(dataset) + elif screenshot: + dataset = LoadScreenshots(source, img_size=imgsz, stride=stride, auto=pt) + else: + dataset = LoadImages(source, img_size=imgsz, stride=stride, auto=pt, vid_stride=vid_stride) + vid_path, vid_writer = [None] * bs, [None] * bs + + # Run inference + model.warmup(imgsz=(1 if pt or model.triton else bs, 3, *imgsz)) # warmup + seen, windows, dt = 0, [], (Profile(device=device), Profile(device=device), Profile(device=device)) + for path, im, im0s, vid_cap, s in dataset: + with dt[0]: + im = torch.from_numpy(im).to(model.device) + im = im.half() if model.fp16 else im.float() # uint8 to fp16/32 + im /= 255 # 0 - 255 to 0.0 - 1.0 + if len(im.shape) == 3: + im = im[None] # expand for batch dim + if model.xml and im.shape[0] > 1: + ims = torch.chunk(im, im.shape[0], 0) + + # Inference + with dt[1]: + visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False + if model.xml and im.shape[0] > 1: + pred = None + for image in ims: + if pred is None: + pred = model(image, augment=augment, visualize=visualize).unsqueeze(0) + else: + pred = torch.cat((pred, model(image, augment=augment, visualize=visualize).unsqueeze(0)), dim=0) + pred = [pred, None] + else: + pred = model(im, augment=augment, visualize=visualize) + # NMS + with dt[2]: + pred = non_max_suppression(pred, conf_thres, iou_thres, classes, agnostic_nms, max_det=max_det) + + # Second-stage classifier (optional) + # pred = utils.general.apply_classifier(pred, classifier_model, im, im0s) + + # Define the path for the CSV file + csv_path = save_dir / "predictions.csv" + + # Create or append to the CSV file + def write_to_csv(image_name, prediction, confidence): + """Writes prediction data for an image to a CSV file, appending if the file exists.""" + data = {"Image Name": image_name, "Prediction": prediction, "Confidence": confidence} + with open(csv_path, mode="a", newline="") as f: + writer = csv.DictWriter(f, fieldnames=data.keys()) + if not csv_path.is_file(): + writer.writeheader() + writer.writerow(data) + + # Process predictions + for i, det in enumerate(pred): # per image + seen += 1 + if webcam: # batch_size >= 1 + p, im0, frame = path[i], im0s[i].copy(), dataset.count + s += f"{i}: " + else: + p, im0, frame = path, im0s.copy(), getattr(dataset, "frame", 0) + + p = Path(p) # to Path + save_path = str(save_dir / p.name) # im.jpg + txt_path = str(save_dir / "labels" / p.stem) + ("" if dataset.mode == "image" else f"_{frame}") # im.txt + s += "{:g}x{:g} ".format(*im.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)) + if len(det): + # Rescale boxes from img_size to im0 size + det[:, :4] = scale_boxes(im.shape[2:], det[:, :4], im0.shape).round() + + # Print results + for c in det[:, 5].unique(): + n = (det[:, 5] == c).sum() # detections per class + s += f"{n} {names[int(c)]}{'s' * (n > 1)}, " # add to string + + # Write results + for *xyxy, conf, cls in reversed(det): + c = int(cls) # integer class + label = names[c] if hide_conf else f"{names[c]}" + confidence = float(conf) + confidence_str = f"{confidence:.2f}" + + if save_csv: + write_to_csv(p.name, label, confidence_str) + + if save_txt: # Write to file + if save_format == 0: + coords = ( + (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() + ) # normalized xywh + else: + coords = (torch.tensor(xyxy).view(1, 4) / gn).view(-1).tolist() # xyxy + line = (cls, *coords, conf) if save_conf else (cls, *coords) # label format + with open(f"{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}") + 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) + + # Stream results + im0 = annotator.result() + if view_img: + if platform.system() == "Linux" and p not in windows: + windows.append(p) + cv2.namedWindow(str(p), cv2.WINDOW_NORMAL | cv2.WINDOW_KEEPRATIO) # allow window resize (Linux) + cv2.resizeWindow(str(p), im0.shape[1], im0.shape[0]) + cv2.imshow(str(p), im0) + cv2.waitKey(1) # 1 millisecond + + # Save results (image with detections) + if save_img: + if dataset.mode == "image": + cv2.imwrite(save_path, im0) + else: # 'video' or 'stream' + if vid_path[i] != save_path: # new video + vid_path[i] = save_path + if isinstance(vid_writer[i], cv2.VideoWriter): + vid_writer[i].release() # release previous video writer + if vid_cap: # video + fps = vid_cap.get(cv2.CAP_PROP_FPS) + w = int(vid_cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + h = int(vid_cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + else: # stream + fps, w, h = 30, im0.shape[1], im0.shape[0] + save_path = str(Path(save_path).with_suffix(".mp4")) # force *.mp4 suffix on results videos + vid_writer[i] = cv2.VideoWriter(save_path, cv2.VideoWriter_fourcc(*"mp4v"), fps, (w, h)) + vid_writer[i].write(im0) + + # Print time (inference-only) + LOGGER.info(f"{s}{'' if len(det) else '(no detections), '}{dt[1].dt * 1E3:.1f}ms") + + # Print results + t = tuple(x.t / 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 "" + LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}") + if update: + strip_optimizer(weights[0]) # update model (to fix SourceChangeWarning) + + +def parse_opt(): + """ + Parse command-line arguments for YOLOv5 detection, allowing custom inference options and model configurations. + + Args: + --weights (str | list[str], optional): Model path or Triton URL. Defaults to ROOT / 'yolov5s.pt'. + --source (str, optional): File/dir/URL/glob/screen/0(webcam). Defaults to ROOT / 'data/images'. + --data (str, optional): Dataset YAML path. Provides dataset configuration information. + --imgsz (list[int], optional): Inference size (height, width). Defaults to [640]. + --conf-thres (float, optional): Confidence threshold. Defaults to 0.25. + --iou-thres (float, optional): NMS IoU threshold. Defaults to 0.45. + --max-det (int, optional): Maximum number of detections per image. Defaults to 1000. + --device (str, optional): CUDA device, i.e., '0' or '0,1,2,3' or 'cpu'. Defaults to "". + --view-img (bool, optional): Flag to display results. Defaults to False. + --save-txt (bool, optional): Flag to save results to *.txt files. Defaults to False. + --save-csv (bool, optional): Flag to save results in CSV format. Defaults to False. + --save-conf (bool, optional): Flag to save confidences in labels saved via --save-txt. Defaults to False. + --save-crop (bool, optional): Flag to save cropped prediction boxes. Defaults to False. + --nosave (bool, optional): Flag to prevent saving images/videos. Defaults to False. + --classes (list[int], optional): List of classes to filter results by, e.g., '--classes 0 2 3'. Defaults to None. + --agnostic-nms (bool, optional): Flag for class-agnostic NMS. Defaults to False. + --augment (bool, optional): Flag for augmented inference. Defaults to False. + --visualize (bool, optional): Flag for visualizing features. Defaults to False. + --update (bool, optional): Flag to update all models in the model directory. Defaults to False. + --project (str, optional): Directory to save results. Defaults to ROOT / 'runs/detect'. + --name (str, optional): Sub-directory name for saving results within --project. Defaults to 'exp'. + --exist-ok (bool, optional): Flag to allow overwriting if the project/name already exists. Defaults to False. + --line-thickness (int, optional): Thickness (in pixels) of bounding boxes. Defaults to 3. + --hide-labels (bool, optional): Flag to hide labels in the output. Defaults to False. + --hide-conf (bool, optional): Flag to hide confidences in the output. Defaults to False. + --half (bool, optional): Flag to use FP16 half-precision inference. Defaults to False. + --dnn (bool, optional): Flag to use OpenCV DNN for ONNX inference. Defaults to False. + --vid-stride (int, optional): Video frame-rate stride, determining the number of frames to skip in between + consecutive frames. Defaults to 1. + + Returns: + argparse.Namespace: Parsed command-line arguments as an argparse.Namespace object. + + Example: + ```python + from ultralytics import YOLOv5 + args = YOLOv5.parse_opt() + ``` + """ + parser = argparse.ArgumentParser() + parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "runs/train/exp9/weights/best.pt", help="model path or triton URL") + parser.add_argument("--source", type=str, default=ROOT / "data/images", help="file/dir/URL/glob/screen/0(webcam)") + parser.add_argument("--data", type=str, default=ROOT / "data/dimo2.yaml", help="(optional) dataset.yaml path") + 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-format", + type=int, + default=0, + help="whether to save boxes coordinates in YOLO format or Pascal-VOC format when save-txt is True, 0 for YOLO and 1 for Pascal-VOC", + ) + parser.add_argument("--save-csv", action="store_true", help="save results in CSV format") + 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("--vid-stride", type=int, default=1, help="video frame-rate stride") + opt = parser.parse_args() + opt.imgsz *= 2 if len(opt.imgsz) == 1 else 1 # expand + print_args(vars(opt)) + return opt + + +def main(opt): + """ + Executes YOLOv5 model inference based on provided command-line arguments, validating dependencies before running. + + Args: + opt (argparse.Namespace): Command-line arguments for YOLOv5 detection. See function `parse_opt` for details. + + Returns: + None + + Note: + This function performs essential pre-execution checks and initiates the YOLOv5 detection process based on user-specified + options. Refer to the usage guide and examples for more information about different sources and formats at: + https://github.com/ultralytics/ultralytics + + Example usage: + + ```python + if __name__ == "__main__": + opt = parse_opt() + main(opt) + ``` + """ + check_requirements(ROOT / "requirements.txt", exclude=("tensorboard", "thop")) + run(**vars(opt)) + + +if __name__ == "__main__": + opt = parse_opt() + main(opt) diff --git a/YOLO/yolov5-master/export.py b/YOLO/yolov5-master/export.py new file mode 100644 index 0000000..f5dc068 --- /dev/null +++ b/YOLO/yolov5-master/export.py @@ -0,0 +1,1546 @@ +# Ultralytics YOLOv5 🚀, AGPL-3.0 license +""" +Export a YOLOv5 PyTorch model to other formats. TensorFlow exports authored by https://github.com/zldrobit. + +Format | `export.py --include` | Model +--- | --- | --- +PyTorch | - | yolov5s.pt +TorchScript | `torchscript` | yolov5s.torchscript +ONNX | `onnx` | yolov5s.onnx +OpenVINO | `openvino` | yolov5s_openvino_model/ +TensorRT | `engine` | yolov5s.engine +CoreML | `coreml` | yolov5s.mlmodel +TensorFlow SavedModel | `saved_model` | yolov5s_saved_model/ +TensorFlow GraphDef | `pb` | yolov5s.pb +TensorFlow Lite | `tflite` | yolov5s.tflite +TensorFlow Edge TPU | `edgetpu` | yolov5s_edgetpu.tflite +TensorFlow.js | `tfjs` | yolov5s_web_model/ +PaddlePaddle | `paddle` | yolov5s_paddle_model/ + +Requirements: + $ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime openvino-dev tensorflow-cpu # CPU + $ pip install -r requirements.txt coremltools onnx onnx-simplifier onnxruntime-gpu openvino-dev tensorflow # GPU + +Usage: + $ python export.py --weights yolov5s.pt --include torchscript onnx openvino engine coreml tflite ... + +Inference: + $ python detect.py --weights yolov5s.pt # PyTorch + yolov5s.torchscript # TorchScript + yolov5s.onnx # ONNX Runtime or OpenCV DNN with --dnn + yolov5s_openvino_model # OpenVINO + yolov5s.engine # TensorRT + yolov5s.mlmodel # CoreML (macOS-only) + yolov5s_saved_model # TensorFlow SavedModel + yolov5s.pb # TensorFlow GraphDef + yolov5s.tflite # TensorFlow Lite + yolov5s_edgetpu.tflite # TensorFlow Edge TPU + yolov5s_paddle_model # PaddlePaddle + +TensorFlow.js: + $ cd .. && git clone https://github.com/zldrobit/tfjs-yolov5-example.git && cd tfjs-yolov5-example + $ npm install + $ ln -s ../../yolov5/yolov5s_web_model public/yolov5s_web_model + $ npm start +""" + +import argparse +import contextlib +import json +import os +import platform +import re +import subprocess +import sys +import time +import warnings +from pathlib import Path + +import pandas as pd +import torch +from torch.utils.mobile_optimizer import optimize_for_mobile + +FILE = Path(__file__).resolve() +ROOT = FILE.parents[0] # YOLOv5 root directory +if str(ROOT) not in sys.path: + sys.path.append(str(ROOT)) # add ROOT to PATH +if platform.system() != "Windows": + ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative + +from models.experimental import attempt_load +from models.yolo import ClassificationModel, Detect, DetectionModel, SegmentationModel +from utils.dataloaders import LoadImages +from utils.general import ( + LOGGER, + Profile, + check_dataset, + check_img_size, + check_requirements, + check_version, + check_yaml, + colorstr, + file_size, + get_default_args, + print_args, + url2file, + yaml_save, +) +from utils.torch_utils import select_device, smart_inference_mode + +MACOS = platform.system() == "Darwin" # macOS environment + + +class iOSModel(torch.nn.Module): + """An iOS-compatible wrapper for YOLOv5 models that normalizes input images based on their dimensions.""" + + def __init__(self, model, im): + """ + Initializes an iOS compatible model with normalization based on image dimensions. + + Args: + model (torch.nn.Module): The PyTorch model to be adapted for iOS compatibility. + im (torch.Tensor): An input tensor representing a batch of images with shape (B, C, H, W). + + Returns: + None: This method does not return any value. + + Notes: + This initializer configures normalization based on the input image dimensions, which is critical for + ensuring the model's compatibility and proper functionality on iOS devices. The normalization step + involves dividing by the image width if the image is square; otherwise, additional conditions might apply. + """ + super().__init__() + b, c, h, w = im.shape # batch, channel, height, width + self.model = model + self.nc = model.nc # number of classes + if w == h: + self.normalize = 1.0 / w + else: + self.normalize = torch.tensor([1.0 / w, 1.0 / h, 1.0 / w, 1.0 / h]) # broadcast (slower, smaller) + # np = model(im)[0].shape[1] # number of points + # self.normalize = torch.tensor([1. / w, 1. / h, 1. / w, 1. / h]).expand(np, 4) # explicit (faster, larger) + + def forward(self, x): + """ + Run a forward pass on the input tensor, returning class confidences and normalized coordinates. + + Args: + x (torch.Tensor): Input tensor containing the image data with shape (batch, channels, height, width). + + Returns: + torch.Tensor: Concatenated tensor with normalized coordinates (xywh), confidence scores (conf), + and class probabilities (cls), having shape (N, 4 + 1 + C), where N is the number of predictions, + and C is the number of classes. + + Examples: + ```python + model = iOSModel(pretrained_model, input_image) + output = model.forward(torch_input_tensor) + ``` + """ + xywh, conf, cls = self.model(x)[0].squeeze().split((4, 1, self.nc), 1) + return cls * conf, xywh * self.normalize # confidence (3780, 80), coordinates (3780, 4) + + +def export_formats(): + r""" + Returns a DataFrame of supported YOLOv5 model export formats and their properties. + + Returns: + pandas.DataFrame: A DataFrame containing supported export formats and their properties. The DataFrame + includes columns for format name, CLI argument suffix, file extension or directory name, and boolean flags + indicating if the export format supports training and detection. + + Examples: + ```python + formats = export_formats() + print(f"Supported export formats:\n{formats}") + ``` + + Notes: + The DataFrame contains the following columns: + - Format: The name of the model format (e.g., PyTorch, TorchScript, ONNX, etc.). + - Include Argument: The argument to use with the export script to include this format. + - File Suffix: File extension or directory name associated with the format. + - Supports Training: Whether the format supports training. + - Supports Detection: Whether the format supports detection. + """ + x = [ + ["PyTorch", "-", ".pt", True, True], + ["TorchScript", "torchscript", ".torchscript", True, True], + ["ONNX", "onnx", ".onnx", True, True], + ["OpenVINO", "openvino", "_openvino_model", True, False], + ["TensorRT", "engine", ".engine", False, True], + ["CoreML", "coreml", ".mlpackage", True, False], + ["TensorFlow SavedModel", "saved_model", "_saved_model", True, True], + ["TensorFlow GraphDef", "pb", ".pb", True, True], + ["TensorFlow Lite", "tflite", ".tflite", True, False], + ["TensorFlow Edge TPU", "edgetpu", "_edgetpu.tflite", False, False], + ["TensorFlow.js", "tfjs", "_web_model", False, False], + ["PaddlePaddle", "paddle", "_paddle_model", True, True], + ] + return pd.DataFrame(x, columns=["Format", "Argument", "Suffix", "CPU", "GPU"]) + + +def try_export(inner_func): + """ + Log success or failure, execution time, and file size for YOLOv5 model export functions wrapped with @try_export. + + Args: + inner_func (Callable): The model export function to be wrapped by the decorator. + + Returns: + Callable: The wrapped function that logs execution details. When executed, this wrapper function returns either: + - Tuple (str | torch.nn.Module): On success — the file path of the exported model and the model instance. + - Tuple (None, None): On failure — None values indicating export failure. + + Examples: + ```python + @try_export + def export_onnx(model, filepath): + # implementation here + pass + + exported_file, exported_model = export_onnx(yolo_model, 'path/to/save/model.onnx') + ``` + + Notes: + For additional requirements and model export formats, refer to the + [Ultralytics YOLOv5 GitHub repository](https://github.com/ultralytics/ultralytics). + """ + inner_args = get_default_args(inner_func) + + def outer_func(*args, **kwargs): + """Logs success/failure and execution details of model export functions wrapped with @try_export decorator.""" + prefix = inner_args["prefix"] + try: + with Profile() as dt: + f, model = inner_func(*args, **kwargs) + LOGGER.info(f"{prefix} export success ✅ {dt.t:.1f}s, saved as {f} ({file_size(f):.1f} MB)") + return f, model + except Exception as e: + LOGGER.info(f"{prefix} export failure ❌ {dt.t:.1f}s: {e}") + return None, None + + return outer_func + + +@try_export +def export_torchscript(model, im, file, optimize, prefix=colorstr("TorchScript:")): + """ + Export a YOLOv5 model to the TorchScript format. + + Args: + model (torch.nn.Module): The YOLOv5 model to be exported. + im (torch.Tensor): Example input tensor to be used for tracing the TorchScript model. + file (Path): File path where the exported TorchScript model will be saved. + optimize (bool): If True, applies optimizations for mobile deployment. + prefix (str): Optional prefix for log messages. Default is 'TorchScript:'. + + Returns: + (str | None, torch.jit.ScriptModule | None): A tuple containing the file path of the exported model + (as a string) and the TorchScript model (as a torch.jit.ScriptModule). If the export fails, both elements + of the tuple will be None. + + Notes: + - This function uses tracing to create the TorchScript model. + - Metadata, including the input image shape, model stride, and class names, is saved in an extra file (`config.txt`) + within the TorchScript model package. + - For mobile optimization, refer to the PyTorch tutorial: https://pytorch.org/tutorials/recipes/mobile_interpreter.html + + Example: + ```python + from pathlib import Path + import torch + from models.experimental import attempt_load + from utils.torch_utils import select_device + + # Load model + weights = 'yolov5s.pt' + device = select_device('') + model = attempt_load(weights, device=device) + + # Example input tensor + im = torch.zeros(1, 3, 640, 640).to(device) + + # Export model + file = Path('yolov5s.torchscript') + export_torchscript(model, im, file, optimize=False) + ``` + """ + LOGGER.info(f"\n{prefix} starting export with torch {torch.__version__}...") + f = file.with_suffix(".torchscript") + + ts = torch.jit.trace(model, im, strict=False) + d = {"shape": im.shape, "stride": int(max(model.stride)), "names": model.names} + extra_files = {"config.txt": json.dumps(d)} # torch._C.ExtraFilesMap() + if optimize: # https://pytorch.org/tutorials/recipes/mobile_interpreter.html + optimize_for_mobile(ts)._save_for_lite_interpreter(str(f), _extra_files=extra_files) + else: + ts.save(str(f), _extra_files=extra_files) + return f, None + + +@try_export +def export_onnx(model, im, file, opset, dynamic, simplify, prefix=colorstr("ONNX:")): + """ + Export a YOLOv5 model to ONNX format with dynamic axes support and optional model simplification. + + Args: + model (torch.nn.Module): The YOLOv5 model to be exported. + im (torch.Tensor): A sample input tensor for model tracing, usually the shape is (1, 3, height, width). + file (pathlib.Path | str): The output file path where the ONNX model will be saved. + opset (int): The ONNX opset version to use for export. + dynamic (bool): If True, enables dynamic axes for batch, height, and width dimensions. + simplify (bool): If True, applies ONNX model simplification for optimization. + prefix (str): A prefix string for logging messages, defaults to 'ONNX:'. + + Returns: + tuple[pathlib.Path | str, None]: The path to the saved ONNX model file and None (consistent with decorator). + + Raises: + ImportError: If required libraries for export (e.g., 'onnx', 'onnx-simplifier') are not installed. + AssertionError: If the simplification check fails. + + Notes: + The required packages for this function can be installed via: + ``` + pip install onnx onnx-simplifier onnxruntime onnxruntime-gpu + ``` + + Example: + ```python + from pathlib import Path + import torch + from models.experimental import attempt_load + from utils.torch_utils import select_device + + # Load model + weights = 'yolov5s.pt' + device = select_device('') + model = attempt_load(weights, map_location=device) + + # Example input tensor + im = torch.zeros(1, 3, 640, 640).to(device) + + # Export model + file_path = Path('yolov5s.onnx') + export_onnx(model, im, file_path, opset=12, dynamic=True, simplify=True) + ``` + """ + check_requirements("onnx>=1.12.0") + import onnx + + LOGGER.info(f"\n{prefix} starting export with onnx {onnx.__version__}...") + f = str(file.with_suffix(".onnx")) + + output_names = ["output0", "output1"] if isinstance(model, SegmentationModel) else ["output0"] + if dynamic: + dynamic = {"images": {0: "batch", 2: "height", 3: "width"}} # shape(1,3,640,640) + if isinstance(model, SegmentationModel): + dynamic["output0"] = {0: "batch", 1: "anchors"} # shape(1,25200,85) + dynamic["output1"] = {0: "batch", 2: "mask_height", 3: "mask_width"} # shape(1,32,160,160) + elif isinstance(model, DetectionModel): + dynamic["output0"] = {0: "batch", 1: "anchors"} # shape(1,25200,85) + + torch.onnx.export( + model.cpu() if dynamic else model, # --dynamic only compatible with cpu + im.cpu() if dynamic else im, + f, + verbose=False, + opset_version=opset, + do_constant_folding=True, # WARNING: DNN inference with torch>=1.12 may require do_constant_folding=False + input_names=["images"], + output_names=output_names, + dynamic_axes=dynamic or None, + ) + + # Checks + model_onnx = onnx.load(f) # load onnx model + onnx.checker.check_model(model_onnx) # check onnx model + + # Metadata + d = {"stride": int(max(model.stride)), "names": model.names} + for k, v in d.items(): + meta = model_onnx.metadata_props.add() + meta.key, meta.value = k, str(v) + onnx.save(model_onnx, f) + + # Simplify + if simplify: + try: + cuda = torch.cuda.is_available() + check_requirements(("onnxruntime-gpu" if cuda else "onnxruntime", "onnxslim")) + import onnxslim + + LOGGER.info(f"{prefix} slimming with onnxslim {onnxslim.__version__}...") + model_onnx = onnxslim.slim(model_onnx) + onnx.save(model_onnx, f) + except Exception as e: + LOGGER.info(f"{prefix} simplifier failure: {e}") + return f, model_onnx + + +@try_export +def export_openvino(file, metadata, half, int8, data, prefix=colorstr("OpenVINO:")): + """ + Export a YOLOv5 model to OpenVINO format with optional FP16 and INT8 quantization. + + Args: + file (Path): Path to the output file where the OpenVINO model will be saved. + metadata (dict): Dictionary including model metadata such as names and strides. + half (bool): If True, export the model with FP16 precision. + int8 (bool): If True, export the model with INT8 quantization. + data (str): Path to the dataset YAML file required for INT8 quantization. + prefix (str): Prefix string for logging purposes (default is "OpenVINO:"). + + Returns: + (str, openvino.runtime.Model | None): The OpenVINO model file path and openvino.runtime.Model object if export is + successful; otherwise, None. + + Notes: + - Requires `openvino-dev` package version 2023.0 or higher. Install with: + `$ pip install openvino-dev>=2023.0` + - For INT8 quantization, also requires `nncf` library version 2.5.0 or higher. Install with: + `$ pip install nncf>=2.5.0` + + Examples: + ```python + from pathlib import Path + from ultralytics import YOLOv5 + + model = YOLOv5('yolov5s.pt') + export_openvino(Path('yolov5s.onnx'), metadata={'names': model.names, 'stride': model.stride}, half=True, + int8=False, data='data.yaml') + ``` + + This will export the YOLOv5 model to OpenVINO with FP16 precision but without INT8 quantization, saving it to + the specified file path. + """ + check_requirements("openvino-dev>=2023.0") # requires openvino-dev: https://pypi.org/project/openvino-dev/ + import openvino.runtime as ov # noqa + from openvino.tools import mo # noqa + + LOGGER.info(f"\n{prefix} starting export with openvino {ov.__version__}...") + f = str(file).replace(file.suffix, f"_{'int8_' if int8 else ''}openvino_model{os.sep}") + f_onnx = file.with_suffix(".onnx") + f_ov = str(Path(f) / file.with_suffix(".xml").name) + + ov_model = mo.convert_model(f_onnx, model_name=file.stem, framework="onnx", compress_to_fp16=half) # export + + if int8: + check_requirements("nncf>=2.5.0") # requires at least version 2.5.0 to use the post-training quantization + import nncf + import numpy as np + + from utils.dataloaders import create_dataloader + + def gen_dataloader(yaml_path, task="train", imgsz=640, workers=4): + """Generates a DataLoader for model training or validation based on the given YAML dataset configuration.""" + data_yaml = check_yaml(yaml_path) + data = check_dataset(data_yaml) + dataloader = create_dataloader( + data[task], imgsz=imgsz, batch_size=1, stride=32, pad=0.5, single_cls=False, rect=False, workers=workers + )[0] + return dataloader + + # noqa: F811 + + def transform_fn(data_item): + """ + Quantization transform function. + + Extracts and preprocess input data from dataloader item for quantization. + + Args: + data_item: Tuple with data item produced by DataLoader during iteration + + Returns: + input_tensor: Input data for quantization + """ + assert data_item[0].dtype == torch.uint8, "input image must be uint8 for the quantization preprocessing" + + img = data_item[0].numpy().astype(np.float32) # uint8 to fp16/32 + img /= 255.0 # 0 - 255 to 0.0 - 1.0 + return np.expand_dims(img, 0) if img.ndim == 3 else img + + ds = gen_dataloader(data) + quantization_dataset = nncf.Dataset(ds, transform_fn) + ov_model = nncf.quantize(ov_model, quantization_dataset, preset=nncf.QuantizationPreset.MIXED) + + ov.serialize(ov_model, f_ov) # save + yaml_save(Path(f) / file.with_suffix(".yaml").name, metadata) # add metadata.yaml + return f, None + + +@try_export +def export_paddle(model, im, file, metadata, prefix=colorstr("PaddlePaddle:")): + """ + Export a YOLOv5 PyTorch model to PaddlePaddle format using X2Paddle, saving the converted model and metadata. + + Args: + model (torch.nn.Module): The YOLOv5 model to be exported. + im (torch.Tensor): Input tensor used for model tracing during export. + file (pathlib.Path): Path to the source file to be converted. + metadata (dict): Additional metadata to be saved alongside the model. + prefix (str): Prefix for logging information. + + Returns: + tuple (str, None): A tuple where the first element is the path to the saved PaddlePaddle model, and the + second element is None. + + Examples: + ```python + from pathlib import Path + import torch + + # Assume 'model' is a pre-trained YOLOv5 model and 'im' is an example input tensor + model = ... # Load your model here + im = torch.randn((1, 3, 640, 640)) # Dummy input tensor for tracing + file = Path("yolov5s.pt") + metadata = {"stride": 32, "names": ["person", "bicycle", "car", "motorbike"]} + + export_paddle(model=model, im=im, file=file, metadata=metadata) + ``` + + Notes: + Ensure that `paddlepaddle` and `x2paddle` are installed, as these are required for the export function. You can + install them via pip: + ``` + $ pip install paddlepaddle x2paddle + ``` + """ + check_requirements(("paddlepaddle", "x2paddle")) + import x2paddle + from x2paddle.convert import pytorch2paddle + + LOGGER.info(f"\n{prefix} starting export with X2Paddle {x2paddle.__version__}...") + f = str(file).replace(".pt", f"_paddle_model{os.sep}") + + pytorch2paddle(module=model, save_dir=f, jit_type="trace", input_examples=[im]) # export + yaml_save(Path(f) / file.with_suffix(".yaml").name, metadata) # add metadata.yaml + return f, None + + +@try_export +def export_coreml(model, im, file, int8, half, nms, mlmodel, prefix=colorstr("CoreML:")): + """ + Export a YOLOv5 model to CoreML format with optional NMS, INT8, and FP16 support. + + Args: + model (torch.nn.Module): The YOLOv5 model to be exported. + im (torch.Tensor): Example input tensor to trace the model. + file (pathlib.Path): Path object where the CoreML model will be saved. + int8 (bool): Flag indicating whether to use INT8 quantization (default is False). + half (bool): Flag indicating whether to use FP16 quantization (default is False). + nms (bool): Flag indicating whether to include Non-Maximum Suppression (default is False). + mlmodel (bool): Flag indicating whether to export as older *.mlmodel format (default is False). + prefix (str): Prefix string for logging purposes (default is 'CoreML:'). + + Returns: + tuple[pathlib.Path | None, None]: The path to the saved CoreML model file, or (None, None) if there is an error. + + Notes: + The exported CoreML model will be saved with a .mlmodel extension. + Quantization is supported only on macOS. + + Example: + ```python + from pathlib import Path + import torch + from models.yolo import Model + model = Model(cfg, ch=3, nc=80) + im = torch.randn(1, 3, 640, 640) + file = Path("yolov5s_coreml") + export_coreml(model, im, file, int8=False, half=False, nms=True, mlmodel=False) + ``` + """ + check_requirements("coremltools") + import coremltools as ct + + LOGGER.info(f"\n{prefix} starting export with coremltools {ct.__version__}...") + if mlmodel: + f = file.with_suffix(".mlmodel") + convert_to = "neuralnetwork" + precision = None + else: + f = file.with_suffix(".mlpackage") + convert_to = "mlprogram" + precision = ct.precision.FLOAT16 if half else ct.precision.FLOAT32 + if nms: + model = iOSModel(model, im) + 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])], + convert_to=convert_to, + compute_precision=precision, + ) + bits, mode = (8, "kmeans") if int8 else (16, "linear") if half else (32, None) + if bits < 32: + if mlmodel: + with warnings.catch_warnings(): + warnings.filterwarnings( + "ignore", category=DeprecationWarning + ) # suppress numpy==1.20 float warning, fixed in coremltools==7.0 + ct_model = ct.models.neural_network.quantization_utils.quantize_weights(ct_model, bits, mode) + elif bits == 8: + op_config = ct.optimize.coreml.OpPalettizerConfig(mode=mode, nbits=bits, weight_threshold=512) + config = ct.optimize.coreml.OptimizationConfig(global_config=op_config) + ct_model = ct.optimize.coreml.palettize_weights(ct_model, config) + ct_model.save(f) + return f, ct_model + + +@try_export +def export_engine( + model, im, file, half, dynamic, simplify, workspace=4, verbose=False, cache="", prefix=colorstr("TensorRT:") +): + """ + Export a YOLOv5 model to TensorRT engine format, requiring GPU and TensorRT>=7.0.0. + + Args: + model (torch.nn.Module): YOLOv5 model to be exported. + im (torch.Tensor): Input tensor of shape (B, C, H, W). + file (pathlib.Path): Path to save the exported model. + half (bool): Set to True to export with FP16 precision. + dynamic (bool): Set to True to enable dynamic input shapes. + simplify (bool): Set to True to simplify the model during export. + workspace (int): Workspace size in GB (default is 4). + verbose (bool): Set to True for verbose logging output. + cache (str): Path to save the TensorRT timing cache. + prefix (str): Log message prefix. + + Returns: + (pathlib.Path, None): Tuple containing the path to the exported model and None. + + Raises: + AssertionError: If executed on CPU instead of GPU. + RuntimeError: If there is a failure in parsing the ONNX file. + + Example: + ```python + from ultralytics import YOLOv5 + import torch + from pathlib import Path + + model = YOLOv5('yolov5s.pt') # Load a pre-trained YOLOv5 model + input_tensor = torch.randn(1, 3, 640, 640).cuda() # example input tensor on GPU + export_path = Path('yolov5s.engine') # export destination + + export_engine(model.model, input_tensor, export_path, half=True, dynamic=True, simplify=True, workspace=8, verbose=True) + ``` + """ + assert im.device.type != "cpu", "export running on CPU but must be on GPU, i.e. `python export.py --device 0`" + try: + import tensorrt as trt + except Exception: + if platform.system() == "Linux": + check_requirements("nvidia-tensorrt", cmds="-U --index-url https://pypi.ngc.nvidia.com") + import tensorrt as trt + + if trt.__version__[0] == "7": # TensorRT 7 handling https://github.com/ultralytics/yolov5/issues/6012 + grid = model.model[-1].anchor_grid + model.model[-1].anchor_grid = [a[..., :1, :1, :] for a in grid] + export_onnx(model, im, file, 12, dynamic, simplify) # opset 12 + model.model[-1].anchor_grid = grid + else: # TensorRT >= 8 + check_version(trt.__version__, "8.0.0", hard=True) # require tensorrt>=8.0.0 + export_onnx(model, im, file, 12, dynamic, simplify) # opset 12 + onnx = file.with_suffix(".onnx") + + LOGGER.info(f"\n{prefix} starting export with TensorRT {trt.__version__}...") + is_trt10 = int(trt.__version__.split(".")[0]) >= 10 # is TensorRT >= 10 + assert onnx.exists(), f"failed to export ONNX file: {onnx}" + f = file.with_suffix(".engine") # TensorRT engine file + logger = trt.Logger(trt.Logger.INFO) + if verbose: + logger.min_severity = trt.Logger.Severity.VERBOSE + + builder = trt.Builder(logger) + config = builder.create_builder_config() + if is_trt10: + config.set_memory_pool_limit(trt.MemoryPoolType.WORKSPACE, workspace << 30) + else: # TensorRT versions 7, 8 + config.max_workspace_size = workspace * 1 << 30 + if cache: # enable timing cache + Path(cache).parent.mkdir(parents=True, exist_ok=True) + buf = Path(cache).read_bytes() if Path(cache).exists() else b"" + timing_cache = config.create_timing_cache(buf) + config.set_timing_cache(timing_cache, ignore_mismatch=True) + flag = 1 << int(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH) + network = builder.create_network(flag) + parser = trt.OnnxParser(network, logger) + if not parser.parse_from_file(str(onnx)): + raise RuntimeError(f"failed to load ONNX file: {onnx}") + + inputs = [network.get_input(i) for i in range(network.num_inputs)] + outputs = [network.get_output(i) for i in range(network.num_outputs)] + for inp in inputs: + LOGGER.info(f'{prefix} input "{inp.name}" with shape{inp.shape} {inp.dtype}') + for out in outputs: + LOGGER.info(f'{prefix} output "{out.name}" with shape{out.shape} {out.dtype}') + + if dynamic: + if im.shape[0] <= 1: + LOGGER.warning(f"{prefix} WARNING ⚠️ --dynamic model requires maximum --batch-size argument") + profile = builder.create_optimization_profile() + for inp in inputs: + profile.set_shape(inp.name, (1, *im.shape[1:]), (max(1, im.shape[0] // 2), *im.shape[1:]), im.shape) + config.add_optimization_profile(profile) + + LOGGER.info(f"{prefix} building FP{16 if builder.platform_has_fast_fp16 and half else 32} engine as {f}") + if builder.platform_has_fast_fp16 and half: + config.set_flag(trt.BuilderFlag.FP16) + + build = builder.build_serialized_network if is_trt10 else builder.build_engine + with build(network, config) as engine, open(f, "wb") as t: + t.write(engine if is_trt10 else engine.serialize()) + if cache: # save timing cache + with open(cache, "wb") as c: + c.write(config.get_timing_cache().serialize()) + return f, None + + +@try_export +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, + keras=False, + prefix=colorstr("TensorFlow SavedModel:"), +): + """ + Export a YOLOv5 model to the TensorFlow SavedModel format, supporting dynamic axes and non-maximum suppression + (NMS). + + Args: + model (torch.nn.Module): The PyTorch model to convert. + im (torch.Tensor): Sample input tensor with shape (B, C, H, W) for tracing. + file (pathlib.Path): File path to save the exported model. + dynamic (bool): Flag to indicate whether dynamic axes should be used. + tf_nms (bool, optional): Enable TensorFlow non-maximum suppression (NMS). Default is False. + agnostic_nms (bool, optional): Enable class-agnostic NMS. Default is False. + topk_per_class (int, optional): Top K detections per class to keep before applying NMS. Default is 100. + topk_all (int, optional): Top K detections across all classes to keep before applying NMS. Default is 100. + iou_thres (float, optional): IoU threshold for NMS. Default is 0.45. + conf_thres (float, optional): Confidence threshold for detections. Default is 0.25. + keras (bool, optional): Save the model in Keras format if True. Default is False. + prefix (str, optional): Prefix for logging messages. Default is "TensorFlow SavedModel:". + + Returns: + tuple[str, tf.keras.Model | None]: A tuple containing the path to the saved model folder and the Keras model instance, + or None if TensorFlow export fails. + + Notes: + - The method supports TensorFlow versions up to 2.15.1. + - TensorFlow NMS may not be supported in older TensorFlow versions. + - If the TensorFlow version exceeds 2.13.1, it might cause issues when exporting to TFLite. + Refer to: https://github.com/ultralytics/yolov5/issues/12489 + + Example: + ```python + model, im = ... # Initialize your PyTorch model and input tensor + export_saved_model(model, im, Path("yolov5_saved_model"), dynamic=True) + ``` + """ + # YOLOv5 TensorFlow SavedModel export + try: + import tensorflow as tf + except Exception: + check_requirements(f"tensorflow{'' if torch.cuda.is_available() else '-macos' if MACOS else '-cpu'}<=2.15.1") + + import tensorflow as tf + from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2 + + from models.tf import TFModel + + LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...") + if tf.__version__ > "2.13.1": + helper_url = "https://github.com/ultralytics/yolov5/issues/12489" + LOGGER.info( + f"WARNING ⚠️ using Tensorflow {tf.__version__} > 2.13.1 might cause issue when exporting the model to tflite {helper_url}" + ) # handling issue https://github.com/ultralytics/yolov5/issues/12489 + 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) + im = tf.zeros((batch_size, *imgsz, ch)) # BHWC order for TensorFlow + _ = tf_model.predict(im, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres) + inputs = tf.keras.Input(shape=(*imgsz, ch), batch_size=None if dynamic else batch_size) + outputs = tf_model.predict(inputs, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres) + keras_model = tf.keras.Model(inputs=inputs, outputs=outputs) + keras_model.trainable = False + keras_model.summary() + if keras: + keras_model.save(f, save_format="tf") + else: + spec = tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype) + m = tf.function(lambda x: keras_model(x)) # full model + m = m.get_concrete_function(spec) + frozen_func = convert_variables_to_constants_v2(m) + tfm = tf.Module() + tfm.__call__ = tf.function(lambda x: frozen_func(x)[:4] if tf_nms else frozen_func(x), [spec]) + tfm.__call__(im) + tf.saved_model.save( + tfm, + f, + options=tf.saved_model.SaveOptions(experimental_custom_gradients=False) + if check_version(tf.__version__, "2.6") + else tf.saved_model.SaveOptions(), + ) + return f, keras_model + + +@try_export +def export_pb(keras_model, file, prefix=colorstr("TensorFlow GraphDef:")): + """ + Export YOLOv5 model to TensorFlow GraphDef (*.pb) format. + + Args: + keras_model (tf.keras.Model): The Keras model to be converted. + file (Path): The output file path where the GraphDef will be saved. + prefix (str): Optional prefix string; defaults to a colored string indicating TensorFlow GraphDef export status. + + Returns: + Tuple[Path, None]: The file path where the GraphDef model was saved and a None placeholder. + + Notes: + For more details, refer to the guide on frozen graphs: https://github.com/leimao/Frozen_Graph_TensorFlow + + Example: + ```python + from pathlib import Path + keras_model = ... # assume an existing Keras model + file = Path("model.pb") + export_pb(keras_model, file) + ``` + """ + 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") + + 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)) + frozen_func = convert_variables_to_constants_v2(m) + 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) + return f, None + + +@try_export +def export_tflite( + keras_model, im, file, int8, per_tensor, data, nms, agnostic_nms, prefix=colorstr("TensorFlow Lite:") +): + # YOLOv5 TensorFlow Lite export + """ + Export a YOLOv5 model to TensorFlow Lite format with optional INT8 quantization and NMS support. + + Args: + keras_model (tf.keras.Model): The Keras model to be exported. + im (torch.Tensor): An input image tensor for normalization and model tracing. + file (Path): The file path to save the TensorFlow Lite model. + int8 (bool): Enables INT8 quantization if True. + per_tensor (bool): If True, disables per-channel quantization. + data (str): Path to the dataset for representative dataset generation in INT8 quantization. + nms (bool): Enables Non-Maximum Suppression (NMS) if True. + agnostic_nms (bool): Enables class-agnostic NMS if True. + prefix (str): Prefix for log messages. + + Returns: + (str | None, tflite.Model | None): The file path of the exported TFLite model and the TFLite model instance, or None + if the export failed. + + Example: + ```python + from pathlib import Path + import torch + import tensorflow as tf + + # Load a Keras model wrapping a YOLOv5 model + keras_model = tf.keras.models.load_model('path/to/keras_model.h5') + + # Example input tensor + im = torch.zeros(1, 3, 640, 640) + + # Export the model + export_tflite(keras_model, im, Path('model.tflite'), int8=True, per_tensor=False, data='data/coco.yaml', + nms=True, agnostic_nms=False) + ``` + + Notes: + - Ensure TensorFlow and TensorFlow Lite dependencies are installed. + - INT8 quantization requires a representative dataset to achieve optimal accuracy. + - TensorFlow Lite models are suitable for efficient inference on mobile and edge devices. + """ + import tensorflow as tf + + 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") + + 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: + from models.tf import representative_dataset_gen + + dataset = LoadImages(check_dataset(check_yaml(data))["train"], img_size=imgsz, auto=False) + converter.representative_dataset = lambda: representative_dataset_gen(dataset, ncalib=100) + 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 = True + if per_tensor: + converter._experimental_disable_per_channel = True + f = str(file).replace(".pt", "-int8.tflite") + if nms or agnostic_nms: + converter.target_spec.supported_ops.append(tf.lite.OpsSet.SELECT_TF_OPS) + + tflite_model = converter.convert() + open(f, "wb").write(tflite_model) + return f, None + + +@try_export +def export_edgetpu(file, prefix=colorstr("Edge TPU:")): + """ + Exports a YOLOv5 model to Edge TPU compatible TFLite format; requires Linux and Edge TPU compiler. + + Args: + file (Path): Path to the YOLOv5 model file to be exported (.pt format). + prefix (str, optional): Prefix for logging messages. Defaults to colorstr("Edge TPU:"). + + Returns: + tuple[Path, None]: Path to the exported Edge TPU compatible TFLite model, None. + + Raises: + AssertionError: If the system is not Linux. + subprocess.CalledProcessError: If any subprocess call to install or run the Edge TPU compiler fails. + + Notes: + To use this function, ensure you have the Edge TPU compiler installed on your Linux system. You can find + installation instructions here: https://coral.ai/docs/edgetpu/compiler/. + + Example: + ```python + from pathlib import Path + file = Path('yolov5s.pt') + export_edgetpu(file) + ``` + """ + cmd = "edgetpu_compiler --version" + help_url = "https://coral.ai/docs/edgetpu/compiler/" + assert platform.system() == "Linux", f"export only supported on Linux. See {help_url}" + if subprocess.run(f"{cmd} > /dev/null 2>&1", shell=True).returncode != 0: + LOGGER.info(f"\n{prefix} export requires Edge TPU compiler. Attempting install from {help_url}") + sudo = subprocess.run("sudo --version >/dev/null", shell=True).returncode == 0 # sudo installed on system + for c in ( + "curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add -", + 'echo "deb https://packages.cloud.google.com/apt coral-edgetpu-stable main" | sudo tee /etc/apt/sources.list.d/coral-edgetpu.list', + "sudo apt-get update", + "sudo apt-get install edgetpu-compiler", + ): + subprocess.run(c if sudo else c.replace("sudo ", ""), shell=True, check=True) + ver = subprocess.run(cmd, shell=True, capture_output=True, check=True).stdout.decode().split()[-1] + + LOGGER.info(f"\n{prefix} starting export with Edge TPU compiler {ver}...") + f = str(file).replace(".pt", "-int8_edgetpu.tflite") # Edge TPU model + f_tfl = str(file).replace(".pt", "-int8.tflite") # TFLite model + + subprocess.run( + [ + "edgetpu_compiler", + "-s", + "-d", + "-k", + "10", + "--out_dir", + str(file.parent), + f_tfl, + ], + check=True, + ) + return f, None + + +@try_export +def export_tfjs(file, int8, prefix=colorstr("TensorFlow.js:")): + """ + Convert a YOLOv5 model to TensorFlow.js format with optional uint8 quantization. + + Args: + file (Path): Path to the YOLOv5 model file to be converted, typically having a ".pt" or ".onnx" extension. + int8 (bool): If True, applies uint8 quantization during the conversion process. + prefix (str): Optional prefix for logging messages, default is 'TensorFlow.js:' with color formatting. + + Returns: + (str, None): Tuple containing the output directory path as a string and None. + + Notes: + - This function requires the `tensorflowjs` package. Install it using: + ```shell + pip install tensorflowjs + ``` + - The converted TensorFlow.js model will be saved in a directory with the "_web_model" suffix appended to the original file name. + - The conversion involves running shell commands that invoke the TensorFlow.js converter tool. + + Example: + ```python + from pathlib import Path + file = Path('yolov5.onnx') + export_tfjs(file, int8=False) + ``` + """ + check_requirements("tensorflowjs") + 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"{f}/model.json" # *.json path + + args = [ + "tensorflowjs_converter", + "--input_format=tf_frozen_model", + "--quantize_uint8" if int8 else "", + "--output_node_names=Identity,Identity_1,Identity_2,Identity_3", + str(f_pb), + f, + ] + subprocess.run([arg for arg in args if arg], check=True) + + json = Path(f_json).read_text() + 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.?.?"}, ' + r'"Identity.?.?": {"name": "Identity.?.?"}, ' + r'"Identity.?.?": {"name": "Identity.?.?"}}}', + r'{"outputs": {"Identity": {"name": "Identity"}, ' + r'"Identity_1": {"name": "Identity_1"}, ' + r'"Identity_2": {"name": "Identity_2"}, ' + r'"Identity_3": {"name": "Identity_3"}}}', + json, + ) + j.write(subst) + return f, None + + +def add_tflite_metadata(file, metadata, num_outputs): + """ + Adds metadata to a TensorFlow Lite (TFLite) model file, supporting multiple outputs according to TensorFlow + guidelines. + + Args: + file (str): Path to the TFLite model file to which metadata will be added. + metadata (dict): Metadata information to be added to the model, structured as required by the TFLite metadata schema. + Common keys include "name", "description", "version", "author", and "license". + num_outputs (int): Number of output tensors the model has, used to configure the metadata properly. + + Returns: + None + + Example: + ```python + metadata = { + "name": "yolov5", + "description": "YOLOv5 object detection model", + "version": "1.0", + "author": "Ultralytics", + "license": "Apache License 2.0" + } + add_tflite_metadata("model.tflite", metadata, num_outputs=4) + ``` + + Note: + TFLite metadata can include information such as model name, version, author, and other relevant details. + For more details on the structure of the metadata, refer to TensorFlow Lite + [metadata guidelines](https://www.tensorflow.org/lite/models/convert/metadata). + """ + with contextlib.suppress(ImportError): + # check_requirements('tflite_support') + from tflite_support import flatbuffers + from tflite_support import metadata as _metadata + from tflite_support import metadata_schema_py_generated as _metadata_fb + + tmp_file = Path("/tmp/meta.txt") + with open(tmp_file, "w") as meta_f: + meta_f.write(str(metadata)) + + model_meta = _metadata_fb.ModelMetadataT() + label_file = _metadata_fb.AssociatedFileT() + label_file.name = tmp_file.name + model_meta.associatedFiles = [label_file] + + subgraph = _metadata_fb.SubGraphMetadataT() + subgraph.inputTensorMetadata = [_metadata_fb.TensorMetadataT()] + subgraph.outputTensorMetadata = [_metadata_fb.TensorMetadataT()] * num_outputs + model_meta.subgraphMetadata = [subgraph] + + b = flatbuffers.Builder(0) + b.Finish(model_meta.Pack(b), _metadata.MetadataPopulator.METADATA_FILE_IDENTIFIER) + metadata_buf = b.Output() + + populator = _metadata.MetadataPopulator.with_model_file(file) + populator.load_metadata_buffer(metadata_buf) + populator.load_associated_files([str(tmp_file)]) + populator.populate() + tmp_file.unlink() + + +def pipeline_coreml(model, im, file, names, y, mlmodel, prefix=colorstr("CoreML Pipeline:")): + """ + Convert a PyTorch YOLOv5 model to CoreML format with Non-Maximum Suppression (NMS), handling different input/output + shapes, and saving the model. + + Args: + model (torch.nn.Module): The YOLOv5 PyTorch model to be converted. + im (torch.Tensor): Example input tensor with shape (N, C, H, W), where N is the batch size, C is the number of channels, + H is the height, and W is the width. + file (Path): Path to save the converted CoreML model. + names (dict[int, str]): Dictionary mapping class indices to class names. + y (torch.Tensor): Output tensor from the PyTorch model's forward pass. + mlmodel (bool): Flag indicating whether to export as older *.mlmodel format (default is False). + prefix (str): Custom prefix for logging messages. + + Returns: + (Path): Path to the saved CoreML model (.mlmodel). + + Raises: + AssertionError: If the number of class names does not match the number of classes in the model. + + Notes: + - This function requires `coremltools` to be installed. + - Running this function on a non-macOS environment might not support some features. + - Flexible input shapes and additional NMS options can be customized within the function. + + Examples: + ```python + from pathlib import Path + import torch + + model = torch.load('yolov5s.pt') # Load YOLOv5 model + im = torch.zeros((1, 3, 640, 640)) # Example input tensor + + names = {0: "person", 1: "bicycle", 2: "car", ...} # Define class names + + y = model(im) # Perform forward pass to get model output + + output_file = Path('yolov5s.mlmodel') # Convert to CoreML + pipeline_coreml(model, im, output_file, names, y) + ``` + """ + import coremltools as ct + from PIL import Image + + f = file.with_suffix(".mlmodel") if mlmodel else file.with_suffix(".mlpackage") + print(f"{prefix} starting pipeline with coremltools {ct.__version__}...") + batch_size, ch, h, w = list(im.shape) # BCHW + t = time.time() + + # YOLOv5 Output shapes + spec = model.get_spec() + out0, out1 = iter(spec.description.output) + if platform.system() == "Darwin": + img = Image.new("RGB", (w, h)) # img(192 width, 320 height) + # img = torch.zeros((*opt.img_size, 3)).numpy() # img size(320,192,3) iDetection + out = model.predict({"image": img}) + out0_shape, out1_shape = out[out0.name].shape, out[out1.name].shape + else: # linux and windows can not run model.predict(), get sizes from pytorch output y + s = tuple(y[0].shape) + out0_shape, out1_shape = (s[1], s[2] - 5), (s[1], 4) # (3780, 80), (3780, 4) + + # Checks + nx, ny = spec.description.input[0].type.imageType.width, spec.description.input[0].type.imageType.height + na, nc = out0_shape + # na, nc = out0.type.multiArrayType.shape # number anchors, classes + assert len(names) == nc, f"{len(names)} names found for nc={nc}" # check + + # Define output shapes (missing) + out0.type.multiArrayType.shape[:] = out0_shape # (3780, 80) + out1.type.multiArrayType.shape[:] = out1_shape # (3780, 4) + # spec.neuralNetwork.preprocessing[0].featureName = '0' + + # Flexible input shapes + # from coremltools.models.neural_network import flexible_shape_utils + # s = [] # shapes + # s.append(flexible_shape_utils.NeuralNetworkImageSize(320, 192)) + # s.append(flexible_shape_utils.NeuralNetworkImageSize(640, 384)) # (height, width) + # flexible_shape_utils.add_enumerated_image_sizes(spec, feature_name='image', sizes=s) + # r = flexible_shape_utils.NeuralNetworkImageSizeRange() # shape ranges + # r.add_height_range((192, 640)) + # r.add_width_range((192, 640)) + # flexible_shape_utils.update_image_size_range(spec, feature_name='image', size_range=r) + + # Print + print(spec.description) + + # Model from spec + weights_dir = None + weights_dir = None if mlmodel else str(f / "Data/com.apple.CoreML/weights") + model = ct.models.MLModel(spec, weights_dir=weights_dir) + + # 3. Create NMS protobuf + nms_spec = ct.proto.Model_pb2.Model() + nms_spec.specificationVersion = 5 + for i in range(2): + decoder_output = model._spec.description.output[i].SerializeToString() + nms_spec.description.input.add() + nms_spec.description.input[i].ParseFromString(decoder_output) + nms_spec.description.output.add() + nms_spec.description.output[i].ParseFromString(decoder_output) + + nms_spec.description.output[0].name = "confidence" + nms_spec.description.output[1].name = "coordinates" + + output_sizes = [nc, 4] + for i in range(2): + ma_type = nms_spec.description.output[i].type.multiArrayType + ma_type.shapeRange.sizeRanges.add() + ma_type.shapeRange.sizeRanges[0].lowerBound = 0 + ma_type.shapeRange.sizeRanges[0].upperBound = -1 + ma_type.shapeRange.sizeRanges.add() + ma_type.shapeRange.sizeRanges[1].lowerBound = output_sizes[i] + ma_type.shapeRange.sizeRanges[1].upperBound = output_sizes[i] + del ma_type.shape[:] + + nms = nms_spec.nonMaximumSuppression + nms.confidenceInputFeatureName = out0.name # 1x507x80 + nms.coordinatesInputFeatureName = out1.name # 1x507x4 + nms.confidenceOutputFeatureName = "confidence" + nms.coordinatesOutputFeatureName = "coordinates" + nms.iouThresholdInputFeatureName = "iouThreshold" + nms.confidenceThresholdInputFeatureName = "confidenceThreshold" + nms.iouThreshold = 0.45 + nms.confidenceThreshold = 0.25 + nms.pickTop.perClass = True + nms.stringClassLabels.vector.extend(names.values()) + nms_model = ct.models.MLModel(nms_spec) + + # 4. Pipeline models together + pipeline = ct.models.pipeline.Pipeline( + input_features=[ + ("image", ct.models.datatypes.Array(3, ny, nx)), + ("iouThreshold", ct.models.datatypes.Double()), + ("confidenceThreshold", ct.models.datatypes.Double()), + ], + output_features=["confidence", "coordinates"], + ) + pipeline.add_model(model) + pipeline.add_model(nms_model) + + # Correct datatypes + pipeline.spec.description.input[0].ParseFromString(model._spec.description.input[0].SerializeToString()) + pipeline.spec.description.output[0].ParseFromString(nms_model._spec.description.output[0].SerializeToString()) + pipeline.spec.description.output[1].ParseFromString(nms_model._spec.description.output[1].SerializeToString()) + + # Update metadata + pipeline.spec.specificationVersion = 5 + pipeline.spec.description.metadata.versionString = "https://github.com/ultralytics/yolov5" + pipeline.spec.description.metadata.shortDescription = "https://github.com/ultralytics/yolov5" + pipeline.spec.description.metadata.author = "glenn.jocher@ultralytics.com" + pipeline.spec.description.metadata.license = "https://github.com/ultralytics/yolov5/blob/master/LICENSE" + pipeline.spec.description.metadata.userDefined.update( + { + "classes": ",".join(names.values()), + "iou_threshold": str(nms.iouThreshold), + "confidence_threshold": str(nms.confidenceThreshold), + } + ) + + # Save the model + model = ct.models.MLModel(pipeline.spec, weights_dir=weights_dir) + model.input_description["image"] = "Input image" + model.input_description["iouThreshold"] = f"(optional) IOU Threshold override (default: {nms.iouThreshold})" + model.input_description["confidenceThreshold"] = ( + f"(optional) Confidence Threshold override (default: {nms.confidenceThreshold})" + ) + model.output_description["confidence"] = 'Boxes × Class confidence (see user-defined metadata "classes")' + model.output_description["coordinates"] = "Boxes × [x, y, width, height] (relative to image size)" + model.save(f) # pipelined + print(f"{prefix} pipeline success ({time.time() - t:.2f}s), saved as {f} ({file_size(f):.1f} MB)") + + +@smart_inference_mode() +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"), # include formats + half=False, # FP16 half-precision export + inplace=False, # set YOLOv5 Detect() inplace=True + keras=False, # use Keras + optimize=False, # TorchScript: optimize for mobile + int8=False, # CoreML/TF INT8 quantization + per_tensor=False, # TF per tensor quantization + dynamic=False, # ONNX/TF/TensorRT: dynamic axes + cache="", # TensorRT: timing cache path + simplify=False, # ONNX: simplify model + mlmodel=False, # CoreML: Export in *.mlmodel format + opset=12, # ONNX: opset version + verbose=False, # TensorRT: verbose log + workspace=4, # TensorRT: workspace size (GB) + nms=False, # TF: add NMS to model + agnostic_nms=False, # TF: add agnostic NMS to model + 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 +): + """ + Exports a YOLOv5 model to specified formats including ONNX, TensorRT, CoreML, and TensorFlow. + + Args: + data (str | Path): Path to the dataset YAML configuration file. Default is 'data/coco128.yaml'. + weights (str | Path): Path to the pretrained model weights file. Default is 'yolov5s.pt'. + imgsz (tuple): Image size as (height, width). Default is (640, 640). + batch_size (int): Batch size for exporting the model. Default is 1. + device (str): Device to run the export on, e.g., '0' for GPU, 'cpu' for CPU. Default is 'cpu'. + include (tuple): Formats to include in the export. Default is ('torchscript', 'onnx'). + half (bool): Flag to export model with FP16 half-precision. Default is False. + inplace (bool): Set the YOLOv5 Detect() module inplace=True. Default is False. + keras (bool): Flag to use Keras for TensorFlow SavedModel export. Default is False. + optimize (bool): Optimize TorchScript model for mobile deployment. Default is False. + int8 (bool): Apply INT8 quantization for CoreML or TensorFlow models. Default is False. + per_tensor (bool): Apply per tensor quantization for TensorFlow models. Default is False. + dynamic (bool): Enable dynamic axes for ONNX, TensorFlow, or TensorRT exports. Default is False. + cache (str): TensorRT timing cache path. Default is an empty string. + simplify (bool): Simplify the ONNX model during export. Default is False. + opset (int): ONNX opset version. Default is 12. + verbose (bool): Enable verbose logging for TensorRT export. Default is False. + workspace (int): TensorRT workspace size in GB. Default is 4. + nms (bool): Add non-maximum suppression (NMS) to the TensorFlow model. Default is False. + agnostic_nms (bool): Add class-agnostic NMS to the TensorFlow model. Default is False. + topk_per_class (int): Top-K boxes per class to keep for TensorFlow.js NMS. Default is 100. + topk_all (int): Top-K boxes for all classes to keep for TensorFlow.js NMS. Default is 100. + iou_thres (float): IoU threshold for NMS. Default is 0.45. + conf_thres (float): Confidence threshold for NMS. Default is 0.25. + mlmodel (bool): Flag to use *.mlmodel for CoreML export. Default is False. + + Returns: + None + + Notes: + - Model export is based on the specified formats in the 'include' argument. + - Be cautious of combinations where certain flags are mutually exclusive, such as `--half` and `--dynamic`. + + Example: + ```python + run( + data="data/coco128.yaml", + weights="yolov5s.pt", + imgsz=(640, 640), + batch_size=1, + device="cpu", + include=("torchscript", "onnx"), + half=False, + inplace=False, + keras=False, + optimize=False, + int8=False, + per_tensor=False, + dynamic=False, + cache="", + simplify=False, + opset=12, + verbose=False, + mlmodel=False, + workspace=4, + nms=False, + agnostic_nms=False, + topk_per_class=100, + topk_all=100, + iou_thres=0.45, + conf_thres=0.25, + ) + ``` + """ + t = time.time() + include = [x.lower() for x in include] # to lowercase + fmts = tuple(export_formats()["Argument"][1:]) # --include arguments + flags = [x in include for x in fmts] + assert sum(flags) == len(include), f"ERROR: Invalid --include {include}, valid --include arguments are {fmts}" + jit, onnx, xml, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle = flags # export booleans + file = Path(url2file(weights) if str(weights).startswith(("http:/", "https:/")) else weights) # PyTorch weights + + # Load PyTorch model + device = select_device(device) + if half: + assert device.type != "cpu" or coreml, "--half only compatible with GPU export, i.e. use --device 0" + assert not dynamic, "--half not compatible with --dynamic, i.e. use either --half or --dynamic but not both" + model = attempt_load(weights, device=device, inplace=True, fuse=True) # load FP32 model + + # Checks + imgsz *= 2 if len(imgsz) == 1 else 1 # expand + if optimize: + assert device.type == "cpu", "--optimize not compatible with cuda devices, i.e. use --device cpu" + + # Input + gs = int(max(model.stride)) # grid size (max stride) + imgsz = [check_img_size(x, gs) for x in imgsz] # verify img_size are gs-multiples + ch = next(model.parameters()).size(1) # require input image channels + im = torch.zeros(batch_size, ch, *imgsz).to(device) # image size(1,3,320,192) BCHW iDetection + + # Update model + model.eval() + for k, m in model.named_modules(): + if isinstance(m, Detect): + m.inplace = inplace + m.dynamic = dynamic + m.export = True + + for _ in range(2): + y = model(im) # dry runs + if half and not coreml: + im, model = im.half(), model.half() # to FP16 + shape = tuple((y[0] if isinstance(y, tuple) else y).shape) # model output shape + metadata = {"stride": int(max(model.stride)), "names": model.names} # model metadata + LOGGER.info(f"\n{colorstr('PyTorch:')} starting from {file} with output shape {shape} ({file_size(file):.1f} MB)") + + # Exports + f = [""] * len(fmts) # exported filenames + warnings.filterwarnings(action="ignore", category=torch.jit.TracerWarning) # suppress TracerWarning + if jit: # TorchScript + f[0], _ = export_torchscript(model, im, file, optimize) + if engine: # TensorRT required before ONNX + f[1], _ = export_engine(model, im, file, half, dynamic, simplify, workspace, verbose, cache) + if onnx or xml: # OpenVINO requires ONNX + f[2], _ = export_onnx(model, im, file, opset, dynamic, simplify) + if xml: # OpenVINO + f[3], _ = export_openvino(file, metadata, half, int8, data) + if coreml: # CoreML + f[4], ct_model = export_coreml(model, im, file, int8, half, nms, mlmodel) + if nms: + pipeline_coreml(ct_model, im, file, model.names, y, mlmodel) + if any((saved_model, pb, tflite, edgetpu, tfjs)): # TensorFlow formats + assert not tflite or not tfjs, "TFLite and TF.js models must be exported separately, please pass only one type." + assert not isinstance(model, ClassificationModel), "ClassificationModel export to TF formats not yet supported." + f[5], s_model = export_saved_model( + model.cpu(), + im, + file, + dynamic, + tf_nms=nms or agnostic_nms or tfjs, + agnostic_nms=agnostic_nms or tfjs, + topk_per_class=topk_per_class, + topk_all=topk_all, + iou_thres=iou_thres, + conf_thres=conf_thres, + keras=keras, + ) + if pb or tfjs: # pb prerequisite to tfjs + f[6], _ = export_pb(s_model, file) + if tflite or edgetpu: + f[7], _ = export_tflite( + s_model, im, file, int8 or edgetpu, per_tensor, data=data, nms=nms, agnostic_nms=agnostic_nms + ) + if edgetpu: + f[8], _ = export_edgetpu(file) + add_tflite_metadata(f[8] or f[7], metadata, num_outputs=len(s_model.outputs)) + if tfjs: + f[9], _ = export_tfjs(file, int8) + if paddle: # PaddlePaddle + f[10], _ = export_paddle(model, im, file, metadata) + + # Finish + f = [str(x) for x in f if x] # filter out '' and None + if any(f): + cls, det, seg = (isinstance(model, x) for x in (ClassificationModel, DetectionModel, SegmentationModel)) # type + det &= not seg # segmentation models inherit from SegmentationModel(DetectionModel) + dir = Path("segment" if seg else "classify" if cls else "") + h = "--half" if half else "" # --half FP16 inference arg + s = ( + "# WARNING ⚠️ ClassificationModel not yet supported for PyTorch Hub AutoShape inference" + if cls + else "# WARNING ⚠️ SegmentationModel not yet supported for PyTorch Hub AutoShape inference" + if seg + else "" + ) + LOGGER.info( + f'\nExport complete ({time.time() - t:.1f}s)' + f"\nResults saved to {colorstr('bold', file.parent.resolve())}" + f"\nDetect: python {dir / ('detect.py' if det else 'predict.py')} --weights {f[-1]} {h}" + f"\nValidate: python {dir / 'val.py'} --weights {f[-1]} {h}" + f"\nPyTorch Hub: model = torch.hub.load('ultralytics/yolov5', 'custom', '{f[-1]}') {s}" + f'\nVisualize: https://netron.app' + ) + return f # return list of exported files/dirs + + +def parse_opt(known=False): + """ + Parse command-line options for YOLOv5 model export configurations. + + Args: + known (bool): If True, uses `argparse.ArgumentParser.parse_known_args`; otherwise, uses `argparse.ArgumentParser.parse_args`. + Default is False. + + Returns: + argparse.Namespace: Object containing parsed command-line arguments. + + Example: + ```python + opts = parse_opt() + print(opts.data) + print(opts.weights) + ``` + """ + 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("--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("--keras", action="store_true", help="TF: use Keras") + parser.add_argument("--optimize", action="store_true", help="TorchScript: optimize for mobile") + parser.add_argument("--int8", action="store_true", help="CoreML/TF/OpenVINO INT8 quantization") + parser.add_argument("--per-tensor", action="store_true", help="TF per-tensor quantization") + parser.add_argument("--dynamic", action="store_true", help="ONNX/TF/TensorRT: dynamic axes") + parser.add_argument("--cache", type=str, default="", help="TensorRT: timing cache file path") + parser.add_argument("--simplify", action="store_true", help="ONNX: simplify model") + parser.add_argument("--mlmodel", action="store_true", help="CoreML: Export in *.mlmodel format") + parser.add_argument("--opset", type=int, default=17, help="ONNX: opset version") + parser.add_argument("--verbose", action="store_true", help="TensorRT: verbose log") + parser.add_argument("--workspace", type=int, default=4, help="TensorRT: workspace size (GB)") + parser.add_argument("--nms", action="store_true", help="TF: add NMS to model") + parser.add_argument("--agnostic-nms", action="store_true", help="TF: add agnostic NMS to model") + 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"], + help="torchscript, onnx, openvino, engine, coreml, saved_model, pb, tflite, edgetpu, tfjs, paddle", + ) + opt = parser.parse_known_args()[0] if known else parser.parse_args() + print_args(vars(opt)) + return opt + + +def main(opt): + """Run(**vars(opt)) # Execute the run function with parsed options.""" + for opt.weights in opt.weights if isinstance(opt.weights, list) else [opt.weights]: + run(**vars(opt)) + + +if __name__ == "__main__": + opt = parse_opt() + main(opt) diff --git a/YOLO/yolov5-master/model/11.12.onnx b/YOLO/yolov5-master/model/11.12.onnx new file mode 100644 index 0000000..9fc19f5 Binary files /dev/null and b/YOLO/yolov5-master/model/11.12.onnx differ diff --git a/YOLO/yolov5-master/requirements.txt b/YOLO/yolov5-master/requirements.txt new file mode 100644 index 0000000..dcd23bf --- /dev/null +++ b/YOLO/yolov5-master/requirements.txt @@ -0,0 +1,49 @@ +# YOLOv5 requirements +# Usage: pip install -r requirements.txt + +# Base ------------------------------------------------------------------------ +gitpython>=3.1.30 +matplotlib>=3.3 +numpy>=1.23.5 +opencv-python>=4.1.1 +pillow>=10.3.0 +psutil # system resources +PyYAML>=5.3.1 +requests>=2.32.2 +scipy>=1.4.1 +thop>=0.1.1 # FLOPs computation +torch>=1.8.0 # see https://pytorch.org/get-started/locally (recommended) +torchvision>=0.9.0 +tqdm>=4.66.3 +ultralytics>=8.2.34 # https://ultralytics.com +# protobuf<=3.20.1 # https://github.com/ultralytics/yolov5/issues/8012 + +# Logging --------------------------------------------------------------------- +# tensorboard>=2.4.1 +# clearml>=1.2.0 +# comet + +# Plotting -------------------------------------------------------------------- +pandas>=1.1.4 +seaborn>=0.11.0 + +# Export ---------------------------------------------------------------------- +# coremltools>=6.0 # CoreML export +# onnx>=1.10.0 # ONNX export +# onnx-simplifier>=0.4.1 # ONNX simplifier +# nvidia-pyindex # TensorRT export +# nvidia-tensorrt # TensorRT export +# scikit-learn<=1.1.2 # CoreML quantization +# tensorflow>=2.4.0,<=2.13.1 # TF exports (-cpu, -aarch64, -macos) +# tensorflowjs>=3.9.0 # TF.js export +# openvino-dev>=2023.0 # OpenVINO export + +# Deploy ---------------------------------------------------------------------- +setuptools>=70.0.0 # Snyk vulnerability fix +# tritonclient[all]~=2.24.0 + +# Extras ---------------------------------------------------------------------- +# ipython # interactive notebook +# mss # screenshots +# albumentations>=1.0.3 +# pycocotools>=2.0.6 # COCO mAP diff --git a/YOLO/yolov5-master/train.py b/YOLO/yolov5-master/train.py new file mode 100644 index 0000000..1a84872 --- /dev/null +++ b/YOLO/yolov5-master/train.py @@ -0,0 +1,987 @@ +# Ultralytics YOLOv5 🚀, AGPL-3.0 license +""" +Train a YOLOv5 model on a custom dataset. Models and datasets download automatically from the latest YOLOv5 release. + +Usage - Single-GPU training: + $ python train.py --data coco128.yaml --weights yolov5s.pt --img 640 # from pretrained (recommended) + $ python train.py --data coco128.yaml --weights '' --cfg yolov5s.yaml --img 640 # from scratch + +Usage - Multi-GPU DDP training: + $ python -m torch.distributed.run --nproc_per_node 4 --master_port 1 train.py --data coco128.yaml --weights yolov5s.pt --img 640 --device 0,1,2,3 + +Models: https://github.com/ultralytics/yolov5/tree/master/models +Datasets: https://github.com/ultralytics/yolov5/tree/master/data +Tutorial: https://docs.ultralytics.com/yolov5/tutorials/train_custom_data +""" + +import argparse +import math +import os +import random +import subprocess +import sys +import time +from copy import deepcopy +from datetime import datetime, timedelta +from pathlib import Path + +os.environ["GIT_PYTHON_REFRESH"] = "quiet" +try: + import comet_ml # must be imported before torch (if installed) +except ImportError: + comet_ml = None + +import numpy as np +import torch +import torch.distributed as dist +import torch.nn as nn +import yaml +from torch.optim import lr_scheduler +from tqdm import tqdm + +FILE = Path(__file__).resolve() +ROOT = FILE.parents[0] # YOLOv5 root directory +if str(ROOT) not in sys.path: + sys.path.append(str(ROOT)) # add ROOT to PATH +ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative + +import val as validate # for end-of-epoch mAP +from models.experimental import attempt_load +from models.yolo import Model +from utils.autoanchor import check_anchors +from utils.autobatch import check_train_batch_size +from utils.callbacks import Callbacks +from utils.dataloaders import create_dataloader +from utils.downloads import attempt_download, is_url +from utils.general import ( + LOGGER, + TQDM_BAR_FORMAT, + check_amp, + check_dataset, + check_file, + check_git_info, + 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, + yaml_save, +) +from utils.loggers import LOGGERS, Loggers +from utils.loggers.comet.comet_utils import check_comet_resume +from utils.loss import ComputeLoss +from utils.metrics import fitness +from utils.plots import plot_evolve +from utils.torch_utils import ( + EarlyStopping, + ModelEMA, + de_parallel, + select_device, + smart_DDP, + smart_optimizer, + smart_resume, + 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)) +GIT_INFO = check_git_info() + + +def train(hyp, opt, device, callbacks): + """ + Train a YOLOv5 model on a custom dataset using specified hyperparameters, options, and device, managing datasets, + model architecture, loss computation, and optimizer steps. + + Args: + hyp (str | dict): Path to the hyperparameters YAML file or a dictionary of hyperparameters. + opt (argparse.Namespace): Parsed command-line arguments containing training options. + device (torch.device): Device on which training occurs, e.g., 'cuda' or 'cpu'. + callbacks (Callbacks): Callback functions for various training events. + + Returns: + None + + Models and datasets download automatically from the latest YOLOv5 release. + + Example: + Single-GPU training: + ```bash + $ python train.py --data coco128.yaml --weights yolov5s.pt --img 640 # from pretrained (recommended) + $ python train.py --data coco128.yaml --weights '' --cfg yolov5s.yaml --img 640 # from scratch + ``` + + Multi-GPU DDP training: + ```bash + $ python -m torch.distributed.run --nproc_per_node 4 --master_port 1 train.py --data coco128.yaml --weights + yolov5s.pt --img 640 --device 0,1,2,3 + ``` + + For more usage details, refer to: + - Models: https://github.com/ultralytics/yolov5/tree/master/models + - Datasets: https://github.com/ultralytics/yolov5/tree/master/data + - Tutorial: https://docs.ultralytics.com/yolov5/tutorials/train_custom_data + """ + 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, + ) + callbacks.run("on_pretrain_routine_start") + + # Directories + 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" + + # Hyperparameters + if isinstance(hyp, str): + 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())) + opt.hyp = hyp.copy() # for saving hyps to checkpoints + + # Save run settings + if not evolve: + yaml_save(save_dir / "hyp.yaml", hyp) + yaml_save(save_dir / "opt.yaml", vars(opt)) + + # Loggers + data_dict = None + if RANK in {-1, 0}: + include_loggers = list(LOGGERS) + if getattr(opt, "ndjson_console", False): + include_loggers.append("ndjson_console") + if getattr(opt, "ndjson_file", False): + include_loggers.append("ndjson_file") + + loggers = Loggers( + save_dir=save_dir, + weights=weights, + opt=opt, + hyp=hyp, + logger=LOGGER, + include=tuple(include_loggers), + ) + + # Register actions + for k in methods(loggers): + callbacks.register_action(k, callback=getattr(loggers, k)) + + # Process custom dataset artifact link + data_dict = loggers.remote_dataset + if resume: # If resuming runs from remote artifact + weights, epochs, hyp, batch_size = opt.weights, opt.epochs, opt.hyp, opt.batch_size + + # Config + plots = not evolve and not opt.noplots # create plots + cuda = device.type != "cpu" + init_seeds(opt.seed + 1 + RANK, deterministic=True) + 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 = {0: "item"} if single_cls and len(data_dict["names"]) != 1 else data_dict["names"] # class names + 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") + if pretrained: + with torch_distributed_zero_first(LOCAL_RANK): + weights = attempt_download(weights) # download if not found locally + ckpt = torch.load(weights, map_location="cpu") # load checkpoint to CPU to avoid CUDA memory leak + 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 + else: + model = Model(cfg, ch=3, nc=nc, anchors=hyp.get("anchors")).to(device) # create + amp = check_amp(model) # check AMP + + # Freeze + freeze = [f"model.{x}." for x in (freeze if len(freeze) > 1 else range(freeze[0]))] # layers to freeze + for k, v in model.named_parameters(): + v.requires_grad = True # train all layers + # v.register_hook(lambda x: torch.nan_to_num(x)) # NaN to 0 (commented for erratic training results) + if any(x in k for x in freeze): + LOGGER.info(f"freezing {k}") + v.requires_grad = False + + # Image size + gs = max(int(model.stride.max()), 32) # grid size (max stride) + imgsz = check_img_size(opt.imgsz, gs, floor=gs * 2) # verify imgsz is gs-multiple + + # Batch size + if RANK == -1 and batch_size == -1: # single-GPU only, estimate best batch size + batch_size = check_train_batch_size(model, imgsz, amp) + loggers.on_params_update({"batch_size": batch_size}) + + # 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 + optimizer = smart_optimizer(model, opt.optimizer, hyp["lr0"], hyp["momentum"], hyp["weight_decay"]) + + # Scheduler + if opt.cos_lr: + lf = one_cycle(1, hyp["lrf"], epochs) # cosine 1->hyp['lrf'] + else: + + def lf(x): + """Linear learning rate scheduler function with decay calculated by epoch proportion.""" + return (1 - x / epochs) * (1.0 - hyp["lrf"]) + hyp["lrf"] # linear + + scheduler = lr_scheduler.LambdaLR(optimizer, lr_lambda=lf) # plot_lr_scheduler(optimizer, scheduler, epochs) + + # EMA + ema = ModelEMA(model) if RANK in {-1, 0} else None + + # Resume + best_fitness, start_epoch = 0.0, 0 + if pretrained: + if resume: + best_fitness, start_epoch, epochs = smart_resume(ckpt, optimizer, ema, weights, epochs, resume) + 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://docs.ultralytics.com/yolov5/tutorials/multi_gpu_training 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()") + + # Trainloader + train_loader, dataset = create_dataloader( + train_path, + imgsz, + batch_size // WORLD_SIZE, + gs, + single_cls, + hyp=hyp, + augment=True, + cache=None if opt.cache == "val" else opt.cache, + rect=opt.rect, + rank=LOCAL_RANK, + workers=workers, + image_weights=opt.image_weights, + quad=opt.quad, + prefix=colorstr("train: "), + shuffle=True, + seed=opt.seed, + ) + labels = np.concatenate(dataset.labels, 0) + mlc = int(labels[:, 0].max()) # max label class + 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 * 2, + pad=0.5, + prefix=colorstr("val: "), + )[0] + + if not resume: + if not opt.noautoanchor: + check_anchors(dataset, model=model, thr=hyp["anchor_t"], imgsz=imgsz) # run AutoAnchor + model.half().float() # pre-reduce anchor precision + + callbacks.run("on_pretrain_routine_end", labels, names) + + # DDP mode + if cuda and RANK != -1: + model = smart_DDP(model) + + # Model attributes + 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 + 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 + model.names = names + + # Start training + t0 = time.time() + nb = len(train_loader) # number of batches + nw = max(round(hyp["warmup_epochs"] * nb), 100) # number of warmup iterations, max(3 epochs, 100 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 + results = (0, 0, 0, 0, 0, 0, 0) # P, R, mAP@.5, mAP@.5-.95, val_loss(box, obj, cls) + scheduler.last_epoch = start_epoch - 1 # do not move + scaler = torch.cuda.amp.GradScaler(enabled=amp) + stopper, stop = EarlyStopping(patience=opt.patience), False + compute_loss = ComputeLoss(model) # init loss class + callbacks.run("on_train_start") + LOGGER.info( + f'Image sizes {imgsz} train, {imgsz} val\n' + f'Using {train_loader.num_workers * WORLD_SIZE} 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 ------------------------------------------------------------------ + callbacks.run("on_train_epoch_start") + model.train() + + # Update image weights (optional, single-GPU only) + if opt.image_weights: + cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 / nc # class weights + iw = labels_to_image_weights(dataset.labels, nc=nc, class_weights=cw) # image weights + dataset.indices = random.choices(range(dataset.n), weights=iw, k=dataset.n) # rand weighted idx + + # Update mosaic border (optional) + # b = int(random.uniform(0.25 * imgsz, 0.75 * imgsz + gs) // gs * gs) + # dataset.mosaic_border = [b - imgsz, -b] # height, width borders + + mloss = torch.zeros(3, device=device) # mean losses + if RANK != -1: + train_loader.sampler.set_epoch(epoch) + pbar = enumerate(train_loader) + LOGGER.info(("\n" + "%11s" * 7) % ("Epoch", "GPU_mem", "box_loss", "obj_loss", "cls_loss", "Instances", "Size")) + if RANK in {-1, 0}: + pbar = tqdm(pbar, total=nb, bar_format=TQDM_BAR_FORMAT) # progress bar + optimizer.zero_grad() + for i, (imgs, targets, paths, _) in pbar: # batch ------------------------------------------------------------- + callbacks.run("on_train_batch_start") + ni = i + nb * epoch # number integrated batches (since train start) + imgs = imgs.to(device, non_blocking=True).float() / 255 # uint8 to float32, 0-255 to 0.0-1.0 + + # Warmup + if ni <= nw: + xi = [0, nw] # x interp + # compute_loss.gr = np.interp(ni, xi, [0.0, 1.0]) # iou loss ratio (obj_loss = 1.0 or iou) + accumulate = max(1, np.interp(ni, xi, [1, nbs / 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 == 0 else 0.0, x["initial_lr"] * lf(epoch)]) + if "momentum" in x: + x["momentum"] = np.interp(ni, xi, [hyp["warmup_momentum"], hyp["momentum"]]) + + # Multi-scale + if opt.multi_scale: + sz = random.randrange(int(imgsz * 0.5), int(imgsz * 1.5) + gs) // gs * gs # size + sf = sz / max(imgs.shape[2:]) # scale factor + if sf != 1: + ns = [math.ceil(x * sf / gs) * gs for x in imgs.shape[2:]] # new shape (stretched to gs-multiple) + imgs = nn.functional.interpolate(imgs, size=ns, mode="bilinear", align_corners=False) + + # Forward + with torch.cuda.amp.autocast(amp): + pred = model(imgs) # forward + loss, loss_items = compute_loss(pred, targets.to(device)) # loss scaled by batch_size + if RANK != -1: + loss *= WORLD_SIZE # gradient averaged between devices in DDP mode + if opt.quad: + loss *= 4.0 + + # Backward + scaler.scale(loss).backward() + + # Optimize - https://pytorch.org/docs/master/notes/amp_examples.html + if ni - last_opt_step >= accumulate: + scaler.unscale_(optimizer) # unscale gradients + torch.nn.utils.clip_grad_norm_(model.parameters(), max_norm=10.0) # clip gradients + scaler.step(optimizer) # optimizer.step + scaler.update() + optimizer.zero_grad() + if ema: + ema.update(model) + last_opt_step = ni + + # 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( + ("%11s" * 2 + "%11.4g" * 5) + % (f"{epoch}/{epochs - 1}", mem, *mloss, targets.shape[0], imgs.shape[-1]) + ) + callbacks.run("on_train_batch_end", model, ni, imgs, targets, paths, list(mloss)) + if callbacks.stop_training: + return + # end batch ------------------------------------------------------------------------------------------------ + + # Scheduler + 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"]) + final_epoch = (epoch + 1 == epochs) or stopper.possible_stop + if not noval or final_epoch: # Calculate mAP + results, maps, _ = validate.run( + data_dict, + batch_size=batch_size // WORLD_SIZE * 2, + imgsz=imgsz, + half=amp, + 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] + stop = stopper(epoch=epoch, fitness=fi) # early stop check + 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) + + # 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(), + "opt": vars(opt), + "git": GIT_INFO, # {remote, branch, commit} if a git repo + "date": datetime.now().isoformat(), + } + + # Save last, best and delete + torch.save(ckpt, last) + if best_fitness == fi: + torch.save(ckpt, best) + if opt.save_period > 0 and epoch % opt.save_period == 0: + torch.save(ckpt, w / f"epoch{epoch}.pt") + del ckpt + callbacks.run("on_model_save", last, epoch, final_epoch, best_fitness, fi) + + # EarlyStopping + if RANK != -1: # if DDP training + broadcast_list = [stop if RANK == 0 else None] + dist.broadcast_object_list(broadcast_list, 0) # broadcast 'stop' to all ranks + if RANK != 0: + stop = broadcast_list[0] + if stop: + break # must break all DDP ranks + + # 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.") + for f in last, best: + if f.exists(): + strip_optimizer(f) # strip optimizers + if f is best: + LOGGER.info(f"\nValidating {f}...") + results, _, _ = validate.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 at iou 0.65 + single_cls=single_cls, + dataloader=val_loader, + save_dir=save_dir, + save_json=is_coco, + verbose=True, + plots=plots, + 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_train_end", last, best, epoch, results) + + torch.cuda.empty_cache() + return results + + +def parse_opt(known=False): + """ + Parse command-line arguments for YOLOv5 training, validation, and testing. + + Args: + known (bool, optional): If True, parses known arguments, ignoring the unknown. Defaults to False. + + Returns: + (argparse.Namespace): Parsed command-line arguments containing options for YOLOv5 execution. + + Example: + ```python + from ultralytics.yolo import parse_opt + opt = parse_opt() + print(opt) + ``` + + Links: + - Models: https://github.com/ultralytics/yolov5/tree/master/models + - Datasets: https://github.com/ultralytics/yolov5/tree/master/data + - Tutorial: https://docs.ultralytics.com/yolov5/tutorials/train_custom_data + """ + 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/dimo3.yaml", help="dataset.yaml path") + parser.add_argument("--hyp", type=str, default=ROOT / "data/hyps/hyp.scratch-low.yaml", help="hyperparameters path") + parser.add_argument("--epochs", type=int, default=100, help="total training epochs") + parser.add_argument("--batch-size", type=int, default=4, 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") + parser.add_argument("--noplots", action="store_true", help="save no plot files") + parser.add_argument("--evolve", type=int, nargs="?", const=300, help="evolve hyperparameters for x generations") + parser.add_argument( + "--evolve_population", type=str, default=ROOT / "data/hyps", help="location for loading population" + ) + parser.add_argument("--resume_evolve", type=str, default=None, help="resume evolve from last generation") + parser.add_argument("--bucket", type=str, default="", help="gsutil bucket") + parser.add_argument("--cache", type=str, nargs="?", const="ram", help="image --cache ram/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("--optimizer", type=str, choices=["SGD", "Adam", "AdamW"], default="SGD", help="optimizer") + parser.add_argument("--sync-bn", action="store_true", help="use SyncBatchNorm, only available in DDP mode") + parser.add_argument("--workers", type=int, default=0, help="max dataloader workers (per RANK in DDP mode)") + 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("--cos-lr", action="store_true", help="cosine LR scheduler") + 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", nargs="+", type=int, default=[0], help="Freeze layers: backbone=10, first3=0 1 2") + parser.add_argument("--save-period", type=int, default=-1, help="Save checkpoint every x epochs (disabled if < 1)") + parser.add_argument("--seed", type=int, default=0, help="Global training seed") + parser.add_argument("--local_rank", type=int, default=-1, help="Automatic DDP Multi-GPU argument, do not modify") + + # Logger arguments + parser.add_argument("--entity", default=None, help="Entity") + parser.add_argument("--upload_dataset", nargs="?", const=True, default=False, help='Upload data, "val" option') + parser.add_argument("--bbox_interval", type=int, default=-1, help="Set bounding-box image logging interval") + parser.add_argument("--artifact_alias", type=str, default="latest", help="Version of dataset artifact to use") + + # NDJSON logging + parser.add_argument("--ndjson-console", action="store_true", help="Log ndjson to console") + parser.add_argument("--ndjson-file", action="store_true", help="Log ndjson to file") + + return parser.parse_known_args()[0] if known else parser.parse_args() + + +def main(opt, callbacks=Callbacks()): + """ + Runs the main entry point for training or hyperparameter evolution with specified options and optional callbacks. + + Args: + opt (argparse.Namespace): The command-line arguments parsed for YOLOv5 training and evolution. + callbacks (ultralytics.utils.callbacks.Callbacks, optional): Callback functions for various training stages. + Defaults to Callbacks(). + + Returns: + None + + Note: + For detailed usage, refer to: + https://github.com/ultralytics/yolov5/tree/master/models + """ + if RANK in {-1, 0}: + print_args(vars(opt)) + check_git_status() + check_requirements(ROOT / "requirements.txt") + + # Resume (from specified or most recent last.pt) + if opt.resume and not check_comet_resume(opt) and not opt.evolve: + last = Path(check_file(opt.resume) if isinstance(opt.resume, str) else get_latest_run()) + opt_yaml = last.parent.parent / "opt.yaml" # train options yaml + opt_data = opt.data # original dataset + if opt_yaml.is_file(): + with open(opt_yaml, errors="ignore") as f: + d = yaml.safe_load(f) + else: + d = torch.load(last, map_location="cpu")["opt"] + opt = argparse.Namespace(**d) # replace + opt.cfg, opt.weights, opt.resume = "", str(last), True # reinstate + if is_url(opt_data): + opt.data = check_file(opt_data) # avoid HUB resume auth timeout + 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" + if opt.evolve: + if opt.project == str(ROOT / "runs/train"): # if default project name, rename to runs/evolve + opt.project = str(ROOT / "runs/evolve") + opt.exist_ok, opt.resume = opt.resume, False # pass resume to exist_ok and disable resume + if opt.name == "cfg": + opt.name = Path(opt.cfg).stem # use model.yaml as name + 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: + msg = "is not compatible with YOLOv5 Multi-GPU DDP training" + assert not opt.image_weights, f"--image-weights {msg}" + assert not opt.evolve, f"--evolve {msg}" + assert opt.batch_size != -1, f"AutoBatch with --batch-size -1 {msg}, please pass a valid --batch-size" + assert opt.batch_size % WORLD_SIZE == 0, f"--batch-size {opt.batch_size} must be multiple of WORLD_SIZE" + assert torch.cuda.device_count() > LOCAL_RANK, "insufficient CUDA devices for DDP command" + torch.cuda.set_device(LOCAL_RANK) + device = torch.device("cuda", LOCAL_RANK) + dist.init_process_group( + backend="nccl" if dist.is_nccl_available() else "gloo", timeout=timedelta(seconds=10800) + ) + + # Train + if not opt.evolve: + train(opt.hyp, opt, device, callbacks) + + # Evolve hyperparameters (optional) + else: + # Hyperparameter evolution metadata (including this hyperparameter True-False, lower_limit, upper_limit) + meta = { + "lr0": (False, 1e-5, 1e-1), # initial learning rate (SGD=1E-2, Adam=1E-3) + "lrf": (False, 0.01, 1.0), # final OneCycleLR learning rate (lr0 * lrf) + "momentum": (False, 0.6, 0.98), # SGD momentum/Adam beta1 + "weight_decay": (False, 0.0, 0.001), # optimizer weight decay + "warmup_epochs": (False, 0.0, 5.0), # warmup epochs (fractions ok) + "warmup_momentum": (False, 0.0, 0.95), # warmup initial momentum + "warmup_bias_lr": (False, 0.0, 0.2), # warmup initial bias lr + "box": (False, 0.02, 0.2), # box loss gain + "cls": (False, 0.2, 4.0), # cls loss gain + "cls_pw": (False, 0.5, 2.0), # cls BCELoss positive_weight + "obj": (False, 0.2, 4.0), # obj loss gain (scale with pixels) + "obj_pw": (False, 0.5, 2.0), # obj BCELoss positive_weight + "iou_t": (False, 0.1, 0.7), # IoU training threshold + "anchor_t": (False, 2.0, 8.0), # anchor-multiple threshold + "anchors": (False, 2.0, 10.0), # anchors per output grid (0 to ignore) + "fl_gamma": (False, 0.0, 2.0), # focal loss gamma (efficientDet default gamma=1.5) + "hsv_h": (True, 0.0, 0.1), # image HSV-Hue augmentation (fraction) + "hsv_s": (True, 0.0, 0.9), # image HSV-Saturation augmentation (fraction) + "hsv_v": (True, 0.0, 0.9), # image HSV-Value augmentation (fraction) + "degrees": (True, 0.0, 45.0), # image rotation (+/- deg) + "translate": (True, 0.0, 0.9), # image translation (+/- fraction) + "scale": (True, 0.0, 0.9), # image scale (+/- gain) + "shear": (True, 0.0, 10.0), # image shear (+/- deg) + "perspective": (True, 0.0, 0.001), # image perspective (+/- fraction), range 0-0.001 + "flipud": (True, 0.0, 1.0), # image flip up-down (probability) + "fliplr": (True, 0.0, 1.0), # image flip left-right (probability) + "mosaic": (True, 0.0, 1.0), # image mosaic (probability) + "mixup": (True, 0.0, 1.0), # image mixup (probability) + "copy_paste": (True, 0.0, 1.0), # segment copy-paste (probability) + } + + # GA configs + pop_size = 50 + mutation_rate_min = 0.01 + mutation_rate_max = 0.5 + crossover_rate_min = 0.5 + crossover_rate_max = 1 + min_elite_size = 2 + max_elite_size = 5 + tournament_size_min = 2 + tournament_size_max = 10 + + 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 opt.noautoanchor: + del hyp["anchors"], meta["anchors"] + 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" + if opt.bucket: + # download evolve.csv if exists + subprocess.run( + [ + "gsutil", + "cp", + f"gs://{opt.bucket}/evolve.csv", + str(evolve_csv), + ] + ) + + # Delete the items in meta dictionary whose first value is False + del_ = [item for item, value_ in meta.items() if value_[0] is False] + hyp_GA = hyp.copy() # Make a copy of hyp dictionary + for item in del_: + del meta[item] # Remove the item from meta dictionary + del hyp_GA[item] # Remove the item from hyp_GA dictionary + + # Set lower_limit and upper_limit arrays to hold the search space boundaries + lower_limit = np.array([meta[k][1] for k in hyp_GA.keys()]) + upper_limit = np.array([meta[k][2] for k in hyp_GA.keys()]) + + # Create gene_ranges list to hold the range of values for each gene in the population + gene_ranges = [(lower_limit[i], upper_limit[i]) for i in range(len(upper_limit))] + + # Initialize the population with initial_values or random values + initial_values = [] + + # If resuming evolution from a previous checkpoint + if opt.resume_evolve is not None: + assert os.path.isfile(ROOT / opt.resume_evolve), "evolve population path is wrong!" + with open(ROOT / opt.resume_evolve, errors="ignore") as f: + evolve_population = yaml.safe_load(f) + for value in evolve_population.values(): + value = np.array([value[k] for k in hyp_GA.keys()]) + initial_values.append(list(value)) + + # If not resuming from a previous checkpoint, generate initial values from .yaml files in opt.evolve_population + else: + yaml_files = [f for f in os.listdir(opt.evolve_population) if f.endswith(".yaml")] + for file_name in yaml_files: + with open(os.path.join(opt.evolve_population, file_name)) as yaml_file: + value = yaml.safe_load(yaml_file) + value = np.array([value[k] for k in hyp_GA.keys()]) + initial_values.append(list(value)) + + # Generate random values within the search space for the rest of the population + if initial_values is None: + population = [generate_individual(gene_ranges, len(hyp_GA)) for _ in range(pop_size)] + elif pop_size > 1: + population = [generate_individual(gene_ranges, len(hyp_GA)) for _ in range(pop_size - len(initial_values))] + for initial_value in initial_values: + population = [initial_value] + population + + # Run the genetic algorithm for a fixed number of generations + list_keys = list(hyp_GA.keys()) + for generation in range(opt.evolve): + if generation >= 1: + save_dict = {} + for i in range(len(population)): + little_dict = {list_keys[j]: float(population[i][j]) for j in range(len(population[i]))} + save_dict[f"gen{str(generation)}number{str(i)}"] = little_dict + + with open(save_dir / "evolve_population.yaml", "w") as outfile: + yaml.dump(save_dict, outfile, default_flow_style=False) + + # Adaptive elite size + elite_size = min_elite_size + int((max_elite_size - min_elite_size) * (generation / opt.evolve)) + # Evaluate the fitness of each individual in the population + fitness_scores = [] + for individual in population: + for key, value in zip(hyp_GA.keys(), individual): + hyp_GA[key] = value + hyp.update(hyp_GA) + results = train(hyp.copy(), opt, device, callbacks) + callbacks = Callbacks() + # Write mutation results + keys = ( + "metrics/precision", + "metrics/recall", + "metrics/mAP_0.5", + "metrics/mAP_0.5:0.95", + "val/box_loss", + "val/obj_loss", + "val/cls_loss", + ) + print_mutation(keys, results, hyp.copy(), save_dir, opt.bucket) + fitness_scores.append(results[2]) + + # Select the fittest individuals for reproduction using adaptive tournament selection + selected_indices = [] + for _ in range(pop_size - elite_size): + # Adaptive tournament size + tournament_size = max( + max(2, tournament_size_min), + int(min(tournament_size_max, pop_size) - (generation / (opt.evolve / 10))), + ) + # Perform tournament selection to choose the best individual + tournament_indices = random.sample(range(pop_size), tournament_size) + tournament_fitness = [fitness_scores[j] for j in tournament_indices] + winner_index = tournament_indices[tournament_fitness.index(max(tournament_fitness))] + selected_indices.append(winner_index) + + # Add the elite individuals to the selected indices + elite_indices = [i for i in range(pop_size) if fitness_scores[i] in sorted(fitness_scores)[-elite_size:]] + selected_indices.extend(elite_indices) + # Create the next generation through crossover and mutation + next_generation = [] + for _ in range(pop_size): + parent1_index = selected_indices[random.randint(0, pop_size - 1)] + parent2_index = selected_indices[random.randint(0, pop_size - 1)] + # Adaptive crossover rate + crossover_rate = max( + crossover_rate_min, min(crossover_rate_max, crossover_rate_max - (generation / opt.evolve)) + ) + if random.uniform(0, 1) < crossover_rate: + crossover_point = random.randint(1, len(hyp_GA) - 1) + child = population[parent1_index][:crossover_point] + population[parent2_index][crossover_point:] + else: + child = population[parent1_index] + # Adaptive mutation rate + mutation_rate = max( + mutation_rate_min, min(mutation_rate_max, mutation_rate_max - (generation / opt.evolve)) + ) + for j in range(len(hyp_GA)): + if random.uniform(0, 1) < mutation_rate: + child[j] += random.uniform(-0.1, 0.1) + child[j] = min(max(child[j], gene_ranges[j][0]), gene_ranges[j][1]) + next_generation.append(child) + # Replace the old population with the new generation + population = next_generation + # Print the best solution found + best_index = fitness_scores.index(max(fitness_scores)) + best_individual = population[best_index] + print("Best solution found:", best_individual) + # Plot results + plot_evolve(evolve_csv) + LOGGER.info( + f'Hyperparameter evolution finished {opt.evolve} generations\n' + f"Results saved to {colorstr('bold', save_dir)}\n" + f'Usage example: $ python train.py --hyp {evolve_yaml}' + ) + + +def generate_individual(input_ranges, individual_length): + """ + Generate an individual with random hyperparameters within specified ranges. + + Args: + input_ranges (list[tuple[float, float]]): List of tuples where each tuple contains the lower and upper bounds + for the corresponding gene (hyperparameter). + individual_length (int): The number of genes (hyperparameters) in the individual. + + Returns: + list[float]: A list representing a generated individual with random gene values within the specified ranges. + + Example: + ```python + input_ranges = [(0.01, 0.1), (0.1, 1.0), (0.9, 2.0)] + individual_length = 3 + individual = generate_individual(input_ranges, individual_length) + print(individual) # Output: [0.035, 0.678, 1.456] (example output) + ``` + + Note: + The individual returned will have a length equal to `individual_length`, with each gene value being a floating-point + number within its specified range in `input_ranges`. + """ + individual = [] + for i in range(individual_length): + lower_bound, upper_bound = input_ranges[i] + individual.append(random.uniform(lower_bound, upper_bound)) + return individual + + +def run(**kwargs): + """ + Execute YOLOv5 training with specified options, allowing optional overrides through keyword arguments. + + Args: + weights (str, optional): Path to initial weights. Defaults to ROOT / 'yolov5s.pt'. + cfg (str, optional): Path to model YAML configuration. Defaults to an empty string. + data (str, optional): Path to dataset YAML configuration. Defaults to ROOT / 'data/coco128.yaml'. + hyp (str, optional): Path to hyperparameters YAML configuration. Defaults to ROOT / 'data/hyps/hyp.scratch-low.yaml'. + epochs (int, optional): Total number of training epochs. Defaults to 100. + batch_size (int, optional): Total batch size for all GPUs. Use -1 for automatic batch size determination. Defaults to 16. + imgsz (int, optional): Image size (pixels) for training and validation. Defaults to 640. + rect (bool, optional): Use rectangular training. Defaults to False. + resume (bool | str, optional): Resume most recent training with an optional path. Defaults to False. + nosave (bool, optional): Only save the final checkpoint. Defaults to False. + noval (bool, optional): Only validate at the final epoch. Defaults to False. + noautoanchor (bool, optional): Disable AutoAnchor. Defaults to False. + noplots (bool, optional): Do not save plot files. Defaults to False. + evolve (int, optional): Evolve hyperparameters for a specified number of generations. Use 300 if provided without a + value. + evolve_population (str, optional): Directory for loading population during evolution. Defaults to ROOT / 'data/ hyps'. + resume_evolve (str, optional): Resume hyperparameter evolution from the last generation. Defaults to None. + bucket (str, optional): gsutil bucket for saving checkpoints. Defaults to an empty string. + cache (str, optional): Cache image data in 'ram' or 'disk'. Defaults to None. + image_weights (bool, optional): Use weighted image selection for training. Defaults to False. + device (str, optional): CUDA device identifier, e.g., '0', '0,1,2,3', or 'cpu'. Defaults to an empty string. + multi_scale (bool, optional): Use multi-scale training, varying image size by ±50%. Defaults to False. + single_cls (bool, optional): Train with multi-class data as single-class. Defaults to False. + optimizer (str, optional): Optimizer type, choices are ['SGD', 'Adam', 'AdamW']. Defaults to 'SGD'. + sync_bn (bool, optional): Use synchronized BatchNorm, only available in DDP mode. Defaults to False. + workers (int, optional): Maximum dataloader workers per rank in DDP mode. Defaults to 8. + project (str, optional): Directory for saving training runs. Defaults to ROOT / 'runs/train'. + name (str, optional): Name for saving the training run. Defaults to 'exp'. + exist_ok (bool, optional): Allow existing project/name without incrementing. Defaults to False. + quad (bool, optional): Use quad dataloader. Defaults to False. + cos_lr (bool, optional): Use cosine learning rate scheduler. Defaults to False. + label_smoothing (float, optional): Label smoothing epsilon value. Defaults to 0.0. + patience (int, optional): Patience for early stopping, measured in epochs without improvement. Defaults to 100. + freeze (list, optional): Layers to freeze, e.g., backbone=10, first 3 layers = [0, 1, 2]. Defaults to [0]. + save_period (int, optional): Frequency in epochs to save checkpoints. Disabled if < 1. Defaults to -1. + seed (int, optional): Global training random seed. Defaults to 0. + local_rank (int, optional): Automatic DDP Multi-GPU argument. Do not modify. Defaults to -1. + + Returns: + None: The function initiates YOLOv5 training or hyperparameter evolution based on the provided options. + + Examples: + ```python + import train + train.run(data='coco128.yaml', imgsz=320, weights='yolov5m.pt') + ``` + + Notes: + - Models: https://github.com/ultralytics/yolov5/tree/master/models + - Datasets: https://github.com/ultralytics/yolov5/tree/master/data + - Tutorial: https://docs.ultralytics.com/yolov5/tutorials/train_custom_data + """ + opt = parse_opt(True) + for k, v in kwargs.items(): + setattr(opt, k, v) + main(opt) + return opt + + +if __name__ == "__main__": + opt = parse_opt() + main(opt) diff --git a/YOLO/yolov5-master/val.py b/YOLO/yolov5-master/val.py new file mode 100644 index 0000000..d792e6a --- /dev/null +++ b/YOLO/yolov5-master/val.py @@ -0,0 +1,604 @@ +# Ultralytics YOLOv5 🚀, AGPL-3.0 license +""" +Validate a trained YOLOv5 detection model on a detection dataset. + +Usage: + $ python val.py --weights yolov5s.pt --data coco128.yaml --img 640 + +Usage - formats: + $ python val.py --weights yolov5s.pt # PyTorch + yolov5s.torchscript # TorchScript + yolov5s.onnx # ONNX Runtime or OpenCV DNN with --dnn + yolov5s_openvino_model # OpenVINO + yolov5s.engine # TensorRT + yolov5s.mlpackage # CoreML (macOS-only) + yolov5s_saved_model # TensorFlow SavedModel + yolov5s.pb # TensorFlow GraphDef + yolov5s.tflite # TensorFlow Lite + yolov5s_edgetpu.tflite # TensorFlow Edge TPU + yolov5s_paddle_model # PaddlePaddle +""" + +import argparse +import json +import os +import subprocess +import sys +from pathlib import Path + +import numpy as np +import torch +from tqdm import tqdm + +FILE = Path(__file__).resolve() +ROOT = FILE.parents[0] # YOLOv5 root directory +if str(ROOT) not in sys.path: + sys.path.append(str(ROOT)) # add ROOT to PATH +ROOT = Path(os.path.relpath(ROOT, Path.cwd())) # relative + +from models.common import DetectMultiBackend +from utils.callbacks import Callbacks +from utils.dataloaders import create_dataloader +from utils.general import ( + LOGGER, + TQDM_BAR_FORMAT, + Profile, + check_dataset, + check_img_size, + check_requirements, + check_yaml, + coco80_to_coco91_class, + colorstr, + increment_path, + non_max_suppression, + print_args, + scale_boxes, + xywh2xyxy, + xyxy2xywh, +) +from utils.metrics import ConfusionMatrix, ap_per_class, box_iou +from utils.plots import output_to_target, plot_images, plot_val_study +from utils.torch_utils import select_device, smart_inference_mode + + +def save_one_txt(predn, save_conf, shape, file): + """ + Saves one detection result to a txt file in normalized xywh format, optionally including confidence. + + Args: + predn (torch.Tensor): Predicted bounding boxes and associated confidence scores and classes in xyxy format, tensor + of shape (N, 6) where N is the number of detections. + save_conf (bool): If True, saves the confidence scores along with the bounding box coordinates. + shape (tuple): Shape of the original image as (height, width). + file (str | Path): File path where the result will be saved. + + Returns: + None + + Notes: + The xyxy bounding box format represents the coordinates (xmin, ymin, xmax, ymax). + The xywh format represents the coordinates (center_x, center_y, width, height) and is normalized by the width and + height of the image. + + Example: + ```python + predn = torch.tensor([[10, 20, 30, 40, 0.9, 1]]) # example prediction + save_one_txt(predn, save_conf=True, shape=(640, 480), file="output.txt") + ``` + """ + gn = torch.tensor(shape)[[1, 0, 1, 0]] # normalization gain whwh + 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") + + +def save_one_json(predn, jdict, path, class_map): + """ + Saves a single JSON detection result, including image ID, category ID, bounding box, and confidence score. + + Args: + predn (torch.Tensor): Predicted detections in xyxy format with shape (n, 6) where n is the number of detections. + The tensor should contain [x_min, y_min, x_max, y_max, confidence, class_id] for each detection. + jdict (list[dict]): List to collect JSON formatted detection results. + path (pathlib.Path): Path object of the image file, used to extract image_id. + class_map (dict[int, int]): Mapping from model class indices to dataset-specific category IDs. + + Returns: + None: Appends detection results as dictionaries to `jdict` list in-place. + + Example: + ```python + predn = torch.tensor([[100, 50, 200, 150, 0.9, 0], [50, 30, 100, 80, 0.8, 1]]) + jdict = [] + path = Path("42.jpg") + class_map = {0: 18, 1: 19} + save_one_json(predn, jdict, path, class_map) + ``` + This will append to `jdict`: + ``` + [ + {'image_id': 42, 'category_id': 18, 'bbox': [125.0, 75.0, 100.0, 100.0], 'score': 0.9}, + {'image_id': 42, 'category_id': 19, 'bbox': [75.0, 55.0, 50.0, 50.0], 'score': 0.8} + ] + ``` + + Notes: + The `bbox` values are formatted as [x, y, width, height], where x and y represent the top-left corner of the box. + """ + image_id = int(path.stem) if path.stem.isnumeric() else path.stem + 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), + } + ) + + +def process_batch(detections, labels, iouv): + """ + Return a correct prediction matrix given detections and labels at various IoU thresholds. + + Args: + detections (np.ndarray): Array of shape (N, 6) where each row corresponds to a detection with format + [x1, y1, x2, y2, conf, class]. + labels (np.ndarray): Array of shape (M, 5) where each row corresponds to a ground truth label with format + [class, x1, y1, x2, y2]. + iouv (np.ndarray): Array of IoU thresholds to evaluate at. + + Returns: + correct (np.ndarray): A binary array of shape (N, len(iouv)) indicating whether each detection is a true positive + for each IoU threshold. There are 10 IoU levels used in the evaluation. + + Example: + ```python + detections = np.array([[50, 50, 200, 200, 0.9, 1], [30, 30, 150, 150, 0.7, 0]]) + labels = np.array([[1, 50, 50, 200, 200]]) + iouv = np.linspace(0.5, 0.95, 10) + correct = process_batch(detections, labels, iouv) + ``` + + Notes: + - This function is used as part of the evaluation pipeline for object detection models. + - IoU (Intersection over Union) is a common evaluation metric for object detection performance. + """ + correct = np.zeros((detections.shape[0], iouv.shape[0])).astype(bool) + iou = box_iou(labels[:, 1:], detections[:, :4]) + correct_class = labels[:, 0:1] == detections[:, 5] + for i in range(len(iouv)): + x = torch.where((iou >= iouv[i]) & correct_class) # IoU > threshold and classes match + if x[0].shape[0]: + matches = torch.cat((torch.stack(x, 1), iou[x[0], x[1]][:, None]), 1).cpu().numpy() # [label, detect, iou] + if x[0].shape[0] > 1: + matches = matches[matches[:, 2].argsort()[::-1]] + matches = matches[np.unique(matches[:, 1], return_index=True)[1]] + # matches = matches[matches[:, 2].argsort()[::-1]] + matches = matches[np.unique(matches[:, 0], return_index=True)[1]] + correct[matches[:, 1].astype(int), i] = True + return torch.tensor(correct, dtype=torch.bool, device=iouv.device) + + +@smart_inference_mode() +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 + max_det=300, # maximum detections per image + task="val", # train, val, test, speed or study + device="", # cuda device, i.e. 0 or 0,1,2,3 or cpu + workers=8, # max dataloader workers (per RANK in DDP mode) + 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 + dnn=False, # use OpenCV DNN for ONNX inference + model=None, + dataloader=None, + save_dir=Path(""), + plots=True, + callbacks=Callbacks(), + compute_loss=None, +): + """ + Evaluates a YOLOv5 model on a dataset and logs performance metrics. + + Args: + data (str | dict): Path to a dataset YAML file or a dataset dictionary. + weights (str | list[str], optional): Path to the model weights file(s). Supports various formats including PyTorch, + TorchScript, ONNX, OpenVINO, TensorRT, CoreML, TensorFlow SavedModel, TensorFlow GraphDef, TensorFlow Lite, + TensorFlow Edge TPU, and PaddlePaddle. + batch_size (int, optional): Batch size for inference. Default is 32. + imgsz (int, optional): Input image size (pixels). Default is 640. + conf_thres (float, optional): Confidence threshold for object detection. Default is 0.001. + iou_thres (float, optional): IoU threshold for Non-Maximum Suppression (NMS). Default is 0.6. + max_det (int, optional): Maximum number of detections per image. Default is 300. + task (str, optional): Task type - 'train', 'val', 'test', 'speed', or 'study'. Default is 'val'. + device (str, optional): Device to use for computation, e.g., '0' or '0,1,2,3' for CUDA or 'cpu' for CPU. Default is ''. + workers (int, optional): Number of dataloader workers. Default is 8. + single_cls (bool, optional): Treat dataset as a single class. Default is False. + augment (bool, optional): Enable augmented inference. Default is False. + verbose (bool, optional): Enable verbose output. Default is False. + save_txt (bool, optional): Save results to *.txt files. Default is False. + save_hybrid (bool, optional): Save label and prediction hybrid results to *.txt files. Default is False. + save_conf (bool, optional): Save confidences in --save-txt labels. Default is False. + save_json (bool, optional): Save a COCO-JSON results file. Default is False. + project (str | Path, optional): Directory to save results. Default is ROOT/'runs/val'. + name (str, optional): Name of the run. Default is 'exp'. + exist_ok (bool, optional): Overwrite existing project/name without incrementing. Default is False. + half (bool, optional): Use FP16 half-precision inference. Default is True. + dnn (bool, optional): Use OpenCV DNN for ONNX inference. Default is False. + model (torch.nn.Module, optional): Model object for training. Default is None. + dataloader (torch.utils.data.DataLoader, optional): Dataloader object. Default is None. + save_dir (Path, optional): Directory to save results. Default is Path(''). + plots (bool, optional): Plot validation images and metrics. Default is True. + callbacks (utils.callbacks.Callbacks, optional): Callbacks for logging and monitoring. Default is Callbacks(). + compute_loss (function, optional): Loss function for training. Default is None. + + Returns: + dict: Contains performance metrics including precision, recall, mAP50, and mAP50-95. + """ + # Initialize/load model and set device + training = model is not None + if training: # called by train.py + device, pt, jit, engine = next(model.parameters()).device, True, False, False # get model device, PyTorch model + half &= device.type != "cpu" # half precision only supported on CUDA + model.half() if half else model.float() + else: # called directly + device = select_device(device, batch_size=batch_size) + + # 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 + + # Load model + model = DetectMultiBackend(weights, device=device, dnn=dnn, data=data, fp16=half) + stride, pt, jit, engine = model.stride, model.pt, model.jit, model.engine + imgsz = check_img_size(imgsz, s=stride) # check image size + half = model.fp16 # FP16 supported on limited backends with CUDA + if engine: + batch_size = model.batch_size + else: + device = model.device + if not (pt or jit): + batch_size = 1 # export.py models default to batch-size 1 + LOGGER.info(f"Forcing --batch-size 1 square inference (1,3,{imgsz},{imgsz}) for non-PyTorch models") + + # Data + data = check_dataset(data) # check + + # Configure + model.eval() + cuda = device.type != "cpu" + is_coco = isinstance(data.get("val"), str) and data["val"].endswith(f"coco{os.sep}val2017.txt") # COCO dataset + nc = 1 if single_cls else int(data["nc"]) # number of classes + iouv = torch.linspace(0.5, 0.95, 10, device=device) # iou vector for mAP@0.5:0.95 + niou = iouv.numel() + + # Dataloader + if not training: + if pt and not single_cls: # check --weights are trained on --data + ncm = model.model.nc + assert ncm == nc, ( + f"{weights} ({ncm} classes) trained on different --data than what you passed ({nc} " + f"classes). Pass correct combination of --weights and --data that are trained together." + ) + model.warmup(imgsz=(1 if pt else batch_size, 3, imgsz, imgsz)) # warmup + pad, rect = (0.0, False) if task == "speed" else (0.5, pt) # square inference for benchmarks + task = task if task in ("train", "val", "test") else "val" # path to train/val/test images + dataloader = create_dataloader( + data[task], + imgsz, + batch_size, + stride, + single_cls, + pad=pad, + rect=rect, + workers=workers, + prefix=colorstr(f"{task}: "), + )[0] + + seen = 0 + confusion_matrix = ConfusionMatrix(nc=nc) + names = model.names if hasattr(model, "names") else model.module.names # get class names + if isinstance(names, (list, tuple)): # old format + names = dict(enumerate(names)) + class_map = coco80_to_coco91_class() if is_coco else list(range(1000)) + s = ("%22s" + "%11s" * 6) % ("Class", "Images", "Instances", "P", "R", "mAP50", "mAP50-95") + tp, fp, p, r, f1, mp, mr, map50, ap50, map = 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0 + dt = Profile(device=device), Profile(device=device), Profile(device=device) # profiling times + loss = torch.zeros(3, device=device) + jdict, stats, ap, ap_class = [], [], [], [] + callbacks.run("on_val_start") + pbar = tqdm(dataloader, desc=s, bar_format=TQDM_BAR_FORMAT) # progress bar + for batch_i, (im, targets, paths, shapes) in enumerate(pbar): + callbacks.run("on_val_batch_start") + with dt[0]: + if cuda: + im = im.to(device, non_blocking=True) + targets = targets.to(device) + im = im.half() if half else im.float() # uint8 to fp16/32 + im /= 255 # 0 - 255 to 0.0 - 1.0 + nb, _, height, width = im.shape # batch size, channels, height, width + + # Inference + with dt[1]: + preds, train_out = model(im) if compute_loss else (model(im, augment=augment), None) + + # Loss + if compute_loss: + loss += compute_loss(train_out, targets)[1] # box, obj, cls + + # NMS + targets[:, 2:] *= torch.tensor((width, height, width, height), device=device) # to pixels + lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_hybrid else [] # for autolabelling + with dt[2]: + preds = non_max_suppression( + preds, conf_thres, iou_thres, labels=lb, multi_label=True, agnostic=single_cls, max_det=max_det + ) + + # Metrics + for si, pred in enumerate(preds): + labels = targets[targets[:, 0] == si, 1:] + nl, npr = labels.shape[0], pred.shape[0] # number of labels, predictions + path, shape = Path(paths[si]), shapes[si][0] + correct = torch.zeros(npr, niou, dtype=torch.bool, device=device) # init + seen += 1 + + if npr == 0: + if nl: + stats.append((correct, *torch.zeros((2, 0), device=device), labels[:, 0])) + if plots: + confusion_matrix.process_batch(detections=None, labels=labels[:, 0]) + continue + + # Predictions + if single_cls: + pred[:, 5] = 0 + predn = pred.clone() + scale_boxes(im[si].shape[1:], predn[:, :4], shape, shapes[si][1]) # native-space pred + + # Evaluate + if nl: + tbox = xywh2xyxy(labels[:, 1:5]) # target boxes + scale_boxes(im[si].shape[1:], tbox, shape, shapes[si][1]) # native-space labels + labelsn = torch.cat((labels[:, 0:1], tbox), 1) # native-space labels + correct = process_batch(predn, labelsn, iouv) + if plots: + confusion_matrix.process_batch(predn, labelsn) + stats.append((correct, pred[:, 4], pred[:, 5], labels[:, 0])) # (correct, conf, pcls, tcls) + + # Save/log + if save_txt: + (save_dir / "labels").mkdir(parents=True, exist_ok=True) + save_one_txt(predn, save_conf, shape, file=save_dir / "labels" / f"{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, im[si]) + + # Plot images + if plots and batch_i < 3: + plot_images(im, targets, paths, save_dir / f"val_batch{batch_i}_labels.jpg", names) # labels + plot_images(im, output_to_target(preds), paths, save_dir / f"val_batch{batch_i}_pred.jpg", names) # pred + + callbacks.run("on_val_batch_end", batch_i, im, targets, paths, shapes, preds) + + # Compute metrics + stats = [torch.cat(x, 0).cpu().numpy() for x in zip(*stats)] # to numpy + if len(stats) and stats[0].any(): + tp, fp, p, r, f1, ap, ap_class = ap_per_class(*stats, plot=plots, save_dir=save_dir, names=names) + ap50, ap = ap[:, 0], ap.mean(1) # AP@0.5, AP@0.5:0.95 + mp, mr, map50, map = p.mean(), r.mean(), ap50.mean(), ap.mean() + nt = np.bincount(stats[3].astype(int), minlength=nc) # number of targets per class + + # Print results + pf = "%22s" + "%11i" * 2 + "%11.3g" * 4 # print format + LOGGER.info(pf % ("all", seen, nt.sum(), mp, mr, map50, map)) + if nt.sum() == 0: + LOGGER.warning(f"WARNING ⚠️ no labels found in {task} set, can not compute metrics without labels") + + # Print results per class + if (verbose or (nc < 50 and not training)) and nc > 1 and len(stats): + for i, c in enumerate(ap_class): + LOGGER.info(pf % (names[c], seen, nt[c], p[i], r[i], ap50[i], ap[i])) + + # Print speeds + t = tuple(x.t / 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) + + # Plots + if plots: + confusion_matrix.plot(save_dir=save_dir, names=list(names.values())) + callbacks.run("on_val_end", nt, tp, fp, p, r, f1, ap, ap50, ap_class, confusion_matrix) + + # 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("../datasets/coco/annotations/instances_val2017.json")) # annotations + if not os.path.exists(anno_json): + anno_json = os.path.join(data["path"], "annotations", "instances_val2017.json") + pred_json = str(save_dir / f"{w}_predictions.json") # predictions + 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>=2.0.6") + 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") + if is_coco: + eval.params.imgIds = [int(Path(x).stem) for x in dataloader.dataset.im_files] # image IDs to evaluate + eval.evaluate() + eval.accumulate() + 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}") + + # 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 "" + LOGGER.info(f"Results saved to {colorstr('bold', save_dir)}{s}") + maps = np.zeros(nc) + map + for i, c in enumerate(ap_class): + maps[c] = ap[i] + return (mp, mr, map50, map, *(loss.cpu() / len(dataloader)).tolist()), maps, t + + +def parse_opt(): + """ + Parse command-line options for configuring YOLOv5 model inference. + + Args: + data (str, optional): Path to the dataset YAML file. Default is 'data/coco128.yaml'. + weights (list[str], optional): List of paths to model weight files. Default is 'yolov5s.pt'. + batch_size (int, optional): Batch size for inference. Default is 32. + imgsz (int, optional): Inference image size in pixels. Default is 640. + conf_thres (float, optional): Confidence threshold for predictions. Default is 0.001. + iou_thres (float, optional): IoU threshold for Non-Max Suppression (NMS). Default is 0.6. + max_det (int, optional): Maximum number of detections per image. Default is 300. + task (str, optional): Task type - options are 'train', 'val', 'test', 'speed', or 'study'. Default is 'val'. + device (str, optional): Device to run the model on. e.g., '0' or '0,1,2,3' or 'cpu'. Default is empty to let the system choose automatically. + workers (int, optional): Maximum number of dataloader workers per rank in DDP mode. Default is 8. + single_cls (bool, optional): If set, treats the dataset as a single-class dataset. Default is False. + augment (bool, optional): If set, performs augmented inference. Default is False. + verbose (bool, optional): If set, reports mAP by class. Default is False. + save_txt (bool, optional): If set, saves results to *.txt files. Default is False. + save_hybrid (bool, optional): If set, saves label+prediction hybrid results to *.txt files. Default is False. + save_conf (bool, optional): If set, saves confidences in --save-txt labels. Default is False. + save_json (bool, optional): If set, saves results to a COCO-JSON file. Default is False. + project (str, optional): Project directory to save results to. Default is 'runs/val'. + name (str, optional): Name of the directory to save results to. Default is 'exp'. + exist_ok (bool, optional): If set, existing directory will not be incremented. Default is False. + half (bool, optional): If set, uses FP16 half-precision inference. Default is False. + dnn (bool, optional): If set, uses OpenCV DNN for ONNX inference. Default is False. + + Returns: + argparse.Namespace: Parsed command-line options. + + Notes: + - The '--data' parameter is checked to ensure it ends with 'coco.yaml' if '--save-json' is set. + - The '--save-txt' option is set to True if '--save-hybrid' is enabled. + - Args are printed using `print_args` to facilitate debugging. + + Example: + To validate a trained YOLOv5 model on a COCO dataset: + ```python + $ python val.py --weights yolov5s.pt --data coco128.yaml --img 640 + ``` + Different model formats could be used instead of `yolov5s.pt`: + ```python + $ python val.py --weights yolov5s.pt yolov5s.torchscript yolov5s.onnx yolov5s_openvino_model yolov5s.engine + ``` + Additional options include saving results in different formats, selecting devices, and more. + """ + parser = argparse.ArgumentParser() + parser.add_argument("--data", type=str, default=ROOT / "data/dimo.yaml", help="dataset.yaml path") + parser.add_argument("--weights", nargs="+", type=str, default=ROOT / "runs/train/exp8/weights/best.pt", help="model path(s)") + parser.add_argument("--batch-size", type=int, default=4, 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("--max-det", type=int, default=300, help="maximum detections per image") + 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("--workers", type=int, default=0, help="max dataloader workers (per RANK in DDP mode)") + 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("--dnn", action="store_true", help="use OpenCV DNN for ONNX inference") + opt = parser.parse_args() + opt.data = check_yaml(opt.data) # check YAML + opt.save_json |= opt.data.endswith("coco.yaml") + opt.save_txt |= opt.save_hybrid + print_args(vars(opt)) + return opt + + +def main(opt): + """ + Executes YOLOv5 tasks like training, validation, testing, speed, and study benchmarks based on provided options. + + Args: + opt (argparse.Namespace): Parsed command-line options. + This includes values for parameters like 'data', 'weights', 'batch_size', 'imgsz', 'conf_thres', + 'iou_thres', 'max_det', 'task', 'device', 'workers', 'single_cls', 'augment', 'verbose', 'save_txt', + 'save_hybrid', 'save_conf', 'save_json', 'project', 'name', 'exist_ok', 'half', and 'dnn', essential + for configuring the YOLOv5 tasks. + + Returns: + None + + Examples: + To validate a trained YOLOv5 model on the COCO dataset with a specific weights file, use: + ```python + $ python val.py --weights yolov5s.pt --data coco128.yaml --img 640 + ``` + """ + check_requirements(ROOT / "requirements.txt", exclude=("tensorboard", "thop")) + + 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 produces invalid results") + if opt.save_hybrid: + LOGGER.info("WARNING ⚠️ --save-hybrid will return high mAP from hybrid labels, not from predictions alone") + run(**vars(opt)) + + else: + weights = opt.weights if isinstance(opt.weights, list) else [opt.weights] + opt.half = torch.cuda.is_available() and opt.device != "cpu" # FP16 for fastest results + if opt.task == "speed": # speed benchmarks + # python val.py --task speed --data coco.yaml --batch 1 --weights yolov5n.pt yolov5s.pt... + opt.conf_thres, opt.iou_thres, opt.save_json = 0.25, 0.45, False + for opt.weights in weights: + run(**vars(opt), plots=False) + + elif opt.task == "study": # speed vs mAP benchmarks + # python val.py --task study --data coco.yaml --iou 0.7 --weights yolov5n.pt yolov5s.pt... + for opt.weights in weights: + f = f"study_{Path(opt.data).stem}_{Path(opt.weights).stem}.txt" # filename to save to + x, y = list(range(256, 1536 + 128, 128)), [] # x axis (image sizes), y axis + for opt.imgsz in x: # img-size + LOGGER.info(f"\nRunning {f} --imgsz {opt.imgsz}...") + r, _, t = run(**vars(opt), plots=False) + y.append(r + t) # results and times + np.savetxt(f, y, fmt="%10.4g") # save + subprocess.run(["zip", "-r", "study.zip", "study_*.txt"]) + plot_val_study(x=x) # plot + else: + raise NotImplementedError(f'--task {opt.task} not in ("train", "val", "test", "speed", "study")') + + +if __name__ == "__main__": + opt = parse_opt() + main(opt)