import os # 1. Toujours importer `spaces` EN PREMIER avant torch (évite le RuntimeError ZeroGPU) try: import spaces HAS_SPACES = True except ImportError: HAS_SPACES = False import cv2 import numpy as np from PIL import Image import torch import torch.nn as nn from torchvision import transforms import gradio as gr # ===================================================================== # 1. CONFIGURATION & ARCHITECTURE DU MODÈLE # ===================================================================== IMG_SIZE = 256 MODEL_WEIGHTS_PATH = "best_casting_model.pth" DEFAULT_THRESHOLD = 0.40 CLASS_NAMES = { 0: "Conforme (OK)", 1: "Défectueuse (DÉFAUT)" } # Initialisation sur CPU au démarrage pour ne pas déclencher d'erreur NVML/ZeroGPU. # (Un CNN 4-blocs 256x256 s'exécute en ~15ms sur CPU). device = torch.device("cpu") class CustomIndustrialCNN(nn.Module): def __init__(self): super(CustomIndustrialCNN, self).__init__() self.features = nn.Sequential( # Bloc 1 : 256x256 -> 128x128 nn.Conv2d(1, 32, kernel_size=3, padding=1), nn.BatchNorm2d(32), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=2, stride=2), # Bloc 2 : 128x128 -> 64x64 nn.Conv2d(32, 64, kernel_size=3, padding=1), nn.BatchNorm2d(64), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=2, stride=2), # Bloc 3 : 64x64 -> 32x32 nn.Conv2d(64, 128, kernel_size=3, padding=1), nn.BatchNorm2d(128), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=2, stride=2), # Bloc 4 : 32x32 -> 16x16 nn.Conv2d(128, 256, kernel_size=3, padding=1), nn.BatchNorm2d(256), nn.ReLU(inplace=True), nn.MaxPool2d(kernel_size=2, stride=2), ) self.global_pool = nn.AdaptiveAvgPool2d((1, 1)) self.classifier = nn.Sequential( nn.Flatten(), nn.Linear(256, 128), nn.ReLU(inplace=True), nn.Dropout(0.5), nn.Linear(128, 1) ) def forward(self, x): x = self.features(x) x = self.global_pool(x) x = self.classifier(x) return x # ===================================================================== # 2. CHARGEMENT DU MODÈLE ET DES POIDS # ===================================================================== model = CustomIndustrialCNN().to(device) if os.path.exists(MODEL_WEIGHTS_PATH): checkpoint = torch.load(MODEL_WEIGHTS_PATH, map_location=device, weights_only=False) if isinstance(checkpoint, dict) and "model_state_dict" in checkpoint: model.load_state_dict(checkpoint["model_state_dict"]) if "optimal_threshold" in checkpoint: DEFAULT_THRESHOLD = float(checkpoint["optimal_threshold"]) else: model.load_state_dict(checkpoint) print(f"Poids chargés avec succès depuis '{MODEL_WEIGHTS_PATH}' sur {device}.") else: print(f"ATTENTION : '{MODEL_WEIGHTS_PATH}' introuvable.") model.eval() # ===================================================================== # 3. PRÉTRAITEMENT # ===================================================================== val_test_transforms = transforms.Compose([ transforms.Grayscale(num_output_channels=1), transforms.Resize((IMG_SIZE, IMG_SIZE)), transforms.ToTensor(), transforms.Normalize(mean=[0.5], std=[0.5]) ]) # ===================================================================== # 4. EXPLICABILITÉ : GRAD-CAM # ===================================================================== def find_last_conv_layer(net): last_conv = None for module in net.modules(): if isinstance(module, nn.Conv2d): last_conv = module return last_conv def generate_gradcam(model, input_tensor, original_pil_img): activations = [] gradients = [] target_layer = find_last_conv_layer(model) if target_layer is None: return np.array(original_pil_img.convert("RGB")), None def forward_hook(module, inp, out): activations.append(out) def backward_hook(module, grad_in, grad_out): gradients.append(grad_out[0]) handle_f = target_layer.register_forward_hook(forward_hook) handle_b = target_layer.register_full_backward_hook(backward_hook) try: model.zero_grad() input_tensor.requires_grad_(True) output = model(input_tensor) logit = output.squeeze() prob_defect = torch.sigmoid(logit).item() logit.backward() grads = gradients[0].detach().cpu().numpy()[0] acts = activations[0].detach().cpu().numpy()[0] weights = np.mean(grads, axis=(1, 2)) cam = np.zeros(acts.shape[1:], dtype=np.float32) for i, w in enumerate(weights): cam += w * acts[i] cam = np.maximum(cam, 0) if np.max(cam) > 0: cam = cam / np.max(cam) orig_w, orig_h = original_pil_img.size cam_resized = cv2.resize(cam, (orig_w, orig_h)) heatmap = np.uint8(255 * cam_resized) heatmap_colored = cv2.applyColorMap(heatmap, cv2.COLORMAP_JET) heatmap_colored = cv2.cvtColor(heatmap_colored, cv2.COLOR_BGR2RGB) orig_rgb = np.array(original_pil_img.convert("L").convert("RGB")) overlay = cv2.addWeighted(orig_rgb, 0.55, heatmap_colored, 0.45, 0) return overlay, prob_defect finally: handle_f.remove() handle_b.remove() # ===================================================================== # 5. FONCTION D'INFÉRENCE GRADIO # ===================================================================== def inspect_casting_part(image: Image.Image, threshold: float): if image is None: return None, None, "
Veuillez charger une image de pièce à inspecter.
" gray_img = image.convert("L") input_tensor = val_test_transforms(gray_img).unsqueeze(0).to(device) gradcam_overlay, prob_defect = generate_gradcam(model, input_tensor, image) if prob_defect is None: with torch.no_grad(): logit = model(input_tensor).squeeze() prob_defect = torch.sigmoid(logit).item() prob_ok = 1.0 - prob_defect confidences = { CLASS_NAMES[1]: float(prob_defect), CLASS_NAMES[0]: float(prob_ok) } is_defective = prob_defect >= threshold if is_defective: badge_color = "#e74c3c" bg_color = "#fdecea" status_title = "⚠️ PIÈCE DÉFECTUEUSE DÉTECTÉE (REJET)" action_msg = ( f"La probabilité d'anomalie ({prob_defect*100:.2f}%) dépasse le seuil de tolérance fixé à " f"{threshold*100:.1f}%.
Action recommandée : Écarter la pièce de la ligne et vérifier la zone " f"rouge/jaune sur la carte Grad-CAM (soufflure, bavure d'usinage ou retrait)." ) else: badge_color = "#2ecc71" bg_color = "#eafaf1" status_title = "✅ PIÈCE CONFORME (VALIDATION)" action_msg = ( f"La probabilité d'anomalie ({prob_defect*100:.2f}%) est inférieure au seuil critique de " f"{threshold*100:.1f}%.
Action recommandée : Autoriser le passage à l'étape d'assemblage suivante." ) diagnostic_html = f"""

{status_title}

{action_msg}

""" return confidences, gradcam_overlay, diagnostic_html # ===================================================================== # 6. INTERFACE GRADIO (Compatible Gradio 6.0+) # ===================================================================== with gr.Blocks(title="Contrôle Qualité - Vision Industrielle 4.0") as demo: gr.Markdown( """ # ⚙️ Contrôle Qualité par Vision Industrielle (Industrie 4.0) **Détection automatique de défauts de moulage sur turbines de pompes submersibles avec explicabilité Grad-CAM.** Chargez une image issue de la caméra d'inspection verticale pour obtenir le diagnostic instantané et localiser les anomalies physiques. """ ) with gr.Row(): with gr.Column(scale=1): input_image = gr.Image( type="pil", label="Image de la pièce (Caméra d'inspection 300x300 / 256x256)", height=320 ) threshold_slider = gr.Slider( minimum=0.05, maximum=0.95, value=DEFAULT_THRESHOLD, step=0.05, label="Seuil de décision (Probabilité minimale de défaut)", info="Abaissez ce seuil (ex: 0.30) pour minimiser les Faux Négatifs (aucune pièce défectueuse tolérée)." ) inspect_btn = gr.Button("🔍 Lancer l'Inspection Qualité", variant="primary", size="lg") with gr.Column(scale=1): diagnostic_box = gr.HTML(label="Décision Industrielle") output_label = gr.Label( num_top_classes=2, label="Scores de Confiance du Modèle" ) gradcam_output = gr.Image( label="Carte d'Activation Grad-CAM (Localisation du défaut)", height=320 ) inspect_btn.click( fn=inspect_casting_part, inputs=[input_image, threshold_slider], outputs=[output_label, gradcam_output, diagnostic_box] ) input_image.change( fn=inspect_casting_part, inputs=[input_image, threshold_slider], outputs=[output_label, gradcam_output, diagnostic_box] ) threshold_slider.release( fn=inspect_casting_part, inputs=[input_image, threshold_slider], outputs=[output_label, gradcam_output, diagnostic_box] ) if os.path.exists("examples") and len(os.listdir("examples")) > 0: example_files = [ [os.path.join("examples", f), DEFAULT_THRESHOLD] for f in sorted(os.listdir("examples")) if f.lower().endswith((".jpeg", ".jpg", ".png")) ] if example_files: gr.Examples( examples=example_files, inputs=[input_image, threshold_slider], outputs=[output_label, gradcam_output, diagnostic_box], fn=inspect_casting_part, cache_examples=False, label="Échantillons de Test (Cliquez sur une image pour tester)" ) if __name__ == "__main__": # En Gradio 6.0+, `theme` est passé dans launch() demo.launch(theme=gr.themes.Soft())