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