Update app.py

app.py

@@ -1,11 +1,37 @@
 import gradio as gr
-import skimage
+import numpy as np
+from PIL import Image
+from transformers import AutoImageProcessor, AutoModelForImageClassification 
+import requests
+labels =  [
+    "None",
+    "Circle",
+    "Triangle",
+    "Square",
+    "Pentagon",
+    "Hexagon"
+] 
+#images = [Image.open(requests.get("https://raw.githubusercontent.com/0-ma/geometric-shape-detector/main/input/exemple_circle.jpg", stream=True).raw), 
+#        Image.open(requests.get("https://raw.githubusercontent.com/0-ma/geometric-shape-detector/main/input/exemple_pentagone.jpg", stream=True).raw)]
+feature_extractor = AutoImageProcessor.from_pretrained('0-ma/swin-geometric-shapes-tiny')
+model = AutoModelForImageClassification.from_pretrained('0-ma/swin-geometric-shapes-tiny')
+
+   
+print(predicted_labels)
+
+
 
 labels = []
 def predict(img):
     img = PILImage.create(img)
-    pred,pred_idx,probs = detector.predict(img)
-    return {labels[i]: float(probs[i]) for i in range(len(labels))}
+
+    inputs = feature_extractor(images=images, return_tensors="pt")
+    logits = model(**inputs)['logits'].cpu().detach().numpy()
+    predictions = np.argmax(logits, axis=1) 
+    predicted_labels = [labels[prediction] for prediction in predictions]
+
+
+    return {"predicted_labels" : predicted_labels , "predictions": predictions}
  
 title = "Geometric Shape Classifier"
 description = "A geometric shape setector."