File size: 2,094 Bytes
f45111e 6341f26 5e5a776 f45111e 5e5a776 0f43111 5e5a776 0f43111 5e5a776 6a15848 5e5a776 82b9d60 5e5a776 1d11911 5e5a776 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 |
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()
|