Update app.py

app.py

@@ -1,52 +1,48 @@
-from fastapi import FastAPI
-from fastapi.responses import RedirectResponse
-import gradio as gr
+from fastapi import FastAPI, UploadFile, Form
+from fastapi.responses import RedirectResponse, FileResponse, JSONResponse
+import os
+import shutil
 from PIL import Image
-import numpy as np
-from transformers import pipeline
+from transformers import ViltProcessor, ViltForQuestionAnswering
 from gtts import gTTS
+import torch
 import tempfile
-import os
-import pytesseract  # ✅ Replacing easyocr
-
-
+import gradio as gr
 
 app = FastAPI()
 
-# Models
-caption_model = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
-vqa_model = pipeline("visual-question-answering", model="dandelin/vilt-b32-finetuned-vqa")
-
+# Load VQA Model
+vqa_processor = ViltProcessor.from_pretrained("dandelin/vilt-b32-finetuned-vqa")
+vqa_model = ViltForQuestionAnswering.from_pretrained("dandelin/vilt-b32-finetuned-vqa")
 
-def process_image_question(image: Image.Image, question: str):
-    if image is None:
-        return "No image uploaded.", None
 
-    try:
-        # Convert image to numpy
-        np_image = np.array(image)
+def answer_question_from_image(image, question):
+    if image is None or not question.strip():
+        return "Please upload an image and ask a question.", None
 
-        # OCR text using pytesseract
-        extracted_text = pytesseract.image_to_string(image)
+    # Process with model
+    inputs = vqa_processor(image, question, return_tensors="pt")
+    with torch.no_grad():
+        outputs = vqa_model(**inputs)
+    predicted_id = outputs.logits.argmax(-1).item()
+    answer = vqa_model.config.id2label[predicted_id]
 
-        # Caption
-        caption = caption_model(image)[0]['generated_text']
-
-        # Visual QA
-        vqa_result = vqa_model(image=image, question=question)
-        answer = vqa_result[0]['answer']
-
-        # Answer as speech
+    # Generate TTS audio
+    try:
         tts = gTTS(text=answer)
         with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as tmp:
             tts.save(tmp.name)
             audio_path = tmp.name
+    except Exception as e:
+        return f"Answer: {answer}\n\n⚠️ Audio generation error: {e}", None
 
-        final_output = f"🖼️ Caption: {caption}\n\n📖 OCR Text:\n{extracted_text}\n\n❓ Answer: {answer}"
-        return final_output, audio_path
+    return answer, audio_path
+
+
+def process_image_question(image: Image.Image, question: str):
+    answer, audio_path = answer_question_from_image(image, question)
+    return answer, audio_path
 
-    except Exception as e:
-        return f"❌ Error: {e}", None
 
 gui = gr.Interface(
     fn=process_image_question,
@@ -55,15 +51,15 @@ gui = gr.Interface(
         gr.Textbox(lines=2, placeholder="Ask a question about the image...", label="Question")
     ],
     outputs=[
-        gr.Textbox(label="Result", lines=10),
+        gr.Textbox(label="Answer", lines=5),
         gr.Audio(label="Answer (Audio)", type="filepath")
     ],
     title="🧠 Image QA with Voice",
-    description="Upload an image and ask any question — even if there's no readable text. The app will use OCR, captioning, visual QA, and read answers aloud."
+    description="Upload an image and ask a question. You'll get a text + spoken answer."
 )
 
 app = gr.mount_gradio_app(app, gui, path="/")
 
 @app.get("/")
 def home():
-    return RedirectResponse(url="/")
+    return RedirectResponse(url="/")