File size: 1,560 Bytes
6b79e42 99f1eb3 6b79e42 0fbc2d1 6b79e42 6767de4 99f1eb3 814da2f 99f1eb3 6767de4 c0cfc30 61b6cd2 814da2f 47da34c 814da2f c0cfc30 814da2f c0cfc30 814da2f c0cfc30 814da2f 47da34c 814da2f 99f1eb3 814da2f |
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 39 40 41 42 43 44 45 46 47 48 |
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
import shutil
import os
app = FastAPI()
model = Model(device='cuda' if torch.cuda.is_available() else 'cpu')
@app.post("/upload/")
async def upload_image(file: UploadFile = File(...), top: int = Form(...), bottom: int = Form(...), left: int = Form(...), right: int = Form(...)):
try:
with open("temp.jpg", "wb") as buffer:
shutil.copyfileobj(file.file, buffer)
# Perform image processing using the model
result_image, _, _ = model.detect_and_align_image("temp.jpg", top, bottom, left, right)
# Save the result image temporarily
result_path = "result.jpg"
cv2.imwrite(result_path, result_image)
# Return the result image
with open(result_path, "rb") as result_buffer:
result_image_bytes = result_buffer.read()
# Remove temporary files
os.remove("temp.jpg")
os.remove(result_path)
return {"result_image": result_image_bytes}
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")
|