Upload app.py

app.py

@@ -0,0 +1,247 @@
+import gdown
+import os
+import torch
+import requests
+import numpy as np
+import numpy.matlib
+import copy
+import cv2
+from PIL import Image
+from typing import List
+import timm
+import gradio as gr
+import torchvision.transforms as transforms
+import zipfile
+
+from pim_module import PluginMoodel  # Assure-toi que ce fichier est présent
+
+# === Décompression automatique du dossier imgs ===
+if not os.path.exists("imgs") and os.path.exists("imgs.zip"):
+    print("Décompression du dossier imgs...")
+    with zipfile.ZipFile("imgs.zip", 'r') as zip_ref:
+        zip_ref.extractall(".")
+    print("Décompression terminée !")
+
+# === Téléchargement automatique depuis Google Drive ===
+if not os.path.exists("weights.pt"):
+    print("Téléchargement des poids depuis Google Drive avec gdown...")
+    file_id = "17RxaEfYeQVKKXThOwqDWM6mHdII5tMpY"
+    url = f"https://drive.google.com/uc?id={file_id}"
+    gdown.download(url, "weights.pt", quiet=False)
+
+
+
+# === Classes
+classes_list = [
+    "Ferrage_Et_Accessoires_Anti_Fausse_Manoeuvre",
+    "Ferrage_Et_Accessoires_Busettes",
+    "Ferrage_Et_Accessoires_Butees",
+    "Ferrage_Et_Accessoires_Chariots",
+    "Ferrage_Et_Accessoires_Charniere",
+    "Ferrage_Et_Accessoires_Compas_limiteur",
+    "Ferrage_Et_Accessoires_Cylindres",
+    "Ferrage_Et_Accessoires_Gaches",
+    "Ferrage_Et_Accessoires_Renvois_D_Angle",
+    "Joints_Et_Consommables_Equerres_Aluminium_Moulees",
+    "Joints_Et_Consommables_Visserie_Inox_Alu",
+    "Poignee_Carre_7_mm",
+    "Poignee_Carre_8_mm",
+    "Poignee_Cremone",
+    "Poignee_Cuvette",
+    "Poignee_De_Tirage",
+    "Poignee_Pour_Levant_Coulissant",
+    "Serrure_Cremone_Multipoints",
+    "Serrure_Cuvette",
+    "Serrure_Gaches",
+    "Serrure_Loqueteau",
+    "Serrure_Pene_Crochet",
+    "Serrure_Pour_Porte",
+    "Serrure_Tringles"
+]
+
+short_classes_list = [
+    "Anti-fausse-manoeuvre",
+    "Busettes",
+    "Butées",
+    "Chariots",
+    "Charnière",
+    "Compas-limiteur",
+    "Cylindres",
+    "Gaches",
+    "Renvois d'angle",
+    "Equerres aluminium moulées",
+    "Visserie inox alu",
+    "Poignée carré  7 mm",
+    "Poignée carré 8 mm",
+    "Poignée crémone",
+    "Poignée cuvette",
+    "Poignée de tirage",
+    "Poignée pour levant coulissant",
+    "Serrure crémone multipoints",
+    "Serrure cuvette",
+    "Serrure gaches",
+    "Loqueteau",
+    "Serrure pene crochet",
+    "Serrure pour porte",
+    "Serrure tringles",
+]
+
+data_size = 384
+fpn_size = 1536
+num_classes = 24
+num_selects = {'layer1': 256, 'layer2': 128, 'layer3': 64, 'layer4': 32}
+features, grads, module_id_mapper = {}, {}, {}
+
+def forward_hook(module, inp_hs, out_hs):
+    layer_id = len(features) + 1
+    module_id_mapper[module] = layer_id
+    features[layer_id] = {"in": inp_hs, "out": out_hs}
+
+def backward_hook(module, inp_grad, out_grad):
+    layer_id = module_id_mapper[module]
+    grads[layer_id] = {"in": inp_grad, "out": out_grad}
+
+def build_model(path: str):
+    backbone = timm.create_model('swin_large_patch4_window12_384_in22k', pretrained=True)
+    model = PluginMoodel(
+        backbone=backbone,
+        return_nodes=None,
+        img_size=data_size,
+        use_fpn=True,
+        fpn_size=fpn_size,
+        proj_type="Linear",
+        upsample_type="Conv",
+        use_selection=True,
+        num_classes=num_classes,
+        num_selects=num_selects,
+        use_combiner=True,
+        comb_proj_size=None
+    )
+    ckpt = torch.load(path, map_location="cpu", weights_only=False)
+    model.load_state_dict(ckpt["model"], strict=False)
+    model.eval()
+
+    for layer in [0, 1, 2, 3]:
+        model.backbone.layers[layer].register_forward_hook(forward_hook)
+        model.backbone.layers[layer].register_full_backward_hook(backward_hook)
+
+    for i in range(1, 5):
+        getattr(model.fpn_down, f'Proj_layer{i}').register_forward_hook(forward_hook)
+        getattr(model.fpn_down, f'Proj_layer{i}').register_full_backward_hook(backward_hook)
+        getattr(model.fpn_up, f'Proj_layer{i}').register_forward_hook(forward_hook)
+        getattr(model.fpn_up, f'Proj_layer{i}').register_full_backward_hook(backward_hook)
+
+    return model
+
+class ImgLoader:
+    def __init__(self, img_size):
+        self.transform = transforms.Compose([
+            transforms.Resize((510, 510), Image.BILINEAR),
+            transforms.CenterCrop((img_size, img_size)),
+            transforms.ToTensor(),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+        ])
+
+    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) -> 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]
+        tmp_out = torch.softmax(tmp_out, dim=-1)
+        sum_out = tmp_out if sum_out is None else sum_out + tmp_out
+
+    with torch.no_grad():
+        smax = torch.softmax(sum_out, dim=-1)
+        A = np.transpose(np.matlib.repmat(smax[0], num_classes, 1)) - np.eye(num_classes)
+        _, _, V = np.linalg.svd(A, full_matrices=True)
+        V = V[num_classes - 1, :]
+        if V[0] < 0:
+            V = -V
+        V = np.log(V)
+        V = V - min(V)
+        V = V / sum(V)
+
+        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")
+img_loader = ImgLoader(data_size)
+
+
+
+def predict_image(image: Image.Image):
+    global features, grads, module_id_mapper
+    features, grads, module_id_mapper = {}, {}, {}
+
+    if image is None:
+        return {}
+#        raise ValueError("Aucune image reçue. Vérifie l'entrée.")
+
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+
+    image_path = "temp.jpg"
+    image.save(image_path)
+
+    img_tensor = img_loader.load(image_path)
+    out = model(img_tensor)
+    top5_dict = cal_backward(out)  # {classe: score}
+
+    gallery_outputs = []
+    for idx, class_name in enumerate(list(top5_dict.keys())):
+        images = [
+            (f"imgs/{class_name}/{class_name}_0001.jpg", f"Exemple {class_name} 1"),
+            (f"imgs/{class_name}/{class_name}_0002.jpg", f"Exemple {class_name} 2"),
+            (f"imgs/{class_name}/{class_name}_0003.jpg", f"Exemple {class_name} 3"),
+        ]
+        gallery_outputs.append(images)
+
+    return top5_dict, *gallery_outputs
+
+
+# === Interface Gradio
+with gr.Blocks(css="""
+.gr-image-upload { display: none !important }
+.gallery-container .gr-box { height: auto !important; padding: 0 !important; }
+""") as demo:
+    with gr.Row():
+        with gr.Column(scale=1):
+            with gr.Tab("Téléversement"):
+                image_input_upload = gr.Image(type="pil", label="Image à classer (upload)", sources=["upload"])
+            with gr.Tab("Webcam"):
+                image_input_webcam = gr.Image(type="pil", label="Image à classer (webcam)", sources=["webcam"])
+
+        with gr.Column(scale=1.5):
+            label_output = gr.Label(label="Prédictions")
+            gallery_outputs = [
+                gr.Gallery(label=f"", columns=3, height=300, container=True, elem_classes=["gallery-container"])
+                for i in range(5)
+            ]
+
+    image_input_upload.change(fn=predict_image, inputs=image_input_upload, outputs=[label_output] + gallery_outputs)
+    image_input_webcam.change(fn=predict_image, inputs=image_input_webcam, outputs=[label_output] + gallery_outputs)
+
+if __name__ == "__main__":
+    demo.launch()
+