fix iteach

iteach_toolkit/DHYOLO/__init__.py

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

iteach_toolkit/DHYOLO/detect.py

@@ -1,290 +0,0 @@
-# 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

@@ -1,1537 +0,0 @@
-# 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 = "[email protected]"
-    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

@@ -1,510 +0,0 @@
-# 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

@@ -1,174 +0,0 @@
-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