File size: 992 Bytes
a014e90
 
c41dc64
21f70e6
fd7d432
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
import os
os.environ["TRANSFORMERS_CACHE"] = "/tmp"
os.environ["HF_HOME"] = "/tmp"
os.environ["TORCH_HOME"] = "/tmp"
from fastapi import FastAPI, UploadFile, File, Form
from PIL import Image, ImageOps
from classification import classify_image
from smolVLMchat import process_chat
from io import BytesIO

app = FastAPI()

@app.post("/classify/")
async def classify(
    file: UploadFile = File(...),
    model_name: str = Form(...) 
):
    """Endpoint to classify an uploaded image with a chosen model."""

    contents = await file.read()

    # Decodează imaginea
    image = Image.open(BytesIO(contents)).convert("RGB")
    image = ImageOps.exif_transpose(image)

    result = classify_image(image, model_name=model_name)

    return result


@app.post("/chat/")
async def chat(text: str = Form(None), image: UploadFile = File(None)):
    """Endpoint to chat with SmolVLM using text, image, or both."""
    response = process_chat(text=text, image=image)
    return {"response": response}