import os import time import torch import timm from torchvision import transforms from PIL import Image from pathlib import Path from ultralytics import YOLO from inference.logger import get_logger os.environ["HF_HUB_OFFLINE"] = "1" log = get_logger("classify") def get_qc_status(predicted_class): if predicted_class == "no_defect": return "Accepted" else: return "Rejected" def load_efficientnet(model_path): log.info("EfficientNet — loading checkpoint from %s", model_path) t0 = time.perf_counter() checkpoint = torch.load(model_path, map_location="cpu") log.debug("EfficientNet — checkpoint read (%.2fs)", time.perf_counter() - t0) config = checkpoint["config"] class_to_idx = checkpoint["class_to_idx"] idx_to_class = {v: k for k, v in class_to_idx.items()} t1 = time.perf_counter() model = timm.create_model(config["backbone"], pretrained=False, num_classes=config["num_classes"]) log.debug("EfficientNet — timm model created (%.2fs)", time.perf_counter() - t1) model.load_state_dict(checkpoint["state_dict"]) model.eval() log.info("EfficientNet — ready total=%.2fs classes=%s", time.perf_counter() - t0, list(idx_to_class.values())) return model, config, idx_to_class def predict_efficientnet(image_path, model, config, idx_to_class): # Convert to grayscale ("L") to strip color, then back to 3-channel ("RGB") img = Image.open(image_path).convert("L").convert("RGB") NORM = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) tf = transforms.Compose([ transforms.Resize((config["img_size"], config["img_size"])), transforms.ToTensor(), NORM ]) x = tf(img).unsqueeze(0) with torch.no_grad(): logits = model(x) probs = torch.softmax(logits, dim=1)[0] pred_idx = logits.argmax(1).item() predicted_class = idx_to_class[pred_idx] confidence = probs[pred_idx].item() return predicted_class, confidence def predict_yolo(image_path, yolo_model): # YOLO returns a list of Results objects. We take the first one since it's one image. results = yolo_model(image_path, verbose=False) result = results[0] # Get top 1 prediction pred_idx = result.probs.top1 confidence = result.probs.top1conf.item() predicted_class = result.names[pred_idx] return predicted_class, confidence import streamlit as st @st.cache_resource def get_efficientnet(): project_root = Path(__file__).parent.parent eff_path = project_root / "models/poc2/efficientnet_b0/best.pt" log.info("get_efficientnet — path=%s exists=%s", eff_path, eff_path.exists()) if not eff_path.exists(): log.error("get_efficientnet — model file not found at %s", eff_path) return None, None, None return load_efficientnet(eff_path) @st.cache_resource def get_yolo(): project_root = Path(__file__).parent.parent yolo_path = project_root / "models/poc2/yolov11_cls/best.pt" log.info("get_yolo (cls) — path=%s exists=%s", yolo_path, yolo_path.exists()) if not yolo_path.exists(): log.error("get_yolo (cls) — model file not found at %s", yolo_path) return None t0 = time.perf_counter() model = YOLO(str(yolo_path)) log.info("get_yolo (cls) — ready total=%.2fs", time.perf_counter() - t0) return model def predict_efficientnet_pil(img, model, config, idx_to_class, min_no_defect_prob=0.0): # Convert to grayscale ("L") to strip color, then back to 3-channel ("RGB") img = img.convert("L").convert("RGB") NORM = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]) tf = transforms.Compose([ transforms.Resize((config["img_size"], config["img_size"])), transforms.ToTensor(), NORM ]) x = tf(img).unsqueeze(0) original_pred = None original_conf = None with torch.no_grad(): logits = model(x) probs = torch.softmax(logits, dim=1)[0] pred_idx = probs.argmax().item() # Apply custom decision threshold logic if requested if min_no_defect_prob > 0.0: no_defect_idx = next((k for k, v in idx_to_class.items() if v == "no_defect"), None) if no_defect_idx is not None and pred_idx == no_defect_idx and probs[pred_idx] < min_no_defect_prob: original_pred = "no_defect" original_conf = probs[pred_idx].item() probs_copy = probs.clone() probs_copy[no_defect_idx] = -1.0 pred_idx = probs_copy.argmax().item() predicted_class = idx_to_class[pred_idx] confidence = probs[pred_idx].item() return predicted_class, confidence, original_pred, original_conf def classify_qc(img, force_model=None, threshold=0.80, min_no_defect_prob=0.0): """ Main entry point for Streamlit. Takes a PIL Image, runs the dual-inference pipeline, and returns the formatted dictionary. """ eff_model, eff_config, eff_idx_to_class = get_efficientnet() yolo_model = get_yolo() orig_pred = None orig_conf = None eff_second_guess = None eff_second_conf = None if force_model == "YOLOv11": if yolo_model is None: return {"status": "Model Not Found", "confidence": 0.0} results = yolo_model(img, verbose=False) result = results[0] pred_idx = result.probs.top1 final_conf = result.probs.top1conf.item() final_class = result.names[pred_idx] source = "YOLOv11 (Forced)" warning_msg = None else: if eff_model is None: return {"status": "Model Not Found", "confidence": 0.0} pred_class, conf, orig_pred, orig_conf = predict_efficientnet_pil(img, eff_model, eff_config, eff_idx_to_class, min_no_defect_prob=min_no_defect_prob) final_class = pred_class final_conf = conf if orig_pred is not None: eff_second_guess = pred_class eff_second_conf = conf source = "EfficientNet (Forced)" if force_model == "EfficientNet" else "EfficientNet (Primary)" warning_msg = None # 2. Fallback check using the dynamic threshold if force_model is None and conf < threshold and yolo_model is not None: warning_msg = f"EfficientNet confidence ({conf*100:.1f}%) was below the {threshold*100:.1f}% threshold. Shifted to YOLOv11 fallback." results = yolo_model(img, verbose=False) result = results[0] pred_idx = result.probs.top1 final_conf = result.probs.top1conf.item() final_class = result.names[pred_idx] source = "YOLOv11 (Fallback)" # Normalize class names from different models/datasets class_mapping = { "ok_front": "no_defect", "ok front": "no_defect", "defective front": "def_front", "defective_front": "def_front", } if final_class in class_mapping: final_class = class_mapping[final_class] # Map to EXACT status UI expects if final_class == "no_defect": status = "QC Approved" else: status = "QC Rejected" # Return dictionary formatted for the UI return { "status": status, "defect": final_class, "confidence": round(final_conf * 100, 2), "source": source, "warning": warning_msg, "original_pred": orig_pred, "original_conf": round(orig_conf * 100, 2) if orig_conf is not None else None, "eff_second_guess": eff_second_guess, "eff_second_conf": round(eff_second_conf * 100, 2) if eff_second_conf is not None else None } if __name__ == "__main__": import argparse parser = argparse.ArgumentParser(description="Run Dual-Inference with YOLO Fallback") parser.add_argument("--image", type=str, help="Path to the image you want to test") parser.add_argument("--eff-model", type=str, default="models/poc2/efficientnet_b0/best.pt", help="Path to primary EfficientNet (.pt)") parser.add_argument("--yolo-model", type=str, default="models/poc2/yolov11_cls/weights/best.pt", help="Path to YOLO Fallback (.pt)") parser.add_argument("--threshold", type=float, default=0.80, help="Confidence threshold below which YOLO will be triggered (default: 0.80)") args = parser.parse_args() if args.image: image_path = Path(args.image) if not image_path.exists(): print(f"Error: Image not found at {image_path}") else: img = Image.open(image_path) res = classify_qc(img, threshold=args.threshold) print("\n=== FINAL RESULTS ===") if res.get("warning"): print(f"WARNING: {res['warning']}") print(f"Source: {res['source']}") print(f"Predicted Defect: {res['defect']}") print(f"QC Status: {res['status']}") print(f"Confidence: {res['confidence']}%") else: print("\nNo --image argument provided. To test an image, run:") print("python inference\\classify.py --image path\\to\\image.jpg")