inference/classify.py

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