Ashrafb commited on
Commit
4a6b773
1 Parent(s): 79c7733

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +2 -54
main.py CHANGED
@@ -1,56 +1,4 @@
1
- from fastapi import FastAPI, File, UploadFile, Form
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'))