Ross - Final submission

app.py

@@ -1,136 +1,60 @@
-
 import gradio as gr
-from transformers import pipeline
-from PIL import Image
-import numpy as np
-
-# Load Models (Explicit)
-classifier = pipeline(
-    "image-classification",
-    model="google/vit-base-patch16-224",
-    use_fast=True
-)
-
-captioner = pipeline(
-    "image-text-to-text",
-    model="Salesforce/blip-image-captioning-base"
-)
-
-sentiment = pipeline(
-    "sentiment-analysis",
-    model="distilbert-base-uncased-finetuned-sst-2-english"
+import torch
+from transformers import (
+    BlipProcessor,
+    BlipForConditionalGeneration,
+    AutoTokenizer,
+    AutoModelForSeq2SeqLM
 )
 
-# Smart Matching Function
-def check_match(label, caption):
-    label = label.lower()
-    caption = caption.lower()
-
-    keyword_map = {
-        "dog": ["dog", "puppy"],
-        "cat": ["cat", "kitten"],
-        "car": ["car", "vehicle"],
-        "pizza": ["pizza"],
-        "bird": ["bird"],
-        "person": ["person", "man", "woman", "people"],
-        "bicycle": ["bike", "bicycle"],
-        "laptop": ["laptop", "computer"],
-        "beach": ["beach", "ocean", "sea"],
-        "mountain": ["mountain"]
-    }
+# Image captioning model
+caption_processor = BlipProcessor.from_pretrained("Salesforce/blip-image-captioning-base")
+caption_model = BlipForConditionalGeneration.from_pretrained("Salesforce/blip-image-captioning-base")
 
-    for key, words in keyword_map.items():
-        if key in label:
-            for word in words:
-                if word in caption:
-                    return True
+# English-to-Spanish translation model
+translation_tokenizer = AutoTokenizer.from_pretrained("Helsinki-NLP/opus-mt-en-es")
+translation_model = AutoModelForSeq2SeqLM.from_pretrained("Helsinki-NLP/opus-mt-en-es")
 
-    return False
+def analyze_image(image):
+    if image is None:
+        return "Please upload an image first.", ""
 
-# Image Processing Function
-def process_image(image):
     try:
-
-        # Image is already PIL — no conversion needed
-
-        # 1. Classification
-        class_result = classifier(image)
-        top_class = class_result[0]['label']
-        confidence = round(class_result[0]['score'], 3)
-
-        # 2. Captioning (FIXED)
-        print(f"Debug: Type of image for captioner: {type(image)}")
-        print(f"Debug: Image info (if PIL): {image.format, image.size, image.mode}")
-        # Pass image as a list and provide a default text prompt
-        caption_result = captioner(images=[image], text="a photo of")
-
-        caption = caption_result[0]['generated_text']
-
-        # Handle empty caption
-        if not caption or len(caption.strip()) == 0:
-            sentiment_label = "N/A"
-            sentiment_score = 0.0
-            comparison_text = "No Caption ❌"
-            caption = "No caption generated."
-        else:
-            # 3. Sentiment
-            sentiment_result = sentiment(caption)[0]
-            sentiment_label = sentiment_result['label']
-            sentiment_score = round(sentiment_result['score'], 3)
-
-            # 4. Comparison
-            match = top_class.lower() in caption.lower()
-            comparison_text = "Match ✅" if match else "No Match ❌"
-
-        return (
-            f"{top_class} (Confidence: {confidence})",
-            caption,
-            f"{sentiment_label} (Score: {sentiment_score})",
-            comparison_text
+        # Generate English caption
+        inputs = caption_processor(image, return_tensors="pt")
+        caption_ids = caption_model.generate(**inputs, max_new_tokens=40)
+        english_caption = caption_processor.decode(caption_ids[0], skip_special_tokens=True)
+
+        # Translate caption to Spanish
+        translation_inputs = translation_tokenizer(
+            english_caption,
+            return_tensors="pt",
+            padding=True
         )
-
-    except Exception as e:
-        return (
-            "Error in classification",
-            "Error in captioning",
-            "Error in sentiment",
-            str(e)
+        translation_ids = translation_model.generate(**translation_inputs, max_new_tokens=60)
+        spanish_caption = translation_tokenizer.decode(
+            translation_ids[0],
+            skip_special_tokens=True
         )
 
-# Create a dummy image for testing the pipelines
-dummy_image_path = "dummy_image.png"
-img_data = np.zeros((100, 100, 3), dtype=np.uint8) # 100x100 black image
-dummy_img = Image.fromarray(img_data)
-dummy_img.save(dummy_image_path)
-
-# Load the dummy image for testing
-img = Image.open(dummy_image_path)
+        return english_caption, spanish_caption
 
-classifier_test_result = classifier(img)
-print("Classifier test result:", classifier_test_result)
-
-# Pass image as a list and provide a default text prompt for testing
-captioner_test_result = captioner(images=[img], text="a photo of")
-print("Captioner test result:", captioner_test_result)
+    except Exception as e:
+        return f"Error: {str(e)}", ""
 
-# Gradio Interface
-interface = gr.Interface(
-    fn=process_image,
-    inputs=gr.Image(type="pil"),
+demo = gr.Interface(
+    fn=analyze_image,
+    inputs=gr.Image(type="pil", label="Upload an Image"),
     outputs=[
-        gr.Text(label="Image Classification"),
-        gr.Text(label="Generated Caption"),
-        gr.Text(label="Caption Sentiment"),
-        gr.Text(label="Model Agreement")
+        gr.Textbox(label="English Image Caption"),
+        gr.Textbox(label="Spanish Translation")
     ],
-    title="Multimodal Image Analysis System",
+    title="Multimodal Image Captioning and Spanish Translation Tool",
     description=(
-        "Upload an image to analyze it using multiple AI models. "
-        "This app compares image classification and captioning results "
-        "and applies NLP to evaluate the caption."
+        "This app uses two AI models. First, it uses a vision-language model to describe "
+        "an uploaded image in English. Then, it uses a language translation model to translate "
+        "that caption into Spanish."
     )
 )
 
-# Launch App
-if __name__ == "__main__":
-    interface.launch(debug=True)
+demo.launch()