app.py

import gradio as gr
import torch
from torchvision import transforms
from PIL import Image
import joblib
import base64
from reportlab.pdfgen import canvas
from reportlab.lib.pagesizes import letter
import os
import tempfile

# Load model and metadata
try:
    model = torch.jit.load("resnet_grocery_model_scripted.pt", map_location="cpu")
    model.eval()
except Exception as e:
    print(f"Error loading model: {str(e)}")
    raise

try:
    assets = joblib.load("deployment_assets.joblib")
    transform = assets['transform']
    class_names = assets['class_names']
    print(f"Loaded {len(class_names)} classes")
except Exception as e:
    print(f"Error loading assets: {str(e)}")
    raise

# Price list
items = [
    'Bisconni Chocolate Chip Cookies 46.8gm', 'Coca Cola Can 250ml', 'Colgate Maximum Cavity Protection 75gm',
    'Fanta 500ml', 'Fresher Guava Nectar 500ml', 'Fruita Vitals Red Grapes 200ml', 'Islamabad Tea 238gm',
    'Kolson Slanty Jalapeno 18gm', 'Kurkure Chutney Chaska 62gm', 'LU Candi Biscuit 60gm', 'LU Oreo Biscuit 19gm',
    'LU Prince Biscuit 55.2gm', 'Lays Masala 34gm', 'Lays Wavy Mexican Chili 34gm', 'Lifebuoy Total Protect Soap 96gm',
    'Lipton Yellow Label Tea 95gm', 'Meezan Ultra Rich Tea 190gm', 'Peek Freans Sooper Biscuit 13.2gm',
    'Safeguard Bar Soap Pure White 175gm', 'Shezan Apple 250ml', 'Sunsilk Shampoo Soft - Smooth 160ml',
    'Super Crisp BBQ 30gm', 'Supreme Tea 95gm', 'Tapal Danedar 95gm', 'Vaseline Healthy White Lotion 100ml'
]

prices = [
    55.20, 31.75, 90.00, 63.50, 50.00, 35.00, 150.00, 15.00, 25.00, 30.00, 10.00, 30.00, 20.00, 20.00, 44.50,
    100.00, 200.00, 10.00, 70.00, 25.00, 120.00, 15.00, 100.00, 100.00, 120.00
]

pricelist = dict(zip(items, prices))

def classify_and_track(files, budget):
    results = []
    total_cost = 0
    receipt_lines = []

    receipt_lines.append("========================================")
    receipt_lines.append("           🧾 OFFICIAL RECEIPT")
    receipt_lines.append("         🛒 SOBRANG TINDANG INA!")
    receipt_lines.append("        📍 Online Smart Checkout")
    receipt_lines.append("========================================")
    receipt_lines.append(" ITEM                          PRICE")
    receipt_lines.append("----------------------------------------")

    for f in files:
        try:
            print(f"Processing file: {f.name}, Exists: {os.path.exists(f.name)}")
            img = Image.open(f.name).convert("RGB")
        except Exception as e:
            print(f"Image loading error for {f.name}: {str(e)}")
            results.append((None, f"Invalid image: {str(e)}"))
            continue

        try:
            input_tensor = transform(img).unsqueeze(0)
            print(f"Input tensor shape: {input_tensor.shape}")
            with torch.no_grad():
                outputs = model(input_tensor)
                print(f"Model output shape: {outputs.shape}")
                predicted_index = torch.argmax(outputs, dim=1).item()
                if predicted_index >= len(class_names):
                    raise ValueError(f"Predicted index {predicted_index} out of range for class_names (len={len(class_names)})")
                predicted_class = class_names[predicted_index]
                price = pricelist.get(predicted_class, 0)
                total_cost += price
                results.append((img, f"{predicted_class} - ₱{price:.2f}"))

                name = predicted_class
                if len(name) > 26:
                    name = name[:23] + "..."
                receipt_lines.append(f" {name:<28} ₱{price:>6.2f}")
        except Exception as e:
            print(f"Inference error for {f.name}: {str(e)}")
            results.append((None, f"Inference error: {str(e)}"))
            continue

    remaining = budget - total_cost
    receipt_lines.append("----------------------------------------")
    receipt_lines.append(f" TOTAL{'':<24}₱{total_cost:>7.2f}")
    receipt_lines.append(f" BUDGET{'':<23}₱{budget:>7.2f}")
    receipt_lines.append(f" REMAINING{'':<20}₱{remaining:>7.2f}")
    receipt_lines.append("========================================")

    percent = min(100, int((total_cost / budget) * 100)) if budget else 0
    status = f"**Total: ₱{total_cost:.2f} / ₱{budget:.2f}**"

    if total_cost <= budget:
        receipt_lines.append(" ✅ You're within budget! 🎉")
        status_msg = f"✅ **Within Budget!** 🎉\n{status}"
        sound_path = "success.mp3"
    else:
        receipt_lines.append(" ❌ Over budget! Remove some snacks!")
        status_msg = f"🚨 **Over Budget!** 💸\n{status}"
        sound_path = "fail.mp3"

    receipt_lines.append("========================================")
    receipt_lines.append(" THANK YOU FOR SHOPPING WITH US!")
    receipt_lines.append("   Come back for smarter buys! 🧾")
    receipt_lines.append("========================================")

    # Convert receipt to markdown
    receipt_text = "```\n" + "\n".join(receipt_lines) + "\n```"

    # Generate PDF
    with tempfile.NamedTemporaryFile(delete=False, suffix=".pdf") as tmp_pdf:
        c = canvas.Canvas(tmp_pdf.name, pagesize=letter)
        width, height = letter
        y = height - 40
        for line in receipt_lines:
            c.drawString(40, y, line)
            y -= 15
        c.save()
        pdf_path = tmp_pdf.name

    return results, status_msg, receipt_text, percent, sound_path, gr.update(value=pdf_path, visible=True), gr.update(value=receipt_text, visible=True)

def reset_ui():
    return [], "", "", 0, None, gr.update(visible=False), gr.update(visible=False)

# Load and encode local background image
try:
    with open("background.jpg", "rb") as img_file:
        b64_string = base64.b64encode(img_file.read()).decode()
except Exception as e:
    print(f"Error loading background image: {str(e)}")
    b64_string = ""

html_style = f"""
<link href="https://fonts.googleapis.com/css2?family=Poppins:wght@400;600;800&display=swap" rel="stylesheet">
<style>
body {{
    background-image: url("data:image/jpg;base64,{b64_string}");
    background-size: cover;
    background-position: center;
    background-repeat: no-repeat;
    font-family: 'Poppins', sans-serif;
    color: #222;
}}

.glass-card {{
    background: rgba(255, 255, 255, 0.75);
    border-radius: 20px;
    box-shadow: 0 8px 32px rgba(31, 38, 135, 0.37);
    backdrop-filter: blur(10px);
    -webkit-backdrop-filter: blur(10px);
    border: 1.5px solid rgba(255, 255, 255, 0.2);
    padding: 2rem;
    margin-bottom: 1.5rem;
}}

h1 {{
    text-align: center;
    font-size: 2.6rem;
    font-weight: 800;
    color: #000 !important;
    margin-bottom: 0.5rem;
    font-family: 'Poppins', sans-serif;
}}

h2 {{
    text-align: center;
    font-size: 0.1rem;
    font-weight: 0.5;
    color: #000 !important;
    font-family: 'Poppins', sans-serif;
}}

.gr-button {{
    font-family: 'Poppins', sans-serif;
    font-weight: 600;
}}

.gr-number input, .gr-file input {{
    font-family: 'Poppins', sans-serif;
}}
</style>
"""

with gr.Blocks(title="Grocery Classifier") as app:
    gr.HTML(html_style)

    with gr.Column(elem_classes="glass-card"):
        gr.Markdown("<h1>🧾 Budget Check! Know Before You Chow 🥗</h1>")
        gr.Markdown("<h2>Snap a snack, upload the pack, and find out if you’re on track — or flat broke. Shopping has never been so dramatic.</h2>")

    with gr.Column(elem_classes="glass-card"):
        with gr.Row():
            file_input = gr.File(file_types=["image"], file_count="multiple", label="📷 Upload Grocery Items")
            budget_input = gr.Number(label="💰 Budget (₱)", value=500.0, precision=2)

        with gr.Row():
            classify_btn = gr.Button("🚀 Classify Items")
            reset_btn = gr.Button("🔄 Reset")

    with gr.Column(elem_classes="glass-card"):
        gallery = gr.Gallery(label="📸 Classified Items", columns=3)
        status_display = gr.Markdown()
        progress_bar = gr.Slider(label="Budget Used (%)", minimum=0, maximum=100, interactive=False)
        receipt_display = gr.Markdown(label="🧾 Grocery Receipt", visible=False)
        download_receipt = gr.File(label="📥 Download Receipt PDF", interactive=False, visible=False)
        audio_output = gr.Audio(interactive=False, autoplay=True, visible=False)

    classify_btn.click(
        fn=classify_and_track,
        inputs=[file_input, budget_input],
        outputs=[gallery, status_display, receipt_display, progress_bar, audio_output, download_receipt, receipt_display]
    )

    reset_btn.click(
        fn=reset_ui,
        outputs=[gallery, status_display, receipt_display, progress_bar, audio_output, download_receipt, receipt_display]
    )

app.launch()