From e999853c94261c4dab96e98739dc2e9335209eb6 Mon Sep 17 00:00:00 2001
From: TY1667 <mengyunhao@mail.ynu.edu.cn>
Date: Tue, 4 Nov 2025 13:18:43 +0800
Subject: [PATCH] fix pylance type warning

---
 utils/dataset.py | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/utils/dataset.py b/utils/dataset.py
index e9984ad..c93b535 100644
--- a/utils/dataset.py
+++ b/utils/dataset.py
@@ -101,6 +101,7 @@ class Dataset(data.Dataset):
         return len(self.filenames)
 
     def load_image(self, i):
+        # FIXME: for png maybe have something different with jpg
         image = cv2.imread(self.filenames[i])
         if image is None:
             raise FileNotFoundError(f"Image Not Found {self.filenames[i]}")
@@ -165,6 +166,13 @@ class Dataset(data.Dataset):
                 x2b = min(shape[1], x2a - x1a)
                 y2b = min(y2a - y1a, shape[0])
 
+            if (
+                isinstance(x1a, type(None))
+                or isinstance(y1a, type(None))
+                or isinstance(x1b, type(None))
+                or isinstance(y1b, type(None))
+            ):
+                raise ValueError("Mosaic calculation error")
             pad_w = x1a - x1b
             pad_h = y1a - y1b
             image4[y1a:y2a, x1a:x2a] = image[y1b:y2b, x1b:x2b]
@@ -205,7 +213,7 @@ class Dataset(data.Dataset):
         path = f"{os.path.dirname(filenames[0])}.cache"
         if os.path.exists(path):
             # XXX: temporarily disable cache
-            os.remove(path)
+            # os.remove(path)
             pass
             # return torch.load(path, weights_only=False)
         x = {}
@@ -217,7 +225,10 @@ class Dataset(data.Dataset):
                     image.verify()  # PIL verify
                 shape = image.size  # image size
                 assert (shape[0] > 9) & (shape[1] > 9), f"image size {shape} <10 pixels"
-                assert image.format.lower() in FORMATS, f"invalid image format {image.format}"
+                if image.format:
+                    assert image.format.lower() in FORMATS, f"invalid image format {image.format}"
+                else:
+                    assert any(filename.lower().endswith(f".{x}") for x in FORMATS), "unknown image format"
 
                 # verify labels
                 a = f"{os.sep}images{os.sep}"