Update app.py

app.py

@@ -1,25 +1,30 @@
 import gradio as gr
-from transformers import pipeline
+from transformers import VisionEncoderDecoderModel, ViTImageProcessor, AutoTokenizer
+from PIL import Image
 
-# Load a simple image classification model
-classifier = pipeline("image-classification", model="google/vit-large-patch16-224-in21k")
+# Load the stronger model
+model_id = "nlpconnect/vit-gpt2-image-captioning"
+model = VisionEncoderDecoderModel.from_pretrained(model_id)
+feature_extractor = ViTImageProcessor.from_pretrained(model_id)
+tokenizer = AutoTokenizer.from_pretrained(model_id)
 
-def classify_image(image):
-    # Run prediction
-    results = classifier(image)
-    # Get top label and confidence
-    top_result = results[0]
-    label = top_result["label"]
-    score = round(top_result["score"] * 100, 2)
-    return f"{label} ({score}%)"
+def classify_better(image):
+    # preprocess
+    if image.mode != "RGB":
+        image = image.convert("RGB")
+    pixel_values = feature_extractor(images=image, return_tensors="pt").pixel_values
+
+    # Generate caption
+    output_ids = model.generate(pixel_values, max_length=20, num_beams=5)
+    caption = tokenizer.decode(output_ids[0], skip_special_tokens=True).strip()
+    return caption
 
-# Simple Gradio UI
 demo = gr.Interface(
-    fn=classify_image,
+    fn=classify_better,
     inputs=gr.Image(type="pil"),
     outputs="text",
-    title="🖼️ Simple Image Classifier",
-    description="what is this food"
+    title="Better Image Captioning",
+    description="Upload an image and the model will try to describe it (better)."
 )
 
 if __name__ == "__main__":