Zhao - Final submission

app.py

@@ -1,98 +1,96 @@
-import torch
-import gradio as gr
-from PIL import Image
-from transformers import (
-    BlipProcessor,
-    BlipForConditionalGeneration,
-    pipeline
-)
-
-# Select device
-device = "cuda" if torch.cuda.is_available() else "cpu"
-torch_dtype = torch.float16 if torch.cuda.is_available() else torch.float32
-
-# Load BLIP captioning model directly
-processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
-caption_model = BlipForConditionalGeneration.from_pretrained(
-    "Salesforce/blip-image-captioning-base",
-    torch_dtype=torch_dtype
-).to(device)
-
-# Load image classification model
-classifier = pipeline(
-    task="image-classification",
-    model="google/vit-base-patch16-224",
-    device=0 if torch.cuda.is_available() else -1
-)
+import os
+import random
 
-print("Models loaded successfully.")
-
-def generate_caption(image):
-    inputs = processor(images=image, return_tensors="pt").to(device, torch_dtype)
-    output = caption_model.generate(**inputs, max_new_tokens=30)
-    caption = processor.decode(output[0], skip_special_tokens=True)
-    return caption
-
-def build_summary(caption: str, label: str) -> str:
-    caption = caption.strip() if caption else "No caption available"
-    label = label.strip() if label else "unknown object"
-
-    return (
-        f"The captioning model describes the image as: {caption}. "
-        f"The image classification model identifies the main subject as: {label}. "
-        f"Taken together, the image appears to focus on this subject or scene."
+import gradio as gr
+import torch
+from datasets import load_dataset
+from transformers import BlipForQuestionAnswering, BlipProcessor
+
+# Always use BlipProcessor + BlipForQuestionAnswering directly.
+# The Transformers pipeline API for BLIP-VQA changed between v4 and v5 and can echo
+# the question instead of returning the answer. Direct model calls are stable across versions.
+MODEL_ID = "Salesforce/blip-vqa-base"
+
+_device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+_processor = BlipProcessor.from_pretrained(MODEL_ID)
+_model = BlipForQuestionAnswering.from_pretrained(MODEL_ID).to(_device)
+_model.eval()
+
+
+def run_vqa(image, question: str) -> str:
+    inputs = _processor(images=image, text=question, return_tensors="pt").to(_device)
+    with torch.no_grad():
+        out = _model.generate(**inputs, max_new_tokens=32)
+    return _processor.decode(out[0], skip_special_tokens=True).strip()
+
+FOOD101_ID = "ethz/food101"
+DEFAULT_VQA_QUESTION = "What kind of food is in this image?"
+
+
+def answer_question(image, question):
+    if not question:
+        return "Please type a question about the image."
+    if image is None:
+        return "Please upload an image (or use Load random sample) before clicking Answer."
+    return run_vqa(image, question)
+
+
+def load_food101_sample(cache: dict | None):
+    """
+    Pull one random example from Hugging Face Datasets (Food-101).
+    Dataset card: https://huggingface.co/datasets/ethz/food101
+    """
+    if cache is None:
+        cache = {"ds": load_dataset(FOOD101_ID, split="validation")}
+    ds = cache["ds"]
+    i = random.randint(0, len(ds) - 1)
+    ex = ds[i]
+    names = ds.features["label"].names
+    gold = names[ex["label"]]
+    note = (
+        f"Loaded from Hugging Face Datasets `{FOOD101_ID}` "
+        f"(split `validation`, index {i}). Reference label: `{gold}`."
     )
+    return ex["image"], DEFAULT_VQA_QUESTION, cache, note
 
-def analyze_image(image):
-    try:
-        if image is None:
-            return (
-                "Please upload an image.",
-                "No classification available.",
-                "Please upload an image first."
-            )
 
-        image = image.convert("RGB")
-
-        # Captioning
-        caption = generate_caption(image)
-        print("CAPTION RESULT:", caption)
-
-        # Classification
-        class_result = classifier(image)
-        print("CLASSIFICATION RESULT:", class_result)
-
-        if isinstance(class_result, list) and len(class_result) > 0:
-            top_label = class_result[0].get("label", "Unknown")
-            top_score = class_result[0].get("score", 0.0)
-            classification_text = f"{top_label} (confidence: {top_score:.4f})"
-        else:
-            top_label = "Unknown"
-            classification_text = "No classification generated."
-
-        summary = build_summary(caption, top_label)
-
-        return caption, classification_text, summary
-
-    except Exception as e:
-        print("ERROR:", str(e))
-        error_text = f"Error: {str(e)}"
-        return error_text, error_text, error_text
-
-demo = gr.Interface(
-    fn=analyze_image,
-    inputs=gr.Image(type="pil", label="Upload an Image"),
-    outputs=[
-        gr.Textbox(label="Generated Caption"),
-        gr.Textbox(label="Top Classification"),
-        gr.Textbox(label="Combined Summary", lines=4)
-    ],
-    title="Image Captioning, Classification, and Summary App",
-    description=(
-        "Upload an image to generate an automatic caption, predict the main image class, "
-        "and produce a short combined summary."
-    ),
-)
+with gr.Blocks() as demo:
+    ds_cache = gr.State(None)
+    gr.Markdown(
+        "## BLIP Visual Question Answering\n"
+        "Gradio app for [Hugging Face Spaces](https://huggingface.co/docs/hub/spaces). "
+        "Upload your own image **or** load a sample image from "
+        f"[Hugging Face Datasets — `{FOOD101_ID}`](https://huggingface.co/datasets/ethz/food101)."
+    )
+    dataset_note = gr.Markdown("")
+    with gr.Row():
+        image_in = gr.Image(type="pil", label="Image")
+        question_in = gr.Textbox(
+            label="Question",
+            placeholder="e.g. What is the weather in this image?",
+        )
+    answer_out = gr.Textbox(label="Answer")
+    with gr.Row():
+        submit_btn = gr.Button("Answer", variant="primary")
+        sample_btn = gr.Button(f"Load random sample ({FOOD101_ID})")
+
+    submit_btn.click(answer_question, [image_in, question_in], [answer_out])
+    sample_btn.click(
+        load_food101_sample,
+        [ds_cache],
+        [image_in, question_in, ds_cache, dataset_note],
+    )
 
 if __name__ == "__main__":
-    demo.launch()
+    # Hugging Face Spaces sets PORT; bind 0.0.0.0 so the Space proxy can reach the app.
+    # Locally: open http://127.0.0.1:<PORT> (default PORT 7860 if unset).
+    _port = int(os.environ.get("PORT", "7860"))
+    # Queue mode is recommended for Spaces and tends to play nicer with Gradio’s async stack.
+    demo.queue()
+    demo.launch(
+        server_name="0.0.0.0",
+        server_port=_port,
+        # Space URL is already public; share=True would add a separate *.gradio.live tunnel (not needed).
+        share=False,
+        show_error=True,
+    )