samarth-24/PCB_Defect_Detection
0
1import gradio as gr2import numpy as np3from PIL import Image4import torch5import torch.nn.functional as F6import torchvision.transforms as transforms7from torchvision import models8import cv2 # Needed for drawing/color conversion9import os # Needed for checking file existence10import platform # For OS detection11import pathlib # For path handling and patching12 13# --- Configuration & Setup ---14# print(f"PyTorch version: {torch._version_}")15# print(f"CUDA available: {torch.cuda.is_available()}")16device = torch.device("cuda" if torch.cuda.is_available() else "cpu")17# print(f"Using device: {device}")18# print(f"Operating System: {platform.system()}")19 20# Paths to model weights (relative to script location)21classifier_path = "pcb_classifier.pth"22yolo_weights_path = 'best.pt'23 24# --- Classifier Loading ---25classifier = None # Initialize26class_names = ['pcb','non_pcb'] # Make sure order matches training27 28try:29 print("Loading Classifier...")30 # Ensure the weights file exists before trying to load model structure31 if not os.path.exists(classifier_path):32 raise FileNotFoundError(f"Classifier weights file not found at: {classifier_path}")33 34 # Load ResNet18 structure35 # Use weights=None for torchvision 0.13+36 # Use pretrained=False for older versions37 try:38 classifier = models.resnet18(weights=None)39 except TypeError:40 print("Using pretrained=False for older torchvision.")41 classifier = models.resnet18(pretrained=False)42 43 # Modify the final layer for 2 classes44 num_ftrs = classifier.fc.in_features45 classifier.fc = torch.nn.Linear(num_ftrs, len(class_names))46 47 # Load the trained weights48 classifier.load_state_dict(torch.load(classifier_path, map_location=device))49 classifier.to(device)50 classifier.eval() # Set to evaluation mode51 print(f"Classifier loaded successfully from {classifier_path}")52 53except FileNotFoundError as e:54 print(f"Error: {e}")55 print("PCB classification will not be available. Please ensure 'pcb_classifier.pth' is in the correct directory.")56 # Optionally exit if classifier is essential:57 # exit()58except Exception as e:59 print(f"Error loading classifier: {e}")60 print("PCB classification will not be available.")61 # Optionally exit:62 # exit()63 64# Image transforms for Classifier (ensure these match classifier training)65transform = transforms.Compose([66 transforms.Resize((224, 224)),67 transforms.ToTensor(),68 # Example normalization, adjust if your training used different values69 transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])70 # Or use the normalization from the original snippet if that's correct:71 # transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])72])73# --- End Classifier Setup ---74 75 76# --- Load YOLOv8 Model using torch.hub ---77yolo_model = None # Initialize as None78yolo_class_names = []79 80# # Apply the pathlib workaround for Windows before loading YOLOv881# if platform.system() == "Windows":82# print("Applying pathlib patch for Windows...")83# try:84# # Store original PosixPath85# temp_posix_path = pathlib.PosixPath86# # Temporarily override PosixPath with WindowsPath87# pathlib.PosixPath = pathlib.WindowsPath88# print("Pathlib patch applied.")89# except AttributeError:90# # Handle cases where PosixPath might not be directly available/needed to patch91# print("Could not apply pathlib patch (PosixPath not directly found or needed). Continuing...")92# pass # Continue without the patch if it fails93 94 95# --- Load YOLOv8 model with ultralytics ---96try:97 from ultralytics import YOLO98 print("Loading YOLOv8 model via ultralytics...")99 yolo_model = YOLO(yolo_weights_path) # path to best.pt100 yolo_model.to(device)101 yolo_model.eval()102 yolo_class_names = yolo_model.names103 print(f"YOLOv8 model loaded successfully from {yolo_weights_path}.")104except Exception as e:105 print(f"Error loading YOLOv8 via ultralytics: {e}")106 yolo_model = None107 108 109# # Only attempt to load YOLO if the weights file exists110# if os.path.exists(yolo_weights_path):111# try:112# print("Loading YOLOv8 model...")113# # Attempt to load the model114# # Set force_reload=True if you suspect cache issues and have internet115# yolo_model = torch.hub.load('ultralytics/YOLOv5', 'custom', path=yolo_weights_path, force_reload=False, trust_repo=True)116# yolo_model.to(device) # Move model to appropriate device117# yolo_model.eval() # Set to evaluation mode118# yolo_class_names = yolo_model.names # Get class names from the loaded YOLO model119# print(f"YOLOv8 model loaded successfully from {yolo_weights_path} via torch.hub.")120# print("YOLO class names:", yolo_class_names)121 122# # More specific error handling123# except Exception as e:124# print(f"Error loading YOLOv8 model via torch.hub: {e}")125# print(f"Ensure internet connection is available if cache is missing/corrupted, or if using force_reload=True.")126# print(f"Ensure '{yolo_weights_path}' is a valid YOLOv8 model file and dependencies (like ultralytics) are installed.")127# print("YOLOv8 defect detection will not be available.")128# yolo_model = None # Ensure it's None if loading fails129 130# # Optional: Restore the original PosixPath if needed later in the script131# # if platform.system() == "Windows" and 'temp_posix_path' in locals():132# # pathlib.PosixPath = temp_posix_path133# # print("Restored original pathlib.PosixPath")134# else:135# print(f"Warning: YOLOv8 weights file not found at: {yolo_weights_path}")136# print("YOLOv8 defect detection will not be available.")137# yolo_model = None138# # --- End YOLOv8 Loading ---139 140 141# --- Main Processing Function ---142def detect_defects(image_pil: Image.Image):143 """144 Classifies the image and runs YOLOv8 defect detection if classified as PCB.145 146 Args:147 image_pil: Input image as a PIL Image object.148 149 Returns:150 Tuple[PIL.Image, str]: Annotated image and results string.151 """152 if image_pil is None:153 return None, "Please upload an image."154 155 # Convert to RGB if necessary (e.g., for PNGs with alpha)156 if image_pil.mode != 'RGB':157 image_pil = image_pil.convert('RGB')158 159 annotated_img_pil = image_pil.copy() # Start with original image160 classification_result_text = "Classification: Model not loaded or failed."161 yolo_result_text = "YOLO: Skipped."162 163 # --- Step 1: Classify (if classifier loaded) ---164 if classifier:165 try:166 input_tensor = transform(image_pil).unsqueeze(0).to(device)167 with torch.no_grad():168 outputs = classifier(input_tensor)169 probs = F.softmax(outputs, dim=1)170 confidence, predicted = torch.max(probs, 1)171 class_label = class_names[predicted.item()]172 class_conf = confidence.item()173 classification_result_text = f"Classified as: {class_label} (Confidence: {class_conf:.4f})"174 print(classification_result_text) # Log to console175 except Exception as e:176 print(f"Error during classification: {e}")177 classification_result_text = f"Error during classification: {e}"178 class_label = None # Ensure YOLO doesn't run if classification fails179 else:180 # Classifier didn't load, cannot proceed with logic that depends on it181 return annotated_img_pil, "Classifier model failed to load. Cannot process image."182 183 # # --- Step 2: Run YOLOv8 only if PCB and model loaded ---184 # if class_label == 'pcb' and yolo_model is not None:185 # try:186 # with torch.no_grad():187 # inference_out = yolo_model(image_pil, imgsz=640, conf=0.5)188 189 # # Unpack list if needed190 # results = inference_out[0] if isinstance(inference_out, list) else inference_out191 192 # # Now render and convert193 # rendered = results.render()194 # if rendered:195 # annotated_img_np = rendered[0]196 # annotated_img_pil = Image.fromarray(197 # cv2.cvtColor(annotated_img_np, cv2.COLOR_BGR2RGB)198 # )199 200 # # Build the text summary from pandas DataFrame201 # detections = results.pandas().xyxy[0]202 # yolo_output_lines = []203 # if not detections.empty:204 # print(f"YOLOv8 detected {len(detections)} potential defects.")205 # for index, row in detections.iterrows():206 # conf = row['confidence']207 # name = row['name'] # Defect class name from the model208 # yolo_output_lines.append(f"- {name} (Confidence: {conf:.2f})")209 # yolo_result_text = "YOLOv8 Detections:\n" + "\n".join(yolo_output_lines)210 # else:211 # print("YOLOv8 detected no defects.")212 # yolo_result_text = "YOLOv8: No defects detected."213 214 # except Exception as e:215 # print(f"Error during YOLOv8 inference or processing: {e}")216 # yolo_result_text = f"Error during YOLOv8 processing: {e}"217 # # Keep original image in case of YOLO error, don't show corrupted render218 # annotated_img_pil = image_pil.copy()219 220 # else: # yolo_model is None221 # print("YOLOv8 model not loaded, skipping defect detection.")222 # yolo_result_text = "YOLOv8: Model not loaded, cannot perform defect detection."223 # # Keep original image224 # annotated_img_pil = image_pil.copy()225 226 # else: # class_label was 'non_pcb'227 # print("Skipping YOLOv8 - Image classified as non_pcb.")228 # yolo_result_text = "YOLOv8: Skipped (Input classified as non-PCB)."229 # # Keep original image230 # annotated_img_pil = image_pil.copy()231 232 # --- Step 2: Run YOLOv8 only if PCB and model loaded ---233 if class_label == 'pcb' and yolo_model is not None:234 try:235 with torch.no_grad():236 out = yolo_model(image_pil, imgsz=640, conf=0.5)237 results = out[0] if isinstance(out, list) else out238 239 # Draw boxes on the image240 annotated_img_np = results.plot() # RGB array241 annotated_img_pil = Image.fromarray(annotated_img_np)242 243 # Extract detections to DataFrame244 df = results.to_df()245 yolo_output_lines = []246 if not df.empty:247 for _, row in df.iterrows():248 yolo_output_lines.append(f"- {row['name']} (Conf: {row['confidence']:.2f})")249 yolo_result_text = "YOLOv8 Detections:\n" + "\n".join(yolo_output_lines)250 else:251 yolo_result_text = "YOLOv8: No defects detected."252 253 except Exception as e:254 print(f"Error during YOLOv8 inference or processing: {e}")255 yolo_result_text = f"Error during YOLOv8 processing: {e}"256 annotated_img_pil = image_pil.copy()257 258 else:259 # Either not PCB or model failed to load260 if class_label != 'pcb':261 print("Skipping YOLOv8 - Image classified as non_pcb.")262 yolo_result_text = "YOLOv8: Skipped (Input classified as non-PCB)."263 else:264 print("YOLOv8 model not loaded, skipping defect detection.")265 yolo_result_text = "YOLOv8: Model not loaded, cannot perform defect detection."266 # Keep original image267 annotated_img_pil = image_pil.copy()268 269 # Combine results for display270 final_text = f"{classification_result_text}\n{yolo_result_text}"271 return annotated_img_pil, final_text272 # --- End Processing Steps ---273 274# --- Gradio App Setup ---275print("Setting up Gradio interface...")276iface = gr.Interface(277 fn=detect_defects,278 inputs=gr.Image(type="pil", label="Upload PCB Image"),279 outputs=[280 gr.Image(type="pil", label="Processed Image"),281 gr.Textbox(label="Classification & YOLO Results", lines=5) # Increased lines282 ],283 title="PCB Defect Detection (Classifier + YOLOv8)",284 description=(285 "Upload an image. It's first classified (PCB/Non-PCB) using a ResNet model.\n"286 "If classified as PCB, a YOLOv8 model attempts to detect defects.\n"287 f"Requires '{os.path.basename(classifier_path)}' and '{os.path.basename(yolo_weights_path)}' to be present.\n"288 f"Using device: {device.type}"289 ),290 # # Add paths to actual example images if available291 # examples=[292 # ["sample_pcb_defective.jpg"], # Assumes this file exists293 # ["sample_pcb_ok.jpg"], # Assumes this file exists294 # ["sample_non_pcb.jpg"] # Assumes this file exists295 # ],296 allow_flagging="never"297)298 299 300print("Launching Gradio app...")301iface.launch(debug=True, share=True)