import gradio as gr
from PIL import Image, ImageChops
from transformers import BlipProcessor, BlipForQuestionAnswering
import torch

# Load BLIP-VQA model
model_name = "Salesforce/blip-vqa-base"
processor = BlipProcessor.from_pretrained(model_name)
model = BlipForQuestionAnswering.from_pretrained(model_name)

valid_classes = ["plastic", "metal", "paper", "cardboard", "glass", "trash"]
base_img = None  # Global variable to store base image

# Function to compute difference image
def get_difference_image(base: Image.Image, trash: Image.Image) -> Image.Image:
    diff = ImageChops.difference(base, trash).convert("RGB")
    # Optional: enhance contrast to highlight difference
    return diff

# Set base image
def set_base(image):
    global base_img
    base_img = image.convert("RGB")
    return "Base image saved successfully."

# Detect trash material
def detect_material(trash_image):
    global base_img
    if base_img is None:
        return "Please set base image first."

    trash_image = trash_image.convert("RGB")
    diff_image = get_difference_image(base_img, trash_image)

    question = "What material is this object? Choose one of: plastic, metal, paper, cardboard, glass, trash."

    inputs = processor(diff_image, question, return_tensors="pt")
    out = model.generate(**inputs)
    answer = processor.decode(out[0], skip_special_tokens=True).lower()

    # Ensure answer is one of the valid classes
    material = next((c for c in valid_classes if c in answer), "trash")
    return material.capitalize()

# Build Gradio UI
set_base_ui = gr.Interface(
    fn=set_base,
    inputs=gr.Image(type="pil", label="Upload Base Image (Empty Bin)"),
    outputs=gr.Textbox(label="Result"),
    title="Set Base Image",
    api_name="/set_base"
)

detect_ui = gr.Interface(
    fn=detect_material,
    inputs=gr.Image(type="pil", label="Upload Trash Image"),
    outputs=gr.Textbox(label="Detected Material"),
    title="Trash Material Detector",
    api_name="/detect_material"
)

demo = gr.TabbedInterface([set_base_ui, detect_ui], ["Set Base", "Detect Trash"])
demo.launch()