diff --git a/app.py b/app.py index 5bbf00a..a10d3ec 100644 --- a/app.py +++ b/app.py @@ -1,7 +1,8 @@ import os import io +import base64 from dotenv import load_dotenv -from fastapi import FastAPI, File, UploadFile, Form, Query +from fastapi import FastAPI, Form from fastapi.responses import JSONResponse, StreamingResponse from PIL import Image import moondream as md @@ -16,14 +17,20 @@ model = md.vl(api_key=api_key) # FastAPI app app = FastAPI() +def decode_base64_image(base64_str: str) -> Image.Image: + try: + image_data = base64.b64decode(base64_str.split(",")[-1]) # strip data URL prefix if present + return Image.open(io.BytesIO(image_data)) + except Exception as e: + raise ValueError("Invalid base64 image") from e + @app.post("/caption") async def generate_caption( - image: UploadFile = File(...), + base64_image: str = Form(...), length: str = Form("short") ): try: - image_bytes = await image.read() - img = Image.open(io.BytesIO(image_bytes)) + img = decode_base64_image(base64_image) response = model.caption(img, length=length) return JSONResponse(content={"caption": response["caption"]}) except Exception as e: @@ -31,13 +38,12 @@ async def generate_caption( @app.post("/query") async def query_image( - image: UploadFile = File(...), + base64_image: str = Form(...), question: str = Form(...), stream: bool = Form(False) ): try: - image_bytes = await image.read() - img = Image.open(io.BytesIO(image_bytes)) + img = decode_base64_image(base64_image) if stream: def generate(): @@ -48,10 +54,7 @@ async def query_image( return StreamingResponse(generate(), media_type="text/plain") else: result = model.query(img, question) - return JSONResponse(content={ - "answer": result["answer"] - # "request_id": result["request_id"] - }) + return JSONResponse(content={"answer": result["answer"]}) except Exception as e: return JSONResponse(status_code=500, content={"error": str(e)}) diff --git a/draaw.jpg b/draaw.jpg new file mode 100644 index 0000000..301dab7 Binary files /dev/null and b/draaw.jpg differ