A Gradio app that classifies product images as defective or non-defective using transfer learning — upload any product image and get instant defect classification with a confidence score.
In manufacturing, product defects that slip through quality control are expensive — rework costs, warranty claims, customer complaints, and in safety-critical industries, regulatory consequences. Traditional quality control relies on human visual inspection — slow, inconsistent, and not scalable.
Computer vision changes this. A trained model can inspect thousands of parts per minute, 24/7, with consistent accuracy that doesn't degrade after an 8-hour shift. This project builds exactly that — a visual defect classifier that works from a single uploaded image.
The dataset contains images of casting products (pump impellers) in two categories: defective and non-defective. Casting defects include roughness, blowholes, cracks, and incomplete fills — all visible in the product surface image.
| Split | OK (Non-Defective) | Defective | Total |
|---|---|---|---|
| Training | 2,875 images | 781 images | 3,656 |
| Test | 453 images | 267 images | 720 |
| Total | 3,328 | 1,048 | 4,376 |
Training a deep CNN from scratch on 4,376 images would produce a mediocre model. Transfer learning takes a model already trained on millions of images (ImageNet) and fine-tunes it on the specific defect classification task. The result: high accuracy with limited data and compute.
import torchvision.models as models import torch.nn as nn # Load pre-trained ResNet-50 model = models.resnet50(pretrained=True) # Freeze all layers for param in model.parameters(): param.requires_grad = False # Replace final layer with binary classifier num_features = model.fc.in_features model.fc = nn.Sequential( nn.Linear(num_features, 256), nn.ReLU(), nn.Dropout(0.4), nn.Linear(256, 2) # defective / non-defective ) # Data augmentation for training train_transforms = transforms.Compose([ transforms.RandomHorizontalFlip(), transforms.RandomRotation(15), transforms.ColorJitter(brightness=0.2, contrast=0.2), transforms.ToTensor(), transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) ])
The app is built with Gradio — the fastest way to turn a PyTorch model into an interactive web demo. No frontend code required. The interface is intentionally simple: upload an image, get a result.
import gradio as gr import torch from PIL import Image def classify_image(image): # Preprocess tensor = preprocess(image).unsqueeze(0) # Predict with torch.no_grad(): output = model(tensor) probs = torch.nn.functional.softmax(output[0], dim=0) confidence = probs.max().item() label = "Non-Defective ✅" if probs.argmax() == 0 else "Defective ❌" return f"{label} — {confidence*100:.1f}% confident" demo = gr.Interface( fn=classify_image, inputs=gr.Image(type="pil", label="Upload Product Image"), outputs=gr.Textbox(label="Classification Result"), title="Manufacturing Defect Detector", description="Upload a product image to check for manufacturing defects." ) demo.launch()