From 22e9afef7401ab6650e45e4937b36ebd210484af Mon Sep 17 00:00:00 2001
From: Glenn Jocher <glenn.jocher@ultralytics.com>
Date: Sat, 13 Jan 2024 22:34:05 +0100
Subject: [PATCH] Python refactor and simplification (#12624)

* Python code cleanup

* Auto-format by Ultralytics actions

---------

Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
---
 export.py                              |   2 +-
 models/common.py                       |  16 ++--
 train.py                               |  32 +++----
 utils/downloads.py                     |   4 +-
 utils/general.py                       |   2 +-
 utils/loggers/__init__.py              |  25 +++---
 utils/loggers/clearml/clearml_utils.py |  50 ++++++-----
 utils/loggers/comet/__init__.py        | 120 ++++++++++++-------------
 utils/loggers/comet/comet_utils.py     |  34 ++++---
 utils/loggers/wandb/wandb_utils.py     |  45 +++++-----
 utils/segment/metrics.py               |   3 +-
 11 files changed, 158 insertions(+), 175 deletions(-)

diff --git a/export.py b/export.py
index 4f269fef005c..99ea15e7030c 100644
--- a/export.py
+++ b/export.py
@@ -546,7 +546,7 @@ def export_tfjs(file, int8, prefix=colorstr("TensorFlow.js:")):
         "--quantize_uint8" if int8 else "",
         "--output_node_names=Identity,Identity_1,Identity_2,Identity_3",
         str(f_pb),
-        str(f),
+        f,
     ]
     subprocess.run([arg for arg in args if arg], check=True)
 
diff --git a/models/common.py b/models/common.py
index 09e7560f4d84..b21b42b00d0d 100644
--- a/models/common.py
+++ b/models/common.py
@@ -859,11 +859,17 @@ def pandas(self):
     def tolist(self):
         # return a list of Detections objects, i.e. 'for result in results.tolist():'
         r = range(self.n)  # iterable
-        x = [Detections([self.ims[i]], [self.pred[i]], [self.files[i]], self.times, self.names, self.s) for i in r]
-        # for d in x:
-        #    for k in ['ims', 'pred', 'xyxy', 'xyxyn', 'xywh', 'xywhn']:
-        #        setattr(d, k, getattr(d, k)[0])  # pop out of list
-        return x
+        return [
+            Detections(
+                [self.ims[i]],
+                [self.pred[i]],
+                [self.files[i]],
+                self.times,
+                self.names,
+                self.s,
+            )
+            for i in r
+        ]
 
     def print(self):
         LOGGER.info(self.__str__())
diff --git a/train.py b/train.py
index 73297d204393..ca284d06df25 100644
--- a/train.py
+++ b/train.py
@@ -682,10 +682,7 @@ def main(opt, callbacks=Callbacks()):
             )
 
         # Delete the items in meta dictionary whose first value is False
-        del_ = []
-        for item in meta.keys():
-            if meta[item][0] is False:
-                del_.append(item)
+        del_ = [item for item, value_ in meta.items() if value_[0] is False]
         hyp_GA = hyp.copy()  # Make a copy of hyp dictionary
         for item in del_:
             del meta[item]  # Remove the item from meta dictionary
@@ -696,9 +693,7 @@ def main(opt, callbacks=Callbacks()):
         upper_limit = np.array([meta[k][2] for k in hyp_GA.keys()])
 
         # Create gene_ranges list to hold the range of values for each gene in the population
-        gene_ranges = []
-        for i in range(len(upper_limit)):
-            gene_ranges.append((lower_limit[i], upper_limit[i]))
+        gene_ranges = [(lower_limit[i], upper_limit[i]) for i in range(len(upper_limit))]
 
         # Initialize the population with initial_values or random values
         initial_values = []
@@ -723,14 +718,11 @@ def main(opt, callbacks=Callbacks()):
 
         # Generate random values within the search space for the rest of the population
         if initial_values is None:
-            population = [generate_individual(gene_ranges, len(hyp_GA)) for i in range(pop_size)]
-        else:
-            if pop_size > 1:
-                population = [
-                    generate_individual(gene_ranges, len(hyp_GA)) for i in range(pop_size - len(initial_values))
-                ]
-                for initial_value in initial_values:
-                    population = [initial_value] + population
+            population = [generate_individual(gene_ranges, len(hyp_GA)) for _ in range(pop_size)]
+        elif pop_size > 1:
+            population = [generate_individual(gene_ranges, len(hyp_GA)) for _ in range(pop_size - len(initial_values))]
+            for initial_value in initial_values:
+                population = [initial_value] + population
 
         # Run the genetic algorithm for a fixed number of generations
         list_keys = list(hyp_GA.keys())
@@ -738,10 +730,8 @@ def main(opt, callbacks=Callbacks()):
             if generation >= 1:
                 save_dict = {}
                 for i in range(len(population)):
-                    little_dict = {}
-                    for j in range(len(population[i])):
-                        little_dict[list_keys[j]] = float(population[i][j])
-                    save_dict["gen" + str(generation) + "number" + str(i)] = little_dict
+                    little_dict = {list_keys[j]: float(population[i][j]) for j in range(len(population[i]))}
+                    save_dict[f"gen{str(generation)}number{str(i)}"] = little_dict
 
                 with open(save_dir / "evolve_population.yaml", "w") as outfile:
                     yaml.dump(save_dict, outfile, default_flow_style=False)
@@ -771,7 +761,7 @@ def main(opt, callbacks=Callbacks()):
 
             # Select the fittest individuals for reproduction using adaptive tournament selection
             selected_indices = []
-            for i in range(pop_size - elite_size):
+            for _ in range(pop_size - elite_size):
                 # Adaptive tournament size
                 tournament_size = max(
                     max(2, tournament_size_min),
@@ -788,7 +778,7 @@ def main(opt, callbacks=Callbacks()):
             selected_indices.extend(elite_indices)
             # Create the next generation through crossover and mutation
             next_generation = []
-            for i in range(pop_size):
+            for _ in range(pop_size):
                 parent1_index = selected_indices[random.randint(0, pop_size - 1)]
                 parent2_index = selected_indices[random.randint(0, pop_size - 1)]
                 # Adaptive crossover rate
diff --git a/utils/downloads.py b/utils/downloads.py
index ee700acb618b..ccab278a1d00 100644
--- a/utils/downloads.py
+++ b/utils/downloads.py
@@ -24,9 +24,7 @@ def is_url(url, check=True):
 def gsutil_getsize(url=""):
     # gs://bucket/file size https://cloud.google.com/storage/docs/gsutil/commands/du
     output = subprocess.check_output(["gsutil", "du", url], shell=True, encoding="utf-8")
-    if output:
-        return int(output.split()[0])
-    return 0
+    return int(output.split()[0]) if output else 0
 
 
 def url_getsize(url="https://ultralytics.com/images/bus.jpg"):
diff --git a/utils/general.py b/utils/general.py
index 47ab656e5a3f..4ba80cec2cb6 100644
--- a/utils/general.py
+++ b/utils/general.py
@@ -188,7 +188,7 @@ class Profile(contextlib.ContextDecorator):
     def __init__(self, t=0.0, device: torch.device = None):
         self.t = t
         self.device = device
-        self.cuda = True if (device and str(device)[:4] == "cuda") else False
+        self.cuda = bool(device and str(device).startswith("cuda"))
 
     def __enter__(self):
         self.start = self.time()
diff --git a/utils/loggers/__init__.py b/utils/loggers/__init__.py
index df67e45c8221..36792979913a 100644
--- a/utils/loggers/__init__.py
+++ b/utils/loggers/__init__.py
@@ -67,9 +67,7 @@ def _json_default(value):
             value = value.item()
         except ValueError:  # "only one element tensors can be converted to Python scalars"
             pass
-    if isinstance(value, float):
-        return value
-    return str(value)
+    return value if isinstance(value, float) else str(value)
 
 
 class Loggers:
@@ -250,12 +248,12 @@ def on_fit_epoch_end(self, vals, epoch, best_fitness, fi):
                 f.write(s + ("%20.5g," * n % tuple([epoch] + vals)).rstrip(",") + "\n")
         if self.ndjson_console or self.ndjson_file:
             json_data = json.dumps(dict(epoch=epoch, **x), default=_json_default)
-            if self.ndjson_console:
-                print(json_data)
-            if self.ndjson_file:
-                file = self.save_dir / "results.ndjson"
-                with open(file, "a") as f:
-                    print(json_data, file=f)
+        if self.ndjson_console:
+            print(json_data)
+        if self.ndjson_file:
+            file = self.save_dir / "results.ndjson"
+            with open(file, "a") as f:
+                print(json_data, file=f)
 
         if self.tb:
             for k, v in x.items():
@@ -370,10 +368,7 @@ def __init__(self, opt, console_logger, include=("tb", "wandb", "clearml")):
         if clearml and "clearml" in self.include:
             try:
                 # Hyp is not available in classification mode
-                if "hyp" not in opt:
-                    hyp = {}
-                else:
-                    hyp = opt.hyp
+                hyp = {} if "hyp" not in opt else opt.hyp
                 self.clearml = ClearmlLogger(opt, hyp)
             except Exception:
                 self.clearml = None
@@ -427,7 +422,9 @@ def log_graph(self, model, imgsz=(640, 640)):
         if self.tb:
             log_tensorboard_graph(self.tb, model, imgsz)
 
-    def log_model(self, model_path, epoch=0, metadata={}):
+    def log_model(self, model_path, epoch=0, metadata=None):
+        if metadata is None:
+            metadata = {}
         # Log model to all loggers
         if self.wandb:
             art = wandb.Artifact(name=f"run_{wandb.run.id}_model", type="model", metadata=metadata)
diff --git a/utils/loggers/clearml/clearml_utils.py b/utils/loggers/clearml/clearml_utils.py
index 8b141d177afd..1bbea61effc2 100644
--- a/utils/loggers/clearml/clearml_utils.py
+++ b/utils/loggers/clearml/clearml_utils.py
@@ -31,7 +31,7 @@ def construct_dataset(clearml_info_string):
             "More than one yaml file was found in the dataset root, cannot determine which one contains "
             "the dataset definition this way."
         )
-    elif len(yaml_filenames) == 0:
+    elif not yaml_filenames:
         raise ValueError(
             "No yaml definition found in dataset root path, check that there is a correct yaml file "
             "inside the dataset root path."
@@ -45,7 +45,7 @@ def construct_dataset(clearml_info_string):
         {"train", "test", "val", "nc", "names"}
     ), "The right keys were not found in the yaml file, make sure it at least has the following keys: ('train', 'test', 'val', 'nc', 'names')"
 
-    data_dict = dict()
+    data_dict = {}
     data_dict["train"] = (
         str((dataset_root_path / dataset_definition["train"]).resolve()) if dataset_definition["train"] else None
     )
@@ -96,7 +96,7 @@ def __init__(self, opt, hyp):
         self.data_dict = None
         if self.clearml:
             self.task = Task.init(
-                project_name=opt.project if not str(opt.project).startswith("runs/") else "YOLOv5",
+                project_name="YOLOv5" if str(opt.project).startswith("runs/") else opt.project,
                 task_name=opt.name if opt.name != "exp" else "Training",
                 tags=["YOLOv5"],
                 output_uri=True,
@@ -202,24 +202,26 @@ def log_image_with_boxes(self, image_path, boxes, class_names, image, conf_thres
         class_names (dict): dict containing mapping of class int to class name
         image (Tensor): A torch tensor containing the actual image data
         """
-        if len(self.current_epoch_logged_images) < self.max_imgs_to_log_per_epoch and self.current_epoch >= 0:
-            # Log every bbox_interval times and deduplicate for any intermittend extra eval runs
-            if self.current_epoch % self.bbox_interval == 0 and image_path not in self.current_epoch_logged_images:
-                im = np.ascontiguousarray(np.moveaxis(image.mul(255).clamp(0, 255).byte().cpu().numpy(), 0, 2))
-                annotator = Annotator(im=im, pil=True)
-                for i, (conf, class_nr, box) in enumerate(zip(boxes[:, 4], boxes[:, 5], boxes[:, :4])):
-                    color = colors(i)
-
-                    class_name = class_names[int(class_nr)]
-                    confidence_percentage = round(float(conf) * 100, 2)
-                    label = f"{class_name}: {confidence_percentage}%"
-
-                    if conf > conf_threshold:
-                        annotator.rectangle(box.cpu().numpy(), outline=color)
-                        annotator.box_label(box.cpu().numpy(), label=label, color=color)
-
-                annotated_image = annotator.result()
-                self.task.get_logger().report_image(
-                    title="Bounding Boxes", series=image_path.name, iteration=self.current_epoch, image=annotated_image
-                )
-                self.current_epoch_logged_images.add(image_path)
+        if (
+            len(self.current_epoch_logged_images) < self.max_imgs_to_log_per_epoch
+            and self.current_epoch >= 0
+            and (self.current_epoch % self.bbox_interval == 0 and image_path not in self.current_epoch_logged_images)
+        ):
+            im = np.ascontiguousarray(np.moveaxis(image.mul(255).clamp(0, 255).byte().cpu().numpy(), 0, 2))
+            annotator = Annotator(im=im, pil=True)
+            for i, (conf, class_nr, box) in enumerate(zip(boxes[:, 4], boxes[:, 5], boxes[:, :4])):
+                color = colors(i)
+
+                class_name = class_names[int(class_nr)]
+                confidence_percentage = round(float(conf) * 100, 2)
+                label = f"{class_name}: {confidence_percentage}%"
+
+                if conf > conf_threshold:
+                    annotator.rectangle(box.cpu().numpy(), outline=color)
+                    annotator.box_label(box.cpu().numpy(), label=label, color=color)
+
+            annotated_image = annotator.result()
+            self.task.get_logger().report_image(
+                title="Bounding Boxes", series=image_path.name, iteration=self.current_epoch, image=annotated_image
+            )
+            self.current_epoch_logged_images.add(image_path)
diff --git a/utils/loggers/comet/__init__.py b/utils/loggers/comet/__init__.py
index bdf81f63982e..cec46f5af1fb 100644
--- a/utils/loggers/comet/__init__.py
+++ b/utils/loggers/comet/__init__.py
@@ -166,34 +166,33 @@ def __init__(self, opt, hyp, run_id=None, job_type="Training", **experiment_kwar
 
     def _get_experiment(self, mode, experiment_id=None):
         if mode == "offline":
+            return (
+                comet_ml.ExistingOfflineExperiment(
+                    previous_experiment=experiment_id,
+                    **self.default_experiment_kwargs,
+                )
+                if experiment_id is not None
+                else comet_ml.OfflineExperiment(
+                    **self.default_experiment_kwargs,
+                )
+            )
+        try:
             if experiment_id is not None:
-                return comet_ml.ExistingOfflineExperiment(
+                return comet_ml.ExistingExperiment(
                     previous_experiment=experiment_id,
                     **self.default_experiment_kwargs,
                 )
 
-            return comet_ml.OfflineExperiment(
-                **self.default_experiment_kwargs,
-            )
+            return comet_ml.Experiment(**self.default_experiment_kwargs)
 
-        else:
-            try:
-                if experiment_id is not None:
-                    return comet_ml.ExistingExperiment(
-                        previous_experiment=experiment_id,
-                        **self.default_experiment_kwargs,
-                    )
-
-                return comet_ml.Experiment(**self.default_experiment_kwargs)
-
-            except ValueError:
-                logger.warning(
-                    "COMET WARNING: "
-                    "Comet credentials have not been set. "
-                    "Comet will default to offline logging. "
-                    "Please set your credentials to enable online logging."
-                )
-                return self._get_experiment("offline", experiment_id)
+        except ValueError:
+            logger.warning(
+                "COMET WARNING: "
+                "Comet credentials have not been set. "
+                "Comet will default to offline logging. "
+                "Please set your credentials to enable online logging."
+            )
+            return self._get_experiment("offline", experiment_id)
 
         return
 
@@ -242,10 +241,7 @@ def check_dataset(self, data_file):
         path = data_config.get("path")
         if path and path.startswith(COMET_PREFIX):
             path = data_config["path"].replace(COMET_PREFIX, "")
-            data_dict = self.download_dataset_artifact(path)
-
-            return data_dict
-
+            return self.download_dataset_artifact(path)
         self.log_asset(self.opt.data, metadata={"type": "data-config-file"})
 
         return check_dataset(data_file)
@@ -269,24 +265,22 @@ def log_predictions(self, image, labelsn, path, shape, predn):
             self.log_image(native_scale_image, name=image_name)
             self.logged_image_names.append(image_name)
 
-        metadata = []
-        for cls, *xyxy in filtered_labels.tolist():
-            metadata.append(
-                {
-                    "label": f"{self.class_names[int(cls)]}-gt",
-                    "score": 100,
-                    "box": {"x": xyxy[0], "y": xyxy[1], "x2": xyxy[2], "y2": xyxy[3]},
-                }
-            )
-        for *xyxy, conf, cls in filtered_detections.tolist():
-            metadata.append(
-                {
-                    "label": f"{self.class_names[int(cls)]}",
-                    "score": conf * 100,
-                    "box": {"x": xyxy[0], "y": xyxy[1], "x2": xyxy[2], "y2": xyxy[3]},
-                }
-            )
-
+        metadata = [
+            {
+                "label": f"{self.class_names[int(cls)]}-gt",
+                "score": 100,
+                "box": {"x": xyxy[0], "y": xyxy[1], "x2": xyxy[2], "y2": xyxy[3]},
+            }
+            for cls, *xyxy in filtered_labels.tolist()
+        ]
+        metadata.extend(
+            {
+                "label": f"{self.class_names[int(cls)]}",
+                "score": conf * 100,
+                "box": {"x": xyxy[0], "y": xyxy[1], "x2": xyxy[2], "y2": xyxy[3]},
+            }
+            for *xyxy, conf, cls in filtered_detections.tolist()
+        )
         self.metadata_dict[image_name] = metadata
         self.logged_images_count += 1
 
@@ -398,9 +392,8 @@ def on_pretrain_routine_end(self, paths):
         for path in paths:
             self.log_asset(str(path))
 
-        if self.upload_dataset:
-            if not self.resume:
-                self.upload_dataset_artifact()
+        if self.upload_dataset and not self.resume:
+            self.upload_dataset_artifact()
 
         return
 
@@ -477,23 +470,22 @@ def on_val_batch_end(self, batch_i, images, targets, paths, shapes, outputs):
         return
 
     def on_val_end(self, nt, tp, fp, p, r, f1, ap, ap50, ap_class, confusion_matrix):
-        if self.comet_log_per_class_metrics:
-            if self.num_classes > 1:
-                for i, c in enumerate(ap_class):
-                    class_name = self.class_names[c]
-                    self.experiment.log_metrics(
-                        {
-                            "mAP@.5": ap50[i],
-                            "mAP@.5:.95": ap[i],
-                            "precision": p[i],
-                            "recall": r[i],
-                            "f1": f1[i],
-                            "true_positives": tp[i],
-                            "false_positives": fp[i],
-                            "support": nt[c],
-                        },
-                        prefix=class_name,
-                    )
+        if self.comet_log_per_class_metrics and self.num_classes > 1:
+            for i, c in enumerate(ap_class):
+                class_name = self.class_names[c]
+                self.experiment.log_metrics(
+                    {
+                        "mAP@.5": ap50[i],
+                        "mAP@.5:.95": ap[i],
+                        "precision": p[i],
+                        "recall": r[i],
+                        "f1": f1[i],
+                        "true_positives": tp[i],
+                        "false_positives": fp[i],
+                        "support": nt[c],
+                    },
+                    prefix=class_name,
+                )
 
         if self.comet_log_confusion_matrix:
             epoch = self.experiment.curr_epoch
diff --git a/utils/loggers/comet/comet_utils.py b/utils/loggers/comet/comet_utils.py
index f7b56dddd5f7..6e8fad68c6cc 100644
--- a/utils/loggers/comet/comet_utils.py
+++ b/utils/loggers/comet/comet_utils.py
@@ -4,7 +4,7 @@
 
 try:
     import comet_ml
-except (ModuleNotFoundError, ImportError):
+except ImportError:
     comet_ml = None
 
 import yaml
@@ -109,14 +109,13 @@ def check_comet_weights(opt):
     if comet_ml is None:
         return
 
-    if isinstance(opt.weights, str):
-        if opt.weights.startswith(COMET_PREFIX):
-            api = comet_ml.API()
-            resource = urlparse(opt.weights)
-            experiment_path = f"{resource.netloc}{resource.path}"
-            experiment = api.get(experiment_path)
-            download_model_checkpoint(opt, experiment)
-            return True
+    if isinstance(opt.weights, str) and opt.weights.startswith(COMET_PREFIX):
+        api = comet_ml.API()
+        resource = urlparse(opt.weights)
+        experiment_path = f"{resource.netloc}{resource.path}"
+        experiment = api.get(experiment_path)
+        download_model_checkpoint(opt, experiment)
+        return True
 
     return None
 
@@ -136,15 +135,14 @@ def check_comet_resume(opt):
     if comet_ml is None:
         return
 
-    if isinstance(opt.resume, str):
-        if opt.resume.startswith(COMET_PREFIX):
-            api = comet_ml.API()
-            resource = urlparse(opt.resume)
-            experiment_path = f"{resource.netloc}{resource.path}"
-            experiment = api.get(experiment_path)
-            set_opt_parameters(opt, experiment)
-            download_model_checkpoint(opt, experiment)
+    if isinstance(opt.resume, str) and opt.resume.startswith(COMET_PREFIX):
+        api = comet_ml.API()
+        resource = urlparse(opt.resume)
+        experiment_path = f"{resource.netloc}{resource.path}"
+        experiment = api.get(experiment_path)
+        set_opt_parameters(opt, experiment)
+        download_model_checkpoint(opt, experiment)
 
-            return True
+        return True
 
     return None
diff --git a/utils/loggers/wandb/wandb_utils.py b/utils/loggers/wandb/wandb_utils.py
index f8d49a33d00f..0af8bda12d85 100644
--- a/utils/loggers/wandb/wandb_utils.py
+++ b/utils/loggers/wandb/wandb_utils.py
@@ -65,28 +65,23 @@ def __init__(self, opt, run_id=None, job_type="Training"):
         self.max_imgs_to_log = 16
         self.data_dict = None
         if self.wandb:
-            self.wandb_run = (
-                wandb.init(
-                    config=opt,
-                    resume="allow",
-                    project="YOLOv5" if opt.project == "runs/train" else Path(opt.project).stem,
-                    entity=opt.entity,
-                    name=opt.name if opt.name != "exp" else None,
-                    job_type=job_type,
-                    id=run_id,
-                    allow_val_change=True,
-                )
-                if not wandb.run
-                else wandb.run
+            self.wandb_run = wandb.run or wandb.init(
+                config=opt,
+                resume="allow",
+                project="YOLOv5" if opt.project == "runs/train" else Path(opt.project).stem,
+                entity=opt.entity,
+                name=opt.name if opt.name != "exp" else None,
+                job_type=job_type,
+                id=run_id,
+                allow_val_change=True,
             )
 
-        if self.wandb_run:
-            if self.job_type == "Training":
-                if isinstance(opt.data, dict):
-                    # This means another dataset manager has already processed the dataset info (e.g. ClearML)
-                    # and they will have stored the already processed dict in opt.data
-                    self.data_dict = opt.data
-                self.setup_training(opt)
+        if self.wandb_run and self.job_type == "Training":
+            if isinstance(opt.data, dict):
+                # This means another dataset manager has already processed the dataset info (e.g. ClearML)
+                # and they will have stored the already processed dict in opt.data
+                self.data_dict = opt.data
+            self.setup_training(opt)
 
     def setup_training(self, opt):
         """
@@ -133,7 +128,7 @@ def log_model(self, path, opt, epoch, fitness_score, best_model=False):
         best_model (boolean) -- Boolean representing if the current checkpoint is the best yet.
         """
         model_artifact = wandb.Artifact(
-            "run_" + wandb.run.id + "_model",
+            f"run_{wandb.run.id}_model",
             type="model",
             metadata={
                 "original_url": str(path),
@@ -146,7 +141,13 @@ def log_model(self, path, opt, epoch, fitness_score, best_model=False):
         )
         model_artifact.add_file(str(path / "last.pt"), name="last.pt")
         wandb.log_artifact(
-            model_artifact, aliases=["latest", "last", "epoch " + str(self.current_epoch), "best" if best_model else ""]
+            model_artifact,
+            aliases=[
+                "latest",
+                "last",
+                f"epoch {str(self.current_epoch)}",
+                "best" if best_model else "",
+            ],
         )
         LOGGER.info(f"Saving model artifact on epoch {epoch + 1}")
 
diff --git a/utils/segment/metrics.py b/utils/segment/metrics.py
index 222a749b5986..7811e7eb364a 100644
--- a/utils/segment/metrics.py
+++ b/utils/segment/metrics.py
@@ -35,7 +35,7 @@ def ap_per_class_box_and_mask(
         tp_m, conf, pred_cls, target_cls, plot=plot, save_dir=save_dir, names=names, prefix="Mask"
     )[2:]
 
-    results = {
+    return {
         "boxes": {
             "p": results_boxes[0],
             "r": results_boxes[1],
@@ -51,7 +51,6 @@ def ap_per_class_box_and_mask(
             "ap_class": results_masks[4],
         },
     }
-    return results
 
 
 class Metric: