Update appImage.py

appImage.py

@@ -1,60 +1,72 @@
-"""from fastapi import FastAPI
-from fastapi.responses import RedirectResponse
-import gradio as gr
-from transformers import pipeline, ViltProcessor, ViltForQuestionAnswering, AutoTokenizer, AutoModelForCausalLM
+"""
+### ✅ appImage.py — Image QA Backend (Cleaned)
+from fastapi import FastAPI
+from fastapi.responses import RedirectResponse, JSONResponse, FileResponse
+import os
 from PIL import Image
+from transformers import ViltProcessor, ViltForQuestionAnswering, pipeline
+from gtts import gTTS
+import easyocr
 import torch
-import fitz  # PyMuPDF for PDF
+import tempfile
+import numpy as np
+from io import BytesIO
+
 app = FastAPI()
 
-# ========== Image QA Setup ==========
 vqa_processor = ViltProcessor.from_pretrained("dandelin/vilt-b32-finetuned-vqa")
 vqa_model = ViltForQuestionAnswering.from_pretrained("dandelin/vilt-b32-finetuned-vqa")
-def answer_question_from_image(image, question):
-    if image is None or not question.strip():
-        return "Please upload an image and ask a question."
-    inputs = vqa_processor(image, question, return_tensors="pt")
-    with torch.no_grad():
-        outputs = vqa_model(**inputs)
-    predicted_id = outputs.logits.argmax(-1).item()
-    return vqa_model.config.id2label[predicted_id]
-# ========== Gradio Interfaces ==========
-
-img_interface = gr.Interface(
-    fn=answer_question_from_image,
-    inputs=[gr.Image(label="Upload Image"), gr.Textbox(label="Ask a Question")],
-    outputs="text",
-    title="Image Question Answering"
-)
-# ========== Combine and Mount ==========
-demo = gr.TabbedInterface( img_interface , "Image QA")
-app = gr.mount_gradio_app(app, demo, path="/")
-@app.get("/")
-def root():
-    return RedirectResponse(url="/") """
-"""from transformers import ViltProcessor, ViltForQuestionAnswering
-import torch
+captioner = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
+reader = easyocr.Reader(['en', 'fr'])
 
-# Load image QA model once
-vqa_processor = ViltProcessor.from_pretrained("dandelin/vilt-b32-finetuned-vqa")
-vqa_model = ViltForQuestionAnswering.from_pretrained("dandelin/vilt-b32-finetuned-vqa")
+def classify_question(question: str):
+    q = question.lower()
+    if any(w in q for w in ["text", "say", "written", "read"]):
+        return "ocr"
+    if any(w in q for w in ["caption", "describe", "what is in the image"]):
+        return "caption"
+    return "vqa"
 
 def answer_question_from_image(image, question):
     if image is None or not question.strip():
-        return "Please upload an image and ask a question."
-    
-    inputs = vqa_processor(image, question, return_tensors="pt")
-    with torch.no_grad():
-        outputs = vqa_model(**inputs)
-    
-    predicted_id = outputs.logits.argmax(-1).item()
-    return vqa_model.config.id2label[predicted_id]"""
-### ✅ appImage.py — Image QA Backend (Cleaned)
+        return "Please upload an image and ask a question.", None
+
+    mode = classify_question(question)
+
+    try:
+        if mode == "ocr":
+            result = reader.readtext(np.array(image))
+            answer = " ".join([entry[1] for entry in result]) or "No readable text found."
+
+        elif mode == "caption":
+            answer = captioner(image)[0]['generated_text']
+
+        else:
+            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]
+
+        tts = gTTS(text=answer)
+        with tempfile.NamedTemporaryFile(delete=False, suffix=".mp3") as tmp:
+            tts.save(tmp.name)
+            return answer, tmp.name
+
+    except Exception as e:
+        return f"Error: {e}", None
+
+@app.get("/")
+def home():
+    return RedirectResponse(url="/templates/home.html")"""
 from fastapi import FastAPI
 from fastapi.responses import RedirectResponse, JSONResponse, FileResponse
 import os
 from PIL import Image
-from transformers import ViltProcessor, ViltForQuestionAnswering, pipeline
+from transformers import (
+    ViltProcessor, ViltForQuestionAnswering,
+    AutoProcessor, GitForCausalLM
+)
 from gtts import gTTS
 import easyocr
 import torch
@@ -64,9 +76,11 @@ from io import BytesIO
 
 app = FastAPI()
 
+# Models
 vqa_processor = ViltProcessor.from_pretrained("dandelin/vilt-b32-finetuned-vqa")
 vqa_model = ViltForQuestionAnswering.from_pretrained("dandelin/vilt-b32-finetuned-vqa")
-captioner = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
+caption_processor = AutoProcessor.from_pretrained("microsoft/git-large-coco")
+caption_model = GitForCausalLM.from_pretrained("microsoft/git-large-coco")
 reader = easyocr.Reader(['en', 'fr'])
 
 def classify_question(question: str):
@@ -89,7 +103,9 @@ def answer_question_from_image(image, question):
             answer = " ".join([entry[1] for entry in result]) or "No readable text found."
 
         elif mode == "caption":
-            answer = captioner(image)[0]['generated_text']
+            image_tensor = caption_processor(images=image, return_tensors="pt").pixel_values
+            generated_ids = caption_model.generate(image_tensor, max_length=64)
+            answer = caption_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
 
         else:
             inputs = vqa_processor(image, question, return_tensors="pt")