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")
|