Lazypanda0103/Unified-Comprehensive-Freshness-Classification

1import gradio as gr2import torch3import timm4import cv25import numpy as np6from ultralytics import YOLO7import torch.nn.functional as F8import os9 10# ----------------------------11# Load Classifier at startup12# ----------------------------13device = "cpu"14 15classifier = timm.create_model("efficientnet_b0", pretrained=False)16classifier.classifier = torch.nn.Linear(classifier.classifier.in_features, 3)17classifier.load_state_dict(torch.load("final_model.pth", map_location=device))18classifier.eval()19 20labels = ["Fresh", "Semi Fresh", "Rotten"]21 22# ----------------------------23# Load YOLO lazily (avoids startup download error)24# ----------------------------25detector = None26 27def get_detector():28    global detector29    if detector is None:30        detector = YOLO("yolov8n.pt")31    return detector32 33# ----------------------------34# Image preprocessing35# ----------------------------36def preprocess(img):37    img = cv2.resize(img, (256, 256))38    img = img / 255.039    img = np.transpose(img, (2, 0, 1))40    tensor = torch.tensor(img, dtype=torch.float32).unsqueeze(0)41    return tensor42 43# ----------------------------44# Prediction pipeline45# ----------------------------46def predict(image):47    if image is None:48        return None, "Please upload an image.", 049 50    img = np.array(image)51    52    try:53        det = get_detector()54        results = det(img)55        if len(results[0].boxes) > 0:56            box = results[0].boxes.xyxy[0].cpu().numpy()57            x1, y1, x2, y2 = map(int, box)58            crop = img[y1:y2, x1:x2]59            cv2.rectangle(img, (x1, y1), (x2, y2), (0, 255, 0), 2)60        else:61            crop = img62    except Exception:63        crop = img64 65    tensor = preprocess(crop)66 67    with torch.no_grad():68        output = classifier(tensor)69        probs = F.softmax(output, dim=1)70        score = torch.max(probs).item() * 10071        label = labels[torch.argmax(probs).item()]72 73    result = f"""### Prediction: **{label}**\nFreshness Score: **{score:.2f}/100**"""74    return img, result, score75 76# ----------------------------77# UI78# ----------------------------79with gr.Blocks(theme=gr.themes.Soft()) as demo:80    gr.Markdown("# 🍎 Food Freshness Detection System")81    gr.Markdown("Upload an image or use your camera to check food freshness.")82 83    with gr.Row():84        image_input = gr.Image(85            sources=["upload", "webcam"],86            type="numpy",87            label="Upload or Capture Image"88        )89        output_image = gr.Image(label="Detection Result")90 91    prediction_text = gr.Markdown()92    confidence = gr.Slider(93        minimum=0,94        maximum=100,95        label="Freshness Score",96        interactive=False97    )98 99    analyze_btn = gr.Button("Analyze Freshness")100    analyze_btn.click(101        fn=predict,102        inputs=image_input,103        outputs=[output_image, prediction_text, confidence]104    )105 106demo.launch()