feat: input files check (filter the unreadable files)

datasets/simple_extractor_dataset.py

@@ -51,28 +51,31 @@ def __getitem__(self, index):
         img_name = self.file_list[index]
         img_path = os.path.join(self.root, img_name)
         img = cv2.imread(img_path, cv2.IMREAD_COLOR)
-        h, w, _ = img.shape
+        if np.any(img):
+            h, w, _ = img.shape
 
-        # Get person center and scale
-        person_center, s = self._box2cs([0, 0, w - 1, h - 1])
-        r = 0
-        trans = get_affine_transform(person_center, s, r, self.input_size)
-        input = cv2.warpAffine(
-            img,
-            trans,
-            (int(self.input_size[1]), int(self.input_size[0])),
-            flags=cv2.INTER_LINEAR,
-            borderMode=cv2.BORDER_CONSTANT,
-            borderValue=(0, 0, 0))
+            # Get person center and scale
+            person_center, s = self._box2cs([0, 0, w - 1, h - 1])
+            r = 0
+            trans = get_affine_transform(person_center, s, r, self.input_size)
+            input = cv2.warpAffine(
+                img,
+                trans,
+                (int(self.input_size[1]), int(self.input_size[0])),
+                flags=cv2.INTER_LINEAR,
+                borderMode=cv2.BORDER_CONSTANT,
+                borderValue=(0, 0, 0))
 
-        input = self.transform(input)
-        meta = {
-            'name': img_name,
-            'center': person_center,
-            'height': h,
-            'width': w,
-            'scale': s,
-            'rotation': r
-        }
+            input = self.transform(input)
+            meta = {
+                'name': img_name,
+                'center': person_center,
+                'height': h,
+                'width': w,
+                'scale': s,
+                'rotation': r
+            }
+
+            return input, meta
+        return 0
 
-        return input, meta

simple_extractor.py

@@ -126,6 +126,8 @@ def main():
     palette = get_palette(num_classes)
     with torch.no_grad():
         for idx, batch in enumerate(tqdm(dataloader)):
+            if len(batch) == 1:
+                continue
             image, meta = batch
             img_name = meta['name'][0]
             c = meta['center'].numpy()[0]