From 6420a1db87460d36fd2141a65659093df27c1996 Mon Sep 17 00:00:00 2001
From: Mohammed Yasin <32206511+Y-T-G@users.noreply.github.com>
Date: Mon, 13 Jan 2025 02:05:20 +0800
Subject: [PATCH] Fix TFLite Segment inference (#13488)

* Fix TFLite Segment inference

* Auto-format by https://ultralytics.com

---------

Co-authored-by: UltralyticsAssistant <web@ultralytics.com>
---
 models/common.py                       | 3 +++
 utils/dataloaders.py                   | 6 +++---
 utils/general.py                       | 6 +++---
 utils/loggers/clearml/clearml_utils.py | 8 +++++---
 utils/torch_utils.py                   | 6 +++---
 5 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/models/common.py b/models/common.py
index 8ad53d5..1ae3280 100644
--- a/models/common.py
+++ b/models/common.py
@@ -750,6 +750,9 @@ class DetectMultiBackend(nn.Module):
                         scale, zero_point = output["quantization"]
                         x = (x.astype(np.float32) - zero_point) * scale  # re-scale
                     y.append(x)
+            if len(y) == 2:  # segment with (det, proto) output order reversed
+                if len(y[1].shape) != 4:
+                    y = list(reversed(y))  # should be y = (1, 116, 8400), (1, 160, 160, 32)
             y = [x if isinstance(x, np.ndarray) else x.numpy() for x in y]
             y[0][..., :4] *= [w, h, w, h]  # xywh normalized to pixels
 
diff --git a/utils/dataloaders.py b/utils/dataloaders.py
index fc5da6b..0824048 100644
--- a/utils/dataloaders.py
+++ b/utils/dataloaders.py
@@ -355,9 +355,9 @@ class LoadImages:
             self._new_video(videos[0])  # new video
         else:
             self.cap = None
-        assert self.nf > 0, (
-            f"No images or videos found in {p}. Supported formats are:\nimages: {IMG_FORMATS}\nvideos: {VID_FORMATS}"
-        )
+        assert (
+            self.nf > 0
+        ), f"No images or videos found in {p}. Supported formats are:\nimages: {IMG_FORMATS}\nvideos: {VID_FORMATS}"
 
     def __iter__(self):
         """Initializes iterator by resetting count and returns the iterator object itself."""
diff --git a/utils/general.py b/utils/general.py
index 41cd203..db6ddac 100644
--- a/utils/general.py
+++ b/utils/general.py
@@ -495,9 +495,9 @@ def check_file(file, suffix=""):
             assert Path(file).exists() and Path(file).stat().st_size > 0, f"File download failed: {url}"  # check
         return file
     elif file.startswith("clearml://"):  # ClearML Dataset ID
-        assert "clearml" in sys.modules, (
-            "ClearML is not installed, so cannot use ClearML dataset. Try running 'pip install clearml'."
-        )
+        assert (
+            "clearml" in sys.modules
+        ), "ClearML is not installed, so cannot use ClearML dataset. Try running 'pip install clearml'."
         return file
     else:  # search
         files = []
diff --git a/utils/loggers/clearml/clearml_utils.py b/utils/loggers/clearml/clearml_utils.py
index 2ddf3a4..fc19c8c 100644
--- a/utils/loggers/clearml/clearml_utils.py
+++ b/utils/loggers/clearml/clearml_utils.py
@@ -41,9 +41,11 @@ def construct_dataset(clearml_info_string):
     with open(yaml_filenames[0]) as f:
         dataset_definition = yaml.safe_load(f)
 
-    assert set(dataset_definition.keys()).issuperset({"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')"
-    )
+    assert set(
+        dataset_definition.keys()
+    ).issuperset(
+        {"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 = {
         "train": (
diff --git a/utils/torch_utils.py b/utils/torch_utils.py
index 29877fa..53e7076 100644
--- a/utils/torch_utils.py
+++ b/utils/torch_utils.py
@@ -121,9 +121,9 @@ def select_device(device="", batch_size=0, newline=True):
         os.environ["CUDA_VISIBLE_DEVICES"] = "-1"  # force torch.cuda.is_available() = False
     elif device:  # non-cpu device requested
         os.environ["CUDA_VISIBLE_DEVICES"] = device  # set environment variable - must be before assert is_available()
-        assert torch.cuda.is_available() and torch.cuda.device_count() >= len(device.replace(",", "")), (
-            f"Invalid CUDA '--device {device}' requested, use '--device cpu' or pass valid CUDA device(s)"
-        )
+        assert torch.cuda.is_available() and torch.cuda.device_count() >= len(
+            device.replace(",", "")
+        ), f"Invalid CUDA '--device {device}' requested, use '--device cpu' or pass valid CUDA device(s)"
 
     if not cpu and not mps and torch.cuda.is_available():  # prefer GPU if available
         devices = device.split(",") if device else "0"  # range(torch.cuda.device_count())  # i.e. 0,1,6,7