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 from pim_module import PluginMoodel # Assure-toi que ce fichier est présent # === 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 = "10nhim7twcKEGB16jVilPQGW0CrKo4jOY" 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()