ultralytics · SangbumChoi · Jul 31, 2024 · Jul 31, 2024 · Aug 11, 2024 · Aug 14, 2024
diff --git a/export.py b/export.py
@@ -449,6 +449,7 @@ 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:
@@ -749,7 +750,7 @@ def export_saved_model(
         import tensorflow as tf
     from tensorflow.python.framework.convert_to_constants import convert_variables_to_constants_v2
 
-    from models.tf import TFModel
+    from models.tf import KerasModel, TFModel
 
     LOGGER.info(f"\n{prefix} starting export with tensorflow {tf.__version__}...")
     if tf.__version__ > "2.13.1":
@@ -764,12 +765,12 @@ def export_saved_model(
     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)
+    outputs = KerasModel(tf_model, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres)(inputs)
     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")
+        keras_model.save(str(file).replace(".pt", ".h5"))
     else:
         spec = tf.TensorSpec(keras_model.inputs[0].shape, keras_model.inputs[0].dtype)
         m = tf.function(lambda x: keras_model(x))  # full model

diff --git a/models/tf.py b/models/tf.py
@@ -652,6 +652,27 @@ def _xywh2xyxy(xywh):
         return tf.concat([x - w / 2, y - h / 2, x + w / 2, y + h / 2], axis=-1)
 
 
+class KerasModel(tf.keras.layers.Layer):
+    def __init__(self, tf_model, tf_nms, agnostic_nms, topk_per_class, topk_all, iou_thres, conf_thres):
+        super(KerasModel, self).__init__()
+        self.tf_model = tf_model
+        self.tf_nms = tf_nms
+        self.agnostic_nms = agnostic_nms
+        self.topk_per_class = topk_per_class
+        self.topk_all = topk_all
+        self.iou_thres = iou_thres
+        self.conf_thres = conf_thres
+
+    def get_config(self):
+        config = super().get_config()
+        return config
+
+    def call(self, inputs):
+        return self.tf_model.predict(
+            inputs, self.tf_nms, self.agnostic_nms, self.topk_per_class, self.topk_all, self.iou_thres, self.conf_thres
+        )
+
+
 class AgnosticNMS(keras.layers.Layer):
     # TF Agnostic NMS
     def call(self, input, topk_all, iou_thres, conf_thres):

diff --git a/utils/augmentations.py b/utils/augmentations.py
@@ -156,7 +156,6 @@ def random_perspective(
 ):
     # torchvision.transforms.RandomAffine(degrees=(-10, 10), translate=(0.1, 0.1), scale=(0.9, 1.1), shear=(-10, 10))
     # targets = [cls, xyxy]
-
     """Applies random perspective transformation to an image, modifying the image and corresponding labels."""
     height = im.shape[0] + border[0] * 2  # shape(h,w,c)
     width = im.shape[1] + border[1] * 2

diff --git a/utils/callbacks.py b/utils/callbacks.py
@@ -64,7 +64,6 @@ def run(self, hook, *args, thread=False, **kwargs):
             thread: (boolean) Run callbacks in daemon thread
             kwargs: Keyword Arguments to receive from YOLOv5
         """
-
         assert hook in self._callbacks, f"hook '{hook}' not found in callbacks {self._callbacks}"
         for logger in self._callbacks[hook]:
             if thread:

diff --git a/utils/dataloaders.py b/utils/dataloaders.py
@@ -1104,7 +1104,8 @@ def extract_boxes(path=DATASETS_DIR / "coco128"):
 def autosplit(path=DATASETS_DIR / "coco128/images", weights=(0.9, 0.1, 0.0), annotated_only=False):
     """Autosplit a dataset into train/val/test splits and save path/autosplit_*.txt files
     Usage: from utils.dataloaders import *; autosplit()
-    Arguments
+
+    Arguments:
         path:            Path to images directory
         weights:         Train, val, test weights (list, tuple)
         annotated_only:  Only use images with an annotated txt file
@@ -1183,7 +1184,7 @@ class HUBDatasetStats:
     """
     Class for generating HUB dataset JSON and `-hub` dataset directory.
 
-    Arguments
+    Arguments:
         path:           Path to data.yaml or data.zip (with data.yaml inside data.zip)
         autodownload:   Attempt to download dataset if not found locally
 
@@ -1314,7 +1315,7 @@ class ClassificationDataset(torchvision.datasets.ImageFolder):
     """
     YOLOv5 Classification Dataset.
 
-    Arguments
+    Arguments:
         root:  Dataset path
         transform:  torchvision transforms, used by default
         album_transform: Albumentations transforms, used if installed

diff --git a/utils/general.py b/utils/general.py
@@ -518,7 +518,6 @@ def check_font(font=FONT, progress=False):
 
 def check_dataset(data, autodownload=True):
     """Validates and/or auto-downloads a dataset, returning its configuration as a dictionary."""
-
     # Download (optional)
     extract_dir = ""
     if isinstance(data, (str, Path)) and (is_zipfile(data) or is_tarfile(data)):
@@ -1023,7 +1022,6 @@ def non_max_suppression(
     Returns:
          list of detections, on (n,6) tensor per image [xyxy, conf, cls]
     """
-
     # Checks
     assert 0 <= conf_thres <= 1, f"Invalid Confidence threshold {conf_thres}, valid values are between 0.0 and 1.0"
     assert 0 <= iou_thres <= 1, f"Invalid IoU {iou_thres}, valid values are between 0.0 and 1.0"

diff --git a/utils/loggers/__init__.py b/utils/loggers/__init__.py
@@ -350,7 +350,8 @@ class GenericLogger:
     """
     YOLOv5 General purpose logger for non-task specific logging
     Usage: from utils.loggers import GenericLogger; logger = GenericLogger(...)
-    Arguments
+
+    Arguments:
         opt:             Run arguments
         console_logger:  Console logger
         include:         loggers to include

diff --git a/utils/loggers/clearml/clearml_utils.py b/utils/loggers/clearml/clearml_utils.py
@@ -80,7 +80,7 @@ def __init__(self, opt, hyp):
         - Initialize ClearML Task, this object will capture the experiment
         - Upload dataset version to ClearML Data if opt.upload_dataset is True
 
-        arguments:
+        Arguments:
         opt (namespace) -- Commandline arguments for this run
         hyp (dict) -- Hyperparameters for this run
 
@@ -133,7 +133,7 @@ def log_scalars(self, metrics, epoch):
         """
         Log scalars/metrics to ClearML.
 
-        arguments:
+        Arguments:
         metrics (dict) Metrics in dict format: {"metrics/mAP": 0.8, ...}
         epoch (int) iteration number for the current set of metrics
         """
@@ -145,7 +145,7 @@ def log_model(self, model_path, model_name, epoch=0):
         """
         Log model weights to ClearML.
 
-        arguments:
+        Arguments:
         model_path (PosixPath or str) Path to the model weights
         model_name (str) Name of the model visible in ClearML
         epoch (int) Iteration / epoch of the model weights
@@ -158,7 +158,7 @@ def log_summary(self, metrics):
         """
         Log final metrics to a summary table.
 
-        arguments:
+        Arguments:
         metrics (dict) Metrics in dict format: {"metrics/mAP": 0.8, ...}
         """
         for k, v in metrics.items():
@@ -168,7 +168,7 @@ def log_plot(self, title, plot_path):
         """
         Log image as plot in the plot section of ClearML.
 
-        arguments:
+        Arguments:
         title (str) Title of the plot
         plot_path (PosixPath or str) Path to the saved image file
         """
@@ -183,7 +183,7 @@ def log_debug_samples(self, files, title="Debug Samples"):
         """
         Log files (images) as debug samples in the ClearML task.
 
-        arguments:
+        Arguments:
         files (List(PosixPath)) a list of file paths in PosixPath format
         title (str) A title that groups together images with the same values
         """
@@ -199,7 +199,7 @@ def log_image_with_boxes(self, image_path, boxes, class_names, image, conf_thres
         """
         Draw the bounding boxes on a single image and report the result as a ClearML debug sample.
 
-        arguments:
+        Arguments:
         image_path (PosixPath) the path the original image file
         boxes (list): list of scaled predictions in the format - [xmin, ymin, xmax, ymax, confidence, class]
         class_names (dict): dict containing mapping of class int to class name

diff --git a/utils/loggers/wandb/wandb_utils.py b/utils/loggers/wandb/wandb_utils.py
@@ -49,7 +49,7 @@ def __init__(self, opt, run_id=None, job_type="Training"):
         - Upload dataset if opt.upload_dataset is True
         - Setup training processes if job_type is 'Training'
 
-        arguments:
+        Arguments:
         opt (namespace) -- Commandline arguments for this run
         run_id (str) -- Run ID of W&B run to be resumed
         job_type (str) -- To set the job_type for this run
@@ -90,7 +90,7 @@ def setup_training(self, opt):
           - Update data_dict, to contain info of previous run if resumed and the paths of dataset artifact if downloaded
           - Setup log_dict, initialize bbox_interval
 
-        arguments:
+        Arguments:
         opt (namespace) -- commandline arguments for this run
 
         """
@@ -120,7 +120,7 @@ def log_model(self, path, opt, epoch, fitness_score, best_model=False):
         """
         Log the model checkpoint as W&B artifact.
 
-        arguments:
+        Arguments:
         path (Path)   -- Path of directory containing the checkpoints
         opt (namespace) -- Command line arguments for this run
         epoch (int)  -- Current epoch number
@@ -159,7 +159,7 @@ def log(self, log_dict):
         """
         Save the metrics to the logging dictionary.
 
-        arguments:
+        Arguments:
         log_dict (Dict) -- metrics/media to be logged in current step
         """
         if self.wandb_run:
@@ -170,7 +170,7 @@ def end_epoch(self):
         """
         Commit the log_dict, model artifacts and Tables to W&B and flush the log_dict.
 
-        arguments:
+        Arguments:
         best_result (boolean): Boolean representing if the result of this evaluation is best or not
         """
         if self.wandb_run:
@@ -197,7 +197,7 @@ def finish_run(self):
 
 @contextmanager
 def all_logging_disabled(highest_level=logging.CRITICAL):
-    """source - https://gist.github.com/simon-weber/7853144
+    """Source - https://gist.github.com/simon-weber/7853144
     A context manager that will prevent any logging messages triggered during the body from being processed.
     :param highest_level: the maximum logging level in use.
       This would only need to be changed if a custom level greater than CRITICAL is defined.

diff --git a/utils/metrics.py b/utils/metrics.py
@@ -41,7 +41,6 @@ def ap_per_class(tp, conf, pred_cls, target_cls, plot=False, save_dir=".", names
     # Returns
         The average precision as computed in py-faster-rcnn.
     """
-
     # Sort by objectness
     i = np.argsort(-conf)
     tp, conf, pred_cls = tp[i], conf[i], pred_cls[i]
@@ -103,7 +102,6 @@ def compute_ap(recall, precision):
     # Returns
         Average precision, precision curve, recall curve
     """
-
     # Append sentinel values to beginning and end
     mrec = np.concatenate(([0.0], recall, [1.0]))
     mpre = np.concatenate(([1.0], precision, [0.0]))
@@ -137,6 +135,7 @@ def process_batch(self, detections, labels):
         Return intersection-over-union (Jaccard index) of boxes.
 
         Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
+
         Arguments:
             detections (Array[N, 6]), x1, y1, x2, y2, conf, class
             labels (Array[M, 5]), class, x1, y1, x2, y2
@@ -233,7 +232,6 @@ def bbox_iou(box1, box2, xywh=True, GIoU=False, DIoU=False, CIoU=False, eps=1e-7
 
     Input shapes are box1(1,4) to box2(n,4).
     """
-
     # Get the coordinates of bounding boxes
     if xywh:  # transform from xywh to xyxy
         (x1, y1, w1, h1), (x2, y2, w2, h2) = box1.chunk(4, -1), box2.chunk(4, -1)
@@ -279,14 +277,15 @@ def box_iou(box1, box2, eps=1e-7):
     Return intersection-over-union (Jaccard index) of boxes.
 
     Both sets of boxes are expected to be in (x1, y1, x2, y2) format.
+
     Arguments:
         box1 (Tensor[N, 4])
         box2 (Tensor[M, 4])
+
     Returns:
         iou (Tensor[N, M]): the NxM matrix containing the pairwise
             IoU values for every element in boxes1 and boxes2
     """
-
     # inter(N,M) = (rb(N,M,2) - lt(N,M,2)).clamp(0).prod(2)
     (a1, a2), (b1, b2) = box1.unsqueeze(1).chunk(2, 2), box2.unsqueeze(0).chunk(2, 2)
     inter = (torch.min(a2, b2) - torch.max(a1, b1)).clamp(0).prod(2)
@@ -304,7 +303,6 @@ def bbox_ioa(box1, box2, eps=1e-7):
     box2:       np.array of shape(nx4)
     returns:    np.array of shape(n)
     """
-
     # Get the coordinates of bounding boxes
     b1_x1, b1_y1, b1_x2, b1_y2 = box1
     b2_x1, b2_y1, b2_x2, b2_y2 = box2.T

diff --git a/utils/segment/augmentations.py b/utils/segment/augmentations.py
@@ -29,7 +29,6 @@ def random_perspective(
 ):
     # torchvision.transforms.RandomAffine(degrees=(-10, 10), translate=(.1, .1), scale=(.9, 1.1), shear=(-10, 10))
     # targets = [cls, xyxy]
-
     """Applies random perspective, rotation, scale, shear, and translation augmentations to an image and targets."""
     height = im.shape[0] + border[0] * 2  # shape(h,w,c)
     width = im.shape[1] + border[1] * 2

diff --git a/utils/segment/general.py b/utils/segment/general.py
@@ -14,7 +14,6 @@ def crop_mask(masks, boxes):
         - masks should be a size [n, h, w] tensor of masks
         - boxes should be a size [n, 4] tensor of bbox coords in relative point form
     """
-
     n, h, w = masks.shape
     x1, y1, x2, y2 = torch.chunk(boxes[:, :, None], 4, 1)  # x1 shape(1,1,n)
     r = torch.arange(w, device=masks.device, dtype=x1.dtype)[None, None, :]  # rows shape(1,w,1)
@@ -33,7 +32,6 @@ def process_mask_upsample(protos, masks_in, bboxes, shape):
 
     return: h, w, n
     """
-
     c, mh, mw = protos.shape  # CHW
     masks = (masks_in @ protos.float().view(c, -1)).sigmoid().view(-1, mh, mw)
     masks = F.interpolate(masks[None], shape, mode="bilinear", align_corners=False)[0]  # CHW
@@ -51,7 +49,6 @@ def process_mask(protos, masks_in, bboxes, shape, upsample=False):
 
     return: h, w, n
     """
-
     c, mh, mw = protos.shape  # CHW
     ih, iw = shape
     masks = (masks_in @ protos.float().view(c, -1)).sigmoid().view(-1, mh, mw)  # CHW

diff --git a/utils/triton.py b/utils/triton.py
@@ -17,10 +17,9 @@ class TritonRemoteModel:
 
     def __init__(self, url: str):
         """
-        Keyword arguments:
+        Keyword Arguments:
         url: Fully qualified address of the Triton server - for e.g. grpc://localhost:8000
         """
-
         parsed_url = urlparse(url)
         if parsed_url.scheme == "grpc":
             from tritonclient.grpc import InferenceServerClient, InferInput