import os import torch import timm from torchvision import transforms from PIL import Image from pathlib import Path from ultralytics import YOLO os.environ["HF_HUB_OFFLINE"] = "1" def get_qc_status(predicted_class): if predicted_class == "no_defect": return "Accepted" else: return "Rejected" def load_efficientnet(model_path): checkpoint = torch.load(model_path, map_location="cpu") config = checkpoint["config"] class_to_idx = checkpoint["class_to_idx"] idx_to_class = {v: k for k, v in class_to_idx.items()} model = timm.create_model(config["backbone"], pretrained=False, num_classes=config["num_classes"]) model.load_state_dict(checkpoint["state_dict"]) model.eval() 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 from functools import lru_cache @lru_cache(maxsize=2) def get_efficientnet(product_type="fmcg_inspection"): project_root = Path(__file__).parent.parent eff_path = project_root / f"models/poc2/{product_type}/efficientnet_b0/best.pt" if not eff_path.exists(): return None, None, None return load_efficientnet(eff_path) @lru_cache(maxsize=2) def get_yolo(product_type="fmcg_inspection"): project_root = Path(__file__).parent.parent yolo_path = project_root / f"models/poc2/{product_type}/yolov11_cls/best.pt" if not yolo_path.exists(): return None return YOLO(str(yolo_path)) 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, product_type="fmcg_inspection", 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(product_type) yolo_model = get_yolo(product_type) 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 detected **{pred_class}** at only {conf*100:.1f}% confidence, which is 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": "non-defective", "ok front": "non-defective", "defective front": "defective", "defective_front": "defective", "no_defect": "non-defective", "def_front": "defective" } if final_class in class_mapping: final_class = class_mapping[final_class] # Map to EXACT status UI expects if final_class == "non-defective": 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")