Update app.py

app.py

@@ -3,6 +3,7 @@ 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",
@@ -14,24 +15,6 @@ model_names = [
     "0-ma/resnet-geometric-shapes",
     "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_dir = "./example"
-example_images = []
-for example_image in os.list_dir(example_dir):
-    example_images+= [os.path.join(example_dir,example_image)]
-
-example_images = [os.path.join(example_dir,example_image) for example_image in os.listdir(example_dir)]
-
-#labels = [example.split("_")[1].split(".")[0] for example in example_images]
 labels = [
     'None',
     'Circle',
@@ -40,6 +23,8 @@ labels = [
     'Pentagone',
     'Hexagone'
 ]
+example_dir = "./example"
+example_images = [os.path.join(example_dir,example_image) for example_image in os.listdir(example_dir)]
 
 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}