+

iteach_toolkit/DHYOLO/__init__.py

@@ -0,0 +1 @@
+from .model import DHYOLODetector

iteach_toolkit/DHYOLO/detect.py

@@ -0,0 +1,290 @@
+# YOLOv5 🚀 by Ultralytics, 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.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
+"""
+
+import argparse
+import csv
+import os
+import platform
+import sys
+from pathlib import Path
+from PIL import Image as PILImg
+import shutil
+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_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 / 'inference',  # save results to project/name
+        name='_dhyolo',  # 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
+):
+    source = str(source)
+    save_img = False #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(), Profile(), Profile())
+    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
+
+        # Inference
+        with dt[1]:
+            visualize = increment_path(save_dir / Path(path).stem, mkdir=True) if visualize else False
+            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):
+            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 += '%gx%g ' % 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
+                        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(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)
+    
+    # need to remove as we are not saving anything
+    shutil.rmtree(project, ignore_errors=True)
+
+    return pred
+
+def parse_opt():
+    parser = argparse.ArgumentParser()
+    parser.add_argument('--weights', nargs='+', type=str, default=ROOT / 'yolov5s.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/coco128.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-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):
+    check_requirements(ROOT / 'requirements.txt', exclude=('tensorboard', 'thop'))
+    run(**vars(opt))
+
+
+if __name__ == '__main__':
+    opt = parse_opt()
+    main(opt)

iteach_toolkit/DHYOLO/export.py

@@ -0,0 +1,1537 @@
+# 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):
+    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():
+    """
+    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.
+            Parameters:
+               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"
+        if half:
+            precision = ct.precision.FLOAT16
+        else:
+            precision = 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, 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.
+        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
+    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())
+    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
+
+    if mlmodel:
+        f = file.with_suffix(".mlmodel")  # filename
+    else:
+        f = file.with_suffix(".mlpackage")  # filename
+
+    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
+    if mlmodel:
+        weights_dir = None
+    else:
+        weights_dir = 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
+    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.
+        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,
+            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
+    im = torch.zeros(batch_size, 3, *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)
+    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("--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)

iteach_toolkit/DHYOLO/hubconf.py

@@ -0,0 +1,510 @@
+# Ultralytics YOLOv5 🚀, AGPL-3.0 license
+"""
+PyTorch Hub models https://pytorch.org/hub/ultralytics_yolov5
+
+Usage:
+    import torch
+    model = torch.hub.load('ultralytics/yolov5', 'yolov5s')  # official model
+    model = torch.hub.load('ultralytics/yolov5:master', 'yolov5s')  # from branch
+    model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5s.pt')  # custom/local model
+    model = torch.hub.load('.', 'custom', 'yolov5s.pt', source='local')  # local repo
+"""
+
+import torch
+
+
+def _create(name, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True, device=None):
+    """
+    Creates or loads a YOLOv5 model, with options for pretrained weights and model customization.
+
+    Args:
+        name (str): Model name (e.g., 'yolov5s') or path to the model checkpoint (e.g., 'path/to/best.pt').
+        pretrained (bool, optional): If True, loads pretrained weights into the model. Defaults to True.
+        channels (int, optional): Number of input channels the model expects. Defaults to 3.
+        classes (int, optional): Number of classes the model is expected to detect. Defaults to 80.
+        autoshape (bool, optional): If True, applies the YOLOv5 .autoshape() wrapper for various input formats. Defaults to True.
+        verbose (bool, optional): If True, prints detailed information during the model creation/loading process. Defaults to True.
+        device (str | torch.device | None, optional): Device to use for model parameters (e.g., 'cpu', 'cuda'). If None, selects
+            the best available device. Defaults to None.
+
+    Returns:
+        (DetectMultiBackend | AutoShape): The loaded YOLOv5 model, potentially wrapped with AutoShape if specified.
+
+    Examples:
+        ```python
+        import torch
+        from ultralytics import _create
+
+        # Load an official YOLOv5s model with pretrained weights
+        model = _create('yolov5s')
+
+        # Load a custom model from a local checkpoint
+        model = _create('path/to/custom_model.pt', pretrained=False)
+
+        # Load a model with specific input channels and classes
+        model = _create('yolov5s', channels=1, classes=10)
+        ```
+
+    Notes:
+        For more information on model loading and customization, visit the
+        [YOLOv5 PyTorch Hub Documentation](https://docs.ultralytics.com/yolov5/tutorials/pytorch_hub_model_loading).
+    """
+    from pathlib import Path
+
+    from models.common import AutoShape, DetectMultiBackend
+    from models.experimental import attempt_load
+    from models.yolo import ClassificationModel, DetectionModel, SegmentationModel
+    from utils.downloads import attempt_download
+    from utils.general import LOGGER, ROOT, check_requirements, intersect_dicts, logging
+    from utils.torch_utils import select_device
+
+    if not verbose:
+        LOGGER.setLevel(logging.WARNING)
+    check_requirements(ROOT / "requirements.txt", exclude=("opencv-python", "tensorboard", "thop"))
+    name = Path(name)
+    path = name.with_suffix(".pt") if name.suffix == "" and not name.is_dir() else name  # checkpoint path
+    try:
+        device = select_device(device)
+        if pretrained and channels == 3 and classes == 80:
+            try:
+                model = DetectMultiBackend(path, device=device, fuse=autoshape)  # detection model
+                if autoshape:
+                    if model.pt and isinstance(model.model, ClassificationModel):
+                        LOGGER.warning(
+                            "WARNING ⚠️ YOLOv5 ClassificationModel is not yet AutoShape compatible. "
+                            "You must pass torch tensors in BCHW to this model, i.e. shape(1,3,224,224)."
+                        )
+                    elif model.pt and isinstance(model.model, SegmentationModel):
+                        LOGGER.warning(
+                            "WARNING ⚠️ YOLOv5 SegmentationModel is not yet AutoShape compatible. "
+                            "You will not be able to run inference with this model."
+                        )
+                    else:
+                        model = AutoShape(model)  # for file/URI/PIL/cv2/np inputs and NMS
+            except Exception:
+                model = attempt_load(path, device=device, fuse=False)  # arbitrary model
+        else:
+            cfg = list((Path(__file__).parent / "models").rglob(f"{path.stem}.yaml"))[0]  # model.yaml path
+            model = DetectionModel(cfg, channels, classes)  # create model
+            if pretrained:
+                ckpt = torch.load(attempt_download(path), map_location=device)  # load
+                csd = ckpt["model"].float().state_dict()  # checkpoint state_dict as FP32
+                csd = intersect_dicts(csd, model.state_dict(), exclude=["anchors"])  # intersect
+                model.load_state_dict(csd, strict=False)  # load
+                if len(ckpt["model"].names) == classes:
+                    model.names = ckpt["model"].names  # set class names attribute
+        if not verbose:
+            LOGGER.setLevel(logging.INFO)  # reset to default
+        return model.to(device)
+
+    except Exception as e:
+        help_url = "https://docs.ultralytics.com/yolov5/tutorials/pytorch_hub_model_loading"
+        s = f"{e}. Cache may be out of date, try `force_reload=True` or see {help_url} for help."
+        raise Exception(s) from e
+
+
+def custom(path="path/to/model.pt", autoshape=True, _verbose=True, device=None):
+    """
+    Loads a custom or local YOLOv5 model from a given path with optional autoshaping and device specification.
+
+    Args:
+        path (str): Path to the custom model file (e.g., 'path/to/model.pt').
+        autoshape (bool): Apply YOLOv5 .autoshape() wrapper to model if True, enabling compatibility with various input
+            types (default is True).
+        _verbose (bool): If True, prints all informational messages to the screen; otherwise, operates silently
+            (default is True).
+        device (str | torch.device | None): Device to load the model on, e.g., 'cpu', 'cuda', torch.device('cuda:0'), etc.
+            (default is None, which automatically selects the best available device).
+
+    Returns:
+        torch.nn.Module: A YOLOv5 model loaded with the specified parameters.
+
+    Notes:
+        For more details on loading models from PyTorch Hub:
+        https://docs.ultralytics.com/yolov5/tutorials/pytorch_hub_model_loading
+
+    Examples:
+        ```python
+        # Load model from a given path with autoshape enabled on the best available device
+        model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5s.pt')
+
+        # Load model from a local path without autoshape on the CPU device
+        model = torch.hub.load('.', 'custom', 'yolov5s.pt', source='local', autoshape=False, device='cpu')
+        ```
+    """
+    return _create(path, autoshape=autoshape, verbose=_verbose, device=device)
+
+
+def yolov5n(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
+    """
+    Instantiates the YOLOv5-nano model with options for pretraining, input channels, class count, autoshaping,
+    verbosity, and device.
+
+    Args:
+        pretrained (bool): If True, loads pretrained weights into the model. Defaults to True.
+        channels (int): Number of input channels for the model. Defaults to 3.
+        classes (int): Number of classes for object detection. Defaults to 80.
+        autoshape (bool): If True, applies the YOLOv5 .autoshape() wrapper to the model for various formats (file/URI/PIL/
+            cv2/np) and non-maximum suppression (NMS) during inference. Defaults to True.
+        _verbose (bool): If True, prints detailed information to the screen. Defaults to True.
+        device (str | torch.device | None): Specifies the device to use for model computation. If None, uses the best device
+            available (i.e., GPU if available, otherwise CPU). Defaults to None.
+
+    Returns:
+        DetectionModel | ClassificationModel | SegmentationModel: The instantiated YOLOv5-nano model, potentially with
+            pretrained weights and autoshaping applied.
+
+    Notes:
+        For further details on loading models from PyTorch Hub, refer to [PyTorch Hub models](https://pytorch.org/hub/
+        ultralytics_yolov5).
+
+    Examples:
+        ```python
+        import torch
+        from ultralytics import yolov5n
+
+        # Load the YOLOv5-nano model with defaults
+        model = yolov5n()
+
+        # Load the YOLOv5-nano model with a specific device
+        model = yolov5n(device='cuda')
+        ```
+    """
+    return _create("yolov5n", pretrained, channels, classes, autoshape, _verbose, device)
+
+
+def yolov5s(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
+    """
+    Create a YOLOv5-small (yolov5s) model with options for pretraining, input channels, class count, autoshaping,
+    verbosity, and device configuration.
+
+    Args:
+        pretrained (bool, optional): Flag to load pretrained weights into the model. Defaults to True.
+        channels (int, optional): Number of input channels. Defaults to 3.
+        classes (int, optional): Number of model classes. Defaults to 80.
+        autoshape (bool, optional): Whether to wrap the model with YOLOv5's .autoshape() for handling various input formats.
+            Defaults to True.
+        _verbose (bool, optional): Flag to print detailed information regarding model loading. Defaults to True.
+        device (str | torch.device | None, optional): Device to use for model computation, can be 'cpu', 'cuda', or
+            torch.device instances. If None, automatically selects the best available device. Defaults to None.
+
+    Returns:
+        torch.nn.Module: The YOLOv5-small model configured and loaded according to the specified parameters.
+
+    Example:
+        ```python
+        import torch
+
+        # Load the official YOLOv5-small model with pretrained weights
+        model = torch.hub.load('ultralytics/yolov5', 'yolov5s')
+
+        # Load the YOLOv5-small model from a specific branch
+        model = torch.hub.load('ultralytics/yolov5:master', 'yolov5s')
+
+        # Load a custom YOLOv5-small model from a local checkpoint
+        model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5s.pt')
+
+        # Load a local YOLOv5-small model specifying source as local repository
+        model = torch.hub.load('.', 'custom', 'yolov5s.pt', source='local')
+        ```
+
+    Notes:
+        For more details on model loading and customization, visit
+        the [YOLOv5 PyTorch Hub Documentation](https://pytorch.org/hub/ultralytics_yolov5).
+    """
+    return _create("yolov5s", pretrained, channels, classes, autoshape, _verbose, device)
+
+
+def yolov5m(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
+    """
+    Instantiates the YOLOv5-medium model with customizable pretraining, channel count, class count, autoshaping,
+    verbosity, and device.
+
+    Args:
+        pretrained (bool, optional): Whether to load pretrained weights into the model. Default is True.
+        channels (int, optional): Number of input channels. Default is 3.
+        classes (int, optional): Number of model classes. Default is 80.
+        autoshape (bool, optional): Apply YOLOv5 .autoshape() wrapper to the model for handling various input formats.
+            Default is True.
+        _verbose (bool, optional): Whether to print detailed information to the screen. Default is True.
+        device (str | torch.device | None, optional): Device specification to use for model parameters (e.g., 'cpu', 'cuda').
+            Default is None.
+
+    Returns:
+        torch.nn.Module: The instantiated YOLOv5-medium model.
+
+    Usage Example:
+        ```python
+        import torch
+
+        model = torch.hub.load('ultralytics/yolov5', 'yolov5m')  # Load YOLOv5-medium from Ultralytics repository
+        model = torch.hub.load('ultralytics/yolov5:master', 'yolov5m')  # Load from the master branch
+        model = torch.hub.load('ultralytics/yolov5', 'custom', 'yolov5m.pt')  # Load a custom/local YOLOv5-medium model
+        model = torch.hub.load('.', 'custom', 'yolov5m.pt', source='local')  # Load from a local repository
+        ```
+
+    For more information, visit https://pytorch.org/hub/ultralytics_yolov5.
+    """
+    return _create("yolov5m", pretrained, channels, classes, autoshape, _verbose, device)
+
+
+def yolov5l(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
+    """
+    Creates YOLOv5-large model with options for pretraining, channels, classes, autoshaping, verbosity, and device
+    selection.
+
+    Args:
+        pretrained (bool): Load pretrained weights into the model. Default is True.
+        channels (int): Number of input channels. Default is 3.
+        classes (int): Number of model classes. Default is 80.
+        autoshape (bool): Apply YOLOv5 .autoshape() wrapper to model. Default is True.
+        _verbose (bool): Print all information to screen. Default is True.
+        device (str | torch.device | None): Device to use for model parameters, e.g., 'cpu', 'cuda', or a torch.device instance.
+            Default is None.
+
+    Returns:
+        YOLOv5 model (torch.nn.Module): The YOLOv5-large model instantiated with specified configurations and possibly
+        pretrained weights.
+
+    Examples:
+        ```python
+        import torch
+        model = torch.hub.load('ultralytics/yolov5', 'yolov5l')
+        ```
+
+    Notes:
+        For additional details, refer to the PyTorch Hub models documentation:
+        https://pytorch.org/hub/ultralytics_yolov5
+    """
+    return _create("yolov5l", pretrained, channels, classes, autoshape, _verbose, device)
+
+
+def yolov5x(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
+    """
+    Perform object detection using the YOLOv5-xlarge model with options for pretraining, input channels, class count,
+    autoshaping, verbosity, and device specification.
+
+    Args:
+        pretrained (bool): If True, loads pretrained weights into the model. Defaults to True.
+        channels (int): Number of input channels for the model. Defaults to 3.
+        classes (int): Number of model classes for object detection. Defaults to 80.
+        autoshape (bool): If True, applies the YOLOv5 .autoshape() wrapper for handling different input formats. Defaults to
+            True.
+        _verbose (bool): If True, prints detailed information during model loading. Defaults to True.
+        device (str | torch.device | None): Device specification for computing the model, e.g., 'cpu', 'cuda:0', torch.device('cuda').
+            Defaults to None.
+
+    Returns:
+        torch.nn.Module: The YOLOv5-xlarge model loaded with the specified parameters, optionally with pretrained weights and
+        autoshaping applied.
+
+    Example:
+        ```python
+        import torch
+        model = torch.hub.load('ultralytics/yolov5', 'yolov5x')
+        ```
+
+    For additional details, refer to the official YOLOv5 PyTorch Hub models documentation:
+    https://pytorch.org/hub/ultralytics_yolov5
+    """
+    return _create("yolov5x", pretrained, channels, classes, autoshape, _verbose, device)
+
+
+def yolov5n6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
+    """
+    Creates YOLOv5-nano-P6 model with options for pretraining, channels, classes, autoshaping, verbosity, and device.
+
+    Args:
+        pretrained (bool, optional): If True, loads pretrained weights into the model. Default is True.
+        channels (int, optional): Number of input channels. Default is 3.
+        classes (int, optional): Number of model classes. Default is 80.
+        autoshape (bool, optional): If True, applies the YOLOv5 .autoshape() wrapper to the model. Default is True.
+        _verbose (bool, optional): If True, prints all information to screen. Default is True.
+        device (str | torch.device | None, optional): Device to use for model parameters. Can be 'cpu', 'cuda', or None.
+            Default is None.
+
+    Returns:
+        torch.nn.Module: YOLOv5-nano-P6 model loaded with the specified configurations.
+
+    Example:
+        ```python
+        import torch
+        model = yolov5n6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device='cuda')
+        ```
+
+    Notes:
+        For more information on PyTorch Hub models, visit: https://pytorch.org/hub/ultralytics_yolov5
+    """
+    return _create("yolov5n6", pretrained, channels, classes, autoshape, _verbose, device)
+
+
+def yolov5s6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
+    """
+    Instantiate the YOLOv5-small-P6 model with options for pretraining, input channels, number of classes, autoshaping,
+    verbosity, and device selection.
+
+    Args:
+        pretrained (bool): If True, loads pretrained weights. Default is True.
+        channels (int): Number of input channels. Default is 3.
+        classes (int): Number of object detection classes. Default is 80.
+        autoshape (bool): If True, applies YOLOv5 .autoshape() wrapper to the model, allowing for varied input formats.
+            Default is True.
+        _verbose (bool): If True, prints detailed information during model loading. Default is True.
+        device (str | torch.device | None): Device specification for model parameters (e.g., 'cpu', 'cuda', or torch.device).
+            Default is None, which selects an available device automatically.
+
+    Returns:
+        torch.nn.Module: The YOLOv5-small-P6 model instance.
+
+    Usage:
+        ```python
+        import torch
+
+        model = torch.hub.load('ultralytics/yolov5', 'yolov5s6')
+        model = torch.hub.load('ultralytics/yolov5:master', 'yolov5s6')  # load from a specific branch
+        model = torch.hub.load('ultralytics/yolov5', 'custom', 'path/to/yolov5s6.pt')  # custom/local model
+        model = torch.hub.load('.', 'custom', 'path/to/yolov5s6.pt', source='local')  # local repo model
+        ```
+
+    Notes:
+        - For more information, refer to the PyTorch Hub models documentation at https://pytorch.org/hub/ultralytics_yolov5
+
+    Raises:
+        Exception: If there is an error during model creation or loading, with a suggestion to visit the YOLOv5
+            tutorials for help.
+    """
+    return _create("yolov5s6", pretrained, channels, classes, autoshape, _verbose, device)
+
+
+def yolov5m6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
+    """
+    Create YOLOv5-medium-P6 model with options for pretraining, channel count, class count, autoshaping, verbosity, and
+    device.
+
+    Args:
+        pretrained (bool): If True, loads pretrained weights. Default is True.
+        channels (int): Number of input channels. Default is 3.
+        classes (int): Number of model classes. Default is 80.
+        autoshape (bool): Apply YOLOv5 .autoshape() wrapper to the model for file/URI/PIL/cv2/np inputs and NMS.
+            Default is True.
+        _verbose (bool): If True, prints detailed information to the screen. Default is True.
+        device (str | torch.device | None): Device to use for model parameters. Default is None, which uses the
+            best available device.
+
+    Returns:
+        torch.nn.Module: The YOLOv5-medium-P6 model.
+
+    Refer to the PyTorch Hub models documentation: https://pytorch.org/hub/ultralytics_yolov5 for additional details.
+
+    Example:
+        ```python
+        import torch
+
+        # Load YOLOv5-medium-P6 model
+        model = torch.hub.load('ultralytics/yolov5', 'yolov5m6')
+        ```
+
+    Notes:
+        - The model can be loaded with pre-trained weights for better performance on specific tasks.
+        - The autoshape feature simplifies input handling by allowing various popular data formats.
+    """
+    return _create("yolov5m6", pretrained, channels, classes, autoshape, _verbose, device)
+
+
+def yolov5l6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
+    """
+    Instantiate the YOLOv5-large-P6 model with options for pretraining, channel and class counts, autoshaping,
+    verbosity, and device selection.
+
+    Args:
+        pretrained (bool, optional): If True, load pretrained weights into the model. Default is True.
+        channels (int, optional): Number of input channels. Default is 3.
+        classes (int, optional): Number of model classes. Default is 80.
+        autoshape (bool, optional): If True, apply YOLOv5 .autoshape() wrapper to the model for input flexibility. Default is True.
+        _verbose (bool, optional): If True, print all information to the screen. Default is True.
+        device (str | torch.device | None, optional): Device to use for model parameters, e.g., 'cpu', 'cuda', or torch.device.
+            If None, automatically selects the best available device. Default is None.
+
+    Returns:
+        torch.nn.Module: The instantiated YOLOv5-large-P6 model.
+
+    Example:
+        ```python
+        import torch
+        model = torch.hub.load('ultralytics/yolov5', 'yolov5l6')  # official model
+        model = torch.hub.load('ultralytics/yolov5:master', 'yolov5l6')  # from specific branch
+        model = torch.hub.load('ultralytics/yolov5', 'custom', 'path/to/yolov5l6.pt')  # custom/local model
+        model = torch.hub.load('.', 'custom', 'path/to/yolov5l6.pt', source='local')  # local repository
+        ```
+
+    Note:
+        Refer to [PyTorch Hub Documentation](https://pytorch.org/hub/ultralytics_yolov5) for additional usage instructions.
+    """
+    return _create("yolov5l6", pretrained, channels, classes, autoshape, _verbose, device)
+
+
+def yolov5x6(pretrained=True, channels=3, classes=80, autoshape=True, _verbose=True, device=None):
+    """
+    Creates the YOLOv5-xlarge-P6 model with options for pretraining, number of input channels, class count, autoshaping,
+    verbosity, and device selection.
+
+    Args:
+        pretrained (bool): If True, loads pretrained weights into the model. Default is True.
+        channels (int): Number of input channels. Default is 3.
+        classes (int): Number of model classes. Default is 80.
+        autoshape (bool): If True, applies YOLOv5 .autoshape() wrapper to the model. Default is True.
+        _verbose (bool): If True, prints all information to the screen. Default is True.
+        device (str | torch.device | None): Device to use for model parameters, can be a string, torch.device object, or
+            None for default device selection. Default is None.
+
+    Returns:
+        torch.nn.Module: The instantiated YOLOv5-xlarge-P6 model.
+
+    Example:
+        ```python
+        import torch
+        model = torch.hub.load('ultralytics/yolov5', 'yolov5x6')  # load the YOLOv5-xlarge-P6 model
+        ```
+
+    Note:
+        For more information on YOLOv5 models, visit the official documentation:
+        https://docs.ultralytics.com/yolov5
+    """
+    return _create("yolov5x6", pretrained, channels, classes, autoshape, _verbose, device)
+
+
+if __name__ == "__main__":
+    import argparse
+    from pathlib import Path
+
+    import numpy as np
+    from PIL import Image
+
+    from utils.general import cv2, print_args
+
+    # Argparser
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--model", type=str, default="yolov5s", help="model name")
+    opt = parser.parse_args()
+    print_args(vars(opt))
+
+    # Model
+    model = _create(name=opt.model, pretrained=True, channels=3, classes=80, autoshape=True, verbose=True)
+    # model = custom(path='path/to/model.pt')  # custom
+
+    # Images
+    imgs = [
+        "data/images/zidane.jpg",  # filename
+        Path("data/images/zidane.jpg"),  # Path
+        "https://ultralytics.com/images/zidane.jpg",  # URI
+        cv2.imread("data/images/bus.jpg")[:, :, ::-1],  # OpenCV
+        Image.open("data/images/bus.jpg"),  # PIL
+        np.zeros((320, 640, 3)),
+    ]  # numpy
+
+    # Inference
+    results = model(imgs, size=320)  # batched inference
+
+    # Results
+    results.print()
+    results.save()

iteach_toolkit/DHYOLO/model.py

@@ -0,0 +1,174 @@
+from .detect import run as run_detection
+import torch
+import cv2
+import logging
+
+# Configure logging
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
+
+class DHYOLODetector:
+    """
+    A class to encapsulate the YOLO model prediction for object detection.
+
+    Attributes:
+        model_path (str): Path to the trained YOLO model weights.
+
+    Methods:
+        predict(image_path, conf_thres, iou_thres, max_det):
+            Runs object detection on the given image using the YOLO model.
+            Logs errors in case of failure and handles exceptions.
+        plot_bboxes():
+            Plots bounding boxes on the input image based on YOLO predictions and returns the modified image.
+    """
+    
+    def __init__(self, model_path):
+        """
+        Initializes DHYOLODetector with the path to the YOLO model weights.
+        
+        Args:
+            model_path (str): Path to the YOLO model weights.
+        """
+        self.model_path = model_path
+        self.image = None
+        self.preds = None
+
+    def predict(self, image_path, conf_thres=0.25, iou_thres=0.45, max_det=1000):
+        """
+        Runs object detection on the provided image using the YOLO model.
+        
+        Args:
+            image_path (str): Path to the image file (file path or a web URL).
+            conf_thres (float): Confidence threshold for YOLO detections (default is 0.25).
+            iou_thres (float): IOU threshold for non-max suppression (default is 0.45).
+            max_det (int): Maximum number of detections per image (default is 1000).
+        
+        Returns:
+            tuple: A tuple containing:
+                - numpy.ndarray: The image read as a NumPy array.
+                - dict: A dictionary with the bounding boxes in xyxy format, confidence scores, and class labels:
+                    - 'boxes': List of bounding boxes in xyxy format.
+                    - 'confidences': List of confidence scores for each detection.
+                    - 'class_labels': List of class labels corresponding to each detection.
+        
+        Raises:
+            FileNotFoundError: If the image file does not exist or cannot be opened.
+            Exception: For any general errors during prediction.
+        """
+        try:
+            logger.info(f"Starting prediction for image: {image_path} with model: {self.model_path}")
+
+            
+            self.image = cv2.imread(image_path) # Read the image as a NumPy array
+            self.image = cv2.cvtColor(self.image, cv2.COLOR_BGR2RGB) # Convert BGR to RGB
+
+            if self.image is None:
+                raise FileNotFoundError(f"Image file not found or cannot be opened: {image_path}")
+
+            # Run YOLO detection with custom parameters
+            result = run_detection(
+                weights=self.model_path,
+                source=image_path,
+                conf_thres=conf_thres,  # Confidence threshold
+                iou_thres=iou_thres,    # IOU threshold
+                max_det=max_det          # Maximum detections
+            )
+
+            self.preds = result[0]  # We are only running for one image
+
+            # Prepare predictions as a dictionary
+            detections_dict = {
+                'boxes': [],          # List of bounding boxes in xyxy format
+                'confidences': [],    # List of confidence scores
+                'class_labels': []     # List of class labels
+            }
+
+            for detection in self.preds:
+                if isinstance(detection, torch.Tensor):
+                    detection = detection.cpu().numpy()
+
+                # Extract details
+                x1, y1, x2, y2, conf, cls = detection[:6]
+                detections_dict['boxes'].append([float(x1), float(y1), float(x2), float(y2)])  # xyxy format
+                detections_dict['confidences'].append(float(conf))
+                detections_dict['class_labels'].append(int(cls))
+
+            logger.info(f"Prediction completed successfully for image: {image_path}")
+            return self.image, detections_dict  # Return image and detections as a tuple
+        
+        except FileNotFoundError as e:
+            logger.error(f"Image file not found: {image_path}. Exception: {e}")
+            raise
+        
+        except Exception as e:
+            logger.error(f"An error occurred during prediction. Exception: {e}")
+            raise
+
+    def plot_bboxes(self, attach_watermark=False):
+        """
+        Plots bounding boxes on the input image based on YOLO predictions and returns the modified image.
+
+        Args:
+            attach_watermark (bool): Whether to attach a watermark text to the image (default is False).
+
+        Returns:
+            tuple: The original image and the image with bounding boxes plotted.
+        """
+        class_labels = {0: "door", 1: "handle"}
+        class_colors = {
+            0: (255, 0, 0),  # Red in RGB format for doors
+            1: (255, 255, 0)  # Yellow in RGB format for handles
+        }
+
+        bbox_img = self.image.copy()  # Create a copy of the original image
+
+        # Check if there are predictions
+        if self.preds is None or len(self.preds) == 0:
+            logger.warning("No predictions to display.")
+            return bbox_img, bbox_img  # Return the original image if no predictions
+
+        # Iterate through detections and plot each bounding box
+        for detection in self.preds:
+            if isinstance(detection, torch.Tensor):
+                detection = detection.cpu().numpy()
+
+            conf = detection[4]
+            x1, y1, x2, y2, _, cls = detection[:6].astype(float)  # Ensure float for bounding box coordinates
+            label = class_labels[int(cls)]
+
+            # Draw the rectangle on the bbox_img
+            cv2.rectangle(bbox_img, (int(x1), int(y1)), (int(x2), int(y2)), class_colors[int(cls)], 2)
+
+            # Prepare text with confidence score
+            text = f'{label} ({conf:.2f})'  # Include confidence score in the text
+            text_size = cv2.getTextSize(text, cv2.FONT_HERSHEY_DUPLEX, 0.5, 1)[0]
+
+            # Set text position directly above the bounding box
+            text_x = int(x1)
+            text_y = int(y1) - 2  # Adjust for a slight overlap with the bounding box
+
+            # Set text color based on class
+            text_color = (0, 0, 0) if cls == 1 else (255, 255, 255)  # Black for handle, white for door
+
+            # Draw a background rectangle for the text
+            cv2.rectangle(bbox_img, (text_x, text_y - text_size[1] - 2), (text_x + text_size[0], text_y), class_colors[int(cls)], cv2.FILLED)
+
+            # Put the label text on the bbox_img
+            cv2.putText(bbox_img, text, (text_x, text_y - 2), cv2.FONT_HERSHEY_DUPLEX, 0.5, text_color, 1, cv2.LINE_AA)
+
+        # Attach watermark if specified
+        if attach_watermark:
+            watermark_text = "Predictions by DH-YOLO"
+            watermark_color = (200, 200, 200)  # Greyish color for watermark
+            watermark_scale = 0.4  # Reduced scale for the watermark text
+            watermark_thickness = 1  # Decreased thickness for the watermark text
+            
+            # Get the text size for positioning
+            text_size = cv2.getTextSize(watermark_text, cv2.FONT_HERSHEY_DUPLEX, watermark_scale, watermark_thickness)[0]
+            text_x = bbox_img.shape[1] - text_size[0] - 10  # 10 pixels from right
+            text_y = bbox_img.shape[0] - 10  # 10 pixels from bottom
+            
+            # Put the watermark text on the bbox_img
+            cv2.putText(bbox_img, watermark_text, (text_x, text_y), cv2.FONT_HERSHEY_DUPLEX, watermark_scale, watermark_color, watermark_thickness, cv2.LINE_AA)
+
+        return self.image, bbox_img  # Return the original image and modified image with bounding boxes