Update appImage.py

appImage.py

@@ -43,12 +43,11 @@ async def caption_from_frontend(file: UploadFile = File(...)):
 @app.get("/")
 def home():
     return RedirectResponse(url="/")"""
-# appImage.py
-from fastapi import UploadFile, File
-from transformers import AutoProcessor, AutoModelForCausalLM
-from transformers import pipeline
+from fastapi import UploadFile
+from transformers import AutoProcessor, AutoModelForCausalLM, pipeline
 from PIL import Image
 import tempfile
+import os
 import torch
 
 # Load model
@@ -63,7 +62,7 @@ except Exception:
 def generate_caption(image_path):
     try:
         if USE_GIT:
-            image = Image.open(image_path)
+            image = Image.open(image_path).convert("RGB")
             inputs = processor(images=image, return_tensors="pt")
             outputs = model.generate(**inputs, max_length=50)
             return processor.batch_decode(outputs, skip_special_tokens=True)[0]
@@ -73,14 +72,26 @@ def generate_caption(image_path):
     except Exception as e:
         return f"Error generating caption: {str(e)}"
 
-async def caption_image(file: UploadFile = File(...)):
+async def caption_image(file: UploadFile):
     try:
-        contents = await file.read()
-        with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp:
+        # Get file extension correctly
+        _, ext = os.path.splitext(file.filename)
+        if ext.lower() not in [".jpg", ".jpeg", ".png", ".bmp", ".gif"]:
+            return {"error": "Unsupported file type"}
+
+        # Save the uploaded image with correct extension
+        with tempfile.NamedTemporaryFile(delete=False, suffix=ext) as tmp:
+            contents = await file.read()
             tmp.write(contents)
             tmp_path = tmp.name
 
+        # Generate caption
         caption = generate_caption(tmp_path)
-        return  caption
+
+        # Handle errors inside generate_caption
+        if caption.startswith("Error"):
+            return {"error": caption}
+        return {"caption": caption}
+
     except Exception as e:
         return {"error": f"Failed to generate caption: {str(e)}"}