Fix TFLite Segment inference (#13488)

* Fix TFLite Segment inference * Auto-format by https://ultralytics.com --------- Co-authored-by: UltralyticsAssistant <[email protected]>
ultralytics · Jan 12, 2025 · 6420a1d · 6420a1d
1 parent 86fd1ab
commit 6420a1d
Show file tree

Hide file tree

Showing 5 changed files with 17 additions and 12 deletions.
diff --git a/models/common.py b/models/common.py
@@ -750,6 +750,9 @@ def forward(self, im, augment=False, visualize=False):
                         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
@@ -355,9 +355,9 @@ def __init__(self, path, img_size=640, stride=32, auto=True, transforms=None, vi
             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
@@ -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
@@ -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
@@ -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