Spaces:

anasAmchaar
/

casting-quality-control

Runtime error

Upload app.py with huggingface_hub

26c9bc9 verified 5 days ago

11.1 kB

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, "<div style='color:gray;'>Veuillez charger une image de pièce à inspecter.</div>"

    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 (<b>{prob_defect*100:.2f}%</b>) dépasse le seuil de tolérance fixé à "
            f"<b>{threshold*100:.1f}%</b>.<br><b>Action recommandée :</b> É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 (<b>{prob_defect*100:.2f}%</b>) est inférieure au seuil critique de "
            f"<b>{threshold*100:.1f}%</b>.<br><b>Action recommandée :</b> Autoriser le passage à l'étape d'assemblage suivante."
        )

    diagnostic_html = f"""
    <div style="padding: 18px; border-radius: 8px; background-color: {bg_color}; border-left: 6px solid {badge_color}; font-family: sans-serif;">
        <h3 style="margin-top: 0; color: {badge_color};">{status_title}</h3>
        <p style="margin-bottom: 0; color: #2c3e50; font-size: 14px; line-height: 1.5;">{action_msg}</p>
    </div>
    """

    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())

	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, "<div style='color:gray;'>Veuillez charger une image de pièce à inspecter.</div>"

	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 (<b>{prob_defect*100:.2f}%</b>) dépasse le seuil de tolérance fixé à "
	f"<b>{threshold*100:.1f}%</b>.<br><b>Action recommandée :</b> É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 (<b>{prob_defect*100:.2f}%</b>) est inférieure au seuil critique de "
	f"<b>{threshold*100:.1f}%</b>.<br><b>Action recommandée :</b> Autoriser le passage à l'étape d'assemblage suivante."
	)

	diagnostic_html = f"""
	<div style="padding: 18px; border-radius: 8px; background-color: {bg_color}; border-left: 6px solid {badge_color}; font-family: sans-serif;">
	<h3 style="margin-top: 0; color: {badge_color};">{status_title}</h3>
	<p style="margin-bottom: 0; color: #2c3e50; font-size: 14px; line-height: 1.5;">{action_msg}</p>
	</div>
	"""

	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())