Update appImage.py

appImage.py

@@ -1,32 +1,60 @@
-# appImage.py
-from transformers import pipeline
-import tempfile, os
+
+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 tempfile
+import numpy as np
+from io import BytesIO
 
-captioner = pipeline("image-to-text", model="nlpconnect/vit-gpt2-image-captioning")
+app = FastAPI()
 
-async def caption_image(file):
-    contents = await file.read()
-    with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp:
-        tmp.write(contents)
-        image_path = tmp.name
+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")
+reader = easyocr.Reader(['en', 'fr'])
 
-    captions = captioner(image_path)
-    caption = captions[0]['generated_text'] if captions else "No caption generated."
+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"
 
-    audio_path = text_to_speech(caption)
+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
 
-    result = {"caption": caption}
-    if audio_path:
-        result["audioUrl"] = f"/files/{os.path.basename(audio_path)}"
-    return result
+    mode = classify_question(question)
 
-def text_to_speech(text: str):
     try:
-        tts = gTTS(text)
-        temp_audio = tempfile.NamedTemporaryFile(delete=False, suffix=".mp3")
-        tts.save(temp_audio.name)
-        return temp_audio.name
-    except:
-        return ""
+        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")