Upload app.py

app.py

@@ -130,15 +130,24 @@ class ImgLoader:
             transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
         ])
 
-    def load(self, path):
-        ori_img = cv2.imread(path)
-        img = copy.deepcopy(ori_img[:, :, ::-1])
-        img = Image.fromarray(img)
+    def load(self, input_img):
+        if isinstance(input_img, str):
+            ori_img = cv2.imread(input_img)
+            img = Image.fromarray(cv2.cvtColor(ori_img, cv2.COLOR_BGR2RGB))
+        elif isinstance(input_img, Image.Image):
+            img = input_img
+        else:
+            raise ValueError("Image invalide")
+
+        if img.mode != "RGB":
+            img = img.convert("RGB")
+
         return self.transform(img).unsqueeze(0)
 
-def cal_backward(out):
+def cal_backward(out) -> dict:
     target_layer_names = ['layer1', 'layer2', 'layer3', 'layer4',
                           'FPN1_layer1', 'FPN1_layer2', 'FPN1_layer3', 'FPN1_layer4', 'comb_outs']
+
     sum_out = None
     for name in target_layer_names:
         tmp_out = out[name].mean(1) if name != "comb_outs" else out[name]
@@ -156,9 +165,12 @@ def cal_backward(out):
         V = V - min(V)
         V = V / sum(V)
 
-        top5 = np.argsort(-V)[:5]
-        accs = -np.sort(-V)[:5]
-    return [f"{classes_list[int(cls)]}: {acc*100:.2f}%" for cls, acc in zip(top5, accs)]
+        top5_indices = np.argsort(-V)[:5]
+        top5_scores = -np.sort(-V)[:5]
+
+    # Construction du dictionnaire pour gr.Label
+    top5_dict = {classes_list[int(idx)]: float(f"{score:.4f}") for idx, score in zip(top5_indices, top5_scores)}
+    return top5_dict
 
 # === Chargement du modèle
 model = build_model("weights.pt")
@@ -167,20 +179,33 @@ img_loader = ImgLoader(data_size)
 def predict_image(image: Image.Image) -> List[str]:
     global features, grads, module_id_mapper
     features, grads, module_id_mapper = {}, {}, {}
-    image_path = "temp_image.jpg"
-    image.save(image_path)
-    img_tensor = img_loader.load(image_path)
+
+    if image is None:
+        return {}
+#        raise ValueError("Aucune image reçue. Vérifie l'entrée.")
+
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+
+    img_tensor = img_loader.load(image)
     out = model(img_tensor)
     return cal_backward(out)
 
 # === Interface Gradio
-demo = gr.Interface(
-    fn=predict_image,
-    inputs=gr.Image(type="pil"),
-    outputs=gr.Textbox(label="Prédictions top 5"),
-    title="Classification OGSO",
-    description="Détection de la catégorie de quincaillerie via Swin Transformer avec architecture personnalisée"
-)
+with gr.Blocks() as demo:
+    with gr.Row():
+        with gr.Column():
+            with gr.Tab("Téléversement"):
+                upload_input = gr.Image(type="pil", label="Image téléchargée", sources=["upload"], show_label=True)
+            with gr.Tab("Webcam"):
+                webcam_input = gr.Image(type="pil", label="Webcam", sources=["webcam"], show_label=True)
+        with gr.Column():
+            output = gr.Label(num_top_classes=5, label="Prédiction")
+
+    # Connexion des callbacks
+    upload_input.change(fn=predict_image, inputs=upload_input, outputs=output)
+    webcam_input.change(fn=predict_image, inputs=webcam_input, outputs=output)
+
 
 if __name__ == "__main__":
     demo.launch()