Update app.py

app.py

@@ -2,7 +2,7 @@ import gradio as gr
 import numpy as np
 from PIL import Image
 from transformers import AutoImageProcessor, AutoModelForImageClassification
-
+import os
 model_names = [
     "0-ma/swin-geometric-shapes-tiny",
     "0-ma/mobilenet-v2-geometric-shapes",
@@ -15,16 +15,29 @@ model_names = [
     "0-ma/vit-geometric-shapes-tiny",
 ]
 
-example_images = [
-    'example/1_None.jpg',
-    'example/2_Circle.jpg',
-    'example/3_Triangle.jpg',
-    'example/4_Square.jpg',
-    'example/5_Pentagone.jpg',
-    'example/6_Hexagone.jpg'
-]
+# example_images = [
+#     'example/1_None.jpg',
+#     'example/2_Circle.jpg',
+#     'example/3_Triangle.jpg',
+#     'example/4_Square.jpg',
+#     'example/5_Pentagone.jpg',
+#     'example/6_Hexagone.jpg'
+# ]
 
-labels = [example.split("_")[1].split(".")[0] for example in example_images]
+example_dir = "./example"
+example_images = []
+for example_image in os.list_dir(example_dir):
+    example_images+= [os.path.join(example_dir,example_image)]
+
+#labels = [example.split("_")[1].split(".")[0] for example in example_images]
+labels = [
+    'None',
+    'Circle',
+    'Triangle',
+    'Square',
+    'Pentagone',
+    'Hexagone'
+]
 
 feature_extractors = {model_name: AutoImageProcessor.from_pretrained(model_name) for model_name in model_names}
 classification_models = {model_name: AutoModelForImageClassification.from_pretrained(model_name) for model_name in model_names}