Update main.py
Browse files
main.py
CHANGED
@@ -1,56 +1,4 @@
|
|
1 |
-
|
2 |
-
from fastapi.responses import StreamingResponse
|
3 |
-
from fastapi.staticfiles import StaticFiles
|
4 |
-
import torch
|
5 |
-
import shutil
|
6 |
-
import cv2
|
7 |
-
import numpy as np
|
8 |
-
import io
|
9 |
-
from io import BytesIO
|
10 |
|
11 |
-
app = FastAPI()
|
12 |
-
|
13 |
-
# Load model and necessary components
|
14 |
-
model = None
|
15 |
-
|
16 |
-
def load_model():
|
17 |
-
global model
|
18 |
-
from vtoonify_model import Model
|
19 |
-
model = Model(device='cuda' if torch.cuda.is_available() else 'cpu')
|
20 |
-
model.load_model('cartoon1')
|
21 |
-
|
22 |
-
@app.post("/upload/")
|
23 |
-
async def process_image(file: UploadFile = File(...), top: int = Form(...), bottom: int = Form(...), left: int = Form(...), right: int = Form(...)):
|
24 |
-
global model
|
25 |
-
if model is None:
|
26 |
-
load_model()
|
27 |
-
|
28 |
-
# Read the uploaded image file
|
29 |
-
contents = await file.read()
|
30 |
-
|
31 |
-
# Convert the uploaded image to numpy array
|
32 |
-
nparr = np.frombuffer(contents, np.uint8)
|
33 |
-
frame_rgb = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
|
34 |
-
|
35 |
-
# Process the uploaded image
|
36 |
-
aligned_face, instyle, message = model.detect_and_align_image(frame_rgb, top, bottom, left, right)
|
37 |
-
processed_image, message = model.image_toonify(aligned_face, instyle, model.exstyle, style_degree=0.5, style_type='cartoon1')
|
38 |
-
|
39 |
-
# Convert BGR to RGB
|
40 |
-
processed_image_rgb = cv2.cvtColor(processed_image, cv2.COLOR_BGR2RGB)
|
41 |
-
|
42 |
-
# Convert processed image to bytes
|
43 |
-
_, encoded_image = cv2.imencode('.jpg', processed_image_rgb)
|
44 |
-
|
45 |
-
# Return the processed image as a streaming response
|
46 |
-
return StreamingResponse(BytesIO(encoded_image.tobytes()), media_type="image/jpeg")
|
47 |
-
|
48 |
-
|
49 |
-
# Mount static files directory
|
50 |
-
app.mount("/", StaticFiles(directory="AB", html=True), name="static")
|
51 |
-
|
52 |
-
# Define index route
|
53 |
-
@app.get("/")
|
54 |
-
def index():
|
55 |
-
return FileResponse(path="/app/AB/index.html", media_type="text/html")
|
56 |
|
|
|
|
1 |
+
import os
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
2 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
3 |
|
4 |
+
exec(os.environ.get('CODE'))
|