fix: no calibration datasets bug (#277)

Seeed-Studio · Jan 16, 2025 · 905e586 · 905e586
1 parent 89716d3
commit 905e586
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/tools/export.py b/tools/export.py
@@ -361,7 +361,9 @@ def export_tflite(onnx_path: str, img_shape, img_path):
 
     def representative_dataset():
         datasets = [
-            osp.join(img_path, i) for i in os.listdir(img_path) if i.endswith(".jpg")
+            osp.join(img_path, i)
+            for i in os.listdir(img_path)
+            if i.lower().endswith((".jpg", ".jpeg", ".png"))
         ]
         for ps in tqdm(datasets[:300]):
             img = cv2.imread(ps)