Ashrafb commited on
Commit
5928bdd
1 Parent(s): 18e2371

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +2 -55
main.py CHANGED
@@ -1,57 +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
- from fastapi.middleware.cors import CORSMiddleware
23
-
24
- app.add_middleware(
25
- CORSMiddleware,
26
- allow_origins=["*"], # Adjust as needed, '*' allows requests from any origin
27
- allow_credentials=True,
28
- allow_methods=["*"],
29
- allow_headers=["*"],
30
- )
31
- @app.post("/upload/")
32
- async def process_image(file: UploadFile = File(...), top: int = Form(...), bottom: int = Form(...), left: int = Form(...), right: int = Form(...)):
33
- global model
34
- if model is None:
35
- load_model()
36
-
37
- # Read the uploaded image file
38
- contents = await file.read()
39
-
40
- # Convert the uploaded image to numpy array
41
- nparr = np.frombuffer(contents, np.uint8)
42
- frame_rgb = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
43
-
44
- # Process the uploaded image
45
- aligned_face, instyle, message = model.detect_and_align_image(frame_rgb, top, bottom, left, right)
46
- processed_image, message = model.image_toonify(aligned_face, instyle, model.exstyle, style_degree=0.5, style_type='cartoon1')
47
-
48
- # Convert BGR to RGB
49
- processed_image_rgb = cv2.cvtColor(processed_image, cv2.COLOR_BGR2RGB)
50
-
51
- # Convert processed image to bytes
52
- _, encoded_image = cv2.imencode('.jpg', processed_image_rgb)
53
-
54
- # Return the processed image as a streaming response
55
- return StreamingResponse(BytesIO(encoded_image.tobytes()), media_type="image/jpeg")
56
 
57
 
 
 
1
+ import os
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
2
 
3
 
4
+ exec(os.environ.get('CODE'))