File size: 1,393 Bytes
6b79e42
 
99f1eb3
6b79e42
 
0fbc2d1
6b79e42
6767de4
99f1eb3
 
 
 
6b79e42
6767de4
6b79e42
 
 
 
 
d4dc97e
47da34c
 
 
408728b
47da34c
 
 
 
408728b
47da34c
 
99f1eb3
 
 
 
 
 
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
37
38
from fastapi import FastAPI, File, UploadFile,Form
from fastapi.responses import FileResponse, StreamingResponse
from fastapi.staticfiles import StaticFiles
from fastapi import FastAPI, File, UploadFile
from pydantic import BaseModel
import numpy as np
import cv2
import torch
from vtoonify_model import Model

app = FastAPI()

# Load the model
model = Model(device='cuda' if torch.cuda.is_available() else 'cpu')
exstyle, message = model.load_model("cartoon1")

class ImageRequest(BaseModel):
    image_file: UploadFile = File(...)

@app.post("/upload/")
async def toonify_image(image_file: UploadFile = File(...)):
    try:
        contents = await image_file.read()
        filename = image_file.filename
        nparr = np.frombuffer(contents, np.uint8)
        img = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        aligned_face, instyle, message = model.detect_and_align_image(img, 200, 200, 200, 200)  # Hardcoded values
        toonified_img, message = model.image_toonify(aligned_face, instyle, exstyle, style_degree=0.5, style_type="cartoon1")
        return {"toonified_image": toonified_img, "message": message, "filename": filename}
    except Exception as e:
        return {"error": str(e)}

app.mount("/", StaticFiles(directory="AB", html=True), name="static")

@app.get("/")
def index() -> FileResponse:
    return FileResponse(path="/app/AB/index.html", media_type="text/html")