img_base64
parent
a2b6730b88
commit
2f7aab5891
25
app.py
25
app.py
|
@ -1,7 +1,8 @@
|
||||||
import os
|
import os
|
||||||
import io
|
import io
|
||||||
|
import base64
|
||||||
from dotenv import load_dotenv
|
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 fastapi.responses import JSONResponse, StreamingResponse
|
||||||
from PIL import Image
|
from PIL import Image
|
||||||
import moondream as md
|
import moondream as md
|
||||||
|
@ -16,14 +17,20 @@ model = md.vl(api_key=api_key)
|
||||||
# FastAPI app
|
# FastAPI app
|
||||||
app = FastAPI()
|
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")
|
@app.post("/caption")
|
||||||
async def generate_caption(
|
async def generate_caption(
|
||||||
image: UploadFile = File(...),
|
base64_image: str = Form(...),
|
||||||
length: str = Form("short")
|
length: str = Form("short")
|
||||||
):
|
):
|
||||||
try:
|
try:
|
||||||
image_bytes = await image.read()
|
img = decode_base64_image(base64_image)
|
||||||
img = Image.open(io.BytesIO(image_bytes))
|
|
||||||
response = model.caption(img, length=length)
|
response = model.caption(img, length=length)
|
||||||
return JSONResponse(content={"caption": response["caption"]})
|
return JSONResponse(content={"caption": response["caption"]})
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
|
@ -31,13 +38,12 @@ async def generate_caption(
|
||||||
|
|
||||||
@app.post("/query")
|
@app.post("/query")
|
||||||
async def query_image(
|
async def query_image(
|
||||||
image: UploadFile = File(...),
|
base64_image: str = Form(...),
|
||||||
question: str = Form(...),
|
question: str = Form(...),
|
||||||
stream: bool = Form(False)
|
stream: bool = Form(False)
|
||||||
):
|
):
|
||||||
try:
|
try:
|
||||||
image_bytes = await image.read()
|
img = decode_base64_image(base64_image)
|
||||||
img = Image.open(io.BytesIO(image_bytes))
|
|
||||||
|
|
||||||
if stream:
|
if stream:
|
||||||
def generate():
|
def generate():
|
||||||
|
@ -48,10 +54,7 @@ async def query_image(
|
||||||
return StreamingResponse(generate(), media_type="text/plain")
|
return StreamingResponse(generate(), media_type="text/plain")
|
||||||
else:
|
else:
|
||||||
result = model.query(img, question)
|
result = model.query(img, question)
|
||||||
return JSONResponse(content={
|
return JSONResponse(content={"answer": result["answer"]})
|
||||||
"answer": result["answer"]
|
|
||||||
# "request_id": result["request_id"]
|
|
||||||
})
|
|
||||||
|
|
||||||
except Exception as e:
|
except Exception as e:
|
||||||
return JSONResponse(status_code=500, content={"error": str(e)})
|
return JSONResponse(status_code=500, content={"error": str(e)})
|
||||||
|
|
Loading…
Reference in New Issue