Update app2.py
Browse files
app2.py
CHANGED
@@ -6,6 +6,8 @@ import cv2
|
|
6 |
import numpy as np
|
7 |
import logging
|
8 |
from io import BytesIO
|
|
|
|
|
9 |
|
10 |
app = FastAPI()
|
11 |
|
@@ -32,30 +34,43 @@ async def process_image(file: UploadFile = File(...), top: int = Form(...), bott
|
|
32 |
|
33 |
# Convert the uploaded image to numpy array
|
34 |
nparr = np.frombuffer(contents, np.uint8)
|
35 |
-
|
36 |
|
37 |
-
if
|
38 |
logging.error("Failed to decode the image.")
|
39 |
return {"error": "Failed to decode the image. Please ensure the file is a valid image format."}
|
40 |
|
41 |
-
logging.info(f"Uploaded image shape: {
|
42 |
|
43 |
-
#
|
44 |
-
|
45 |
-
|
46 |
-
|
47 |
-
return {"error": message}
|
48 |
|
49 |
-
|
50 |
-
|
51 |
-
|
52 |
-
|
|
|
|
|
53 |
|
54 |
-
|
55 |
-
|
|
|
|
|
56 |
|
57 |
-
|
58 |
-
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
59 |
|
60 |
# Mount static files directory
|
61 |
app.mount("/", StaticFiles(directory="AB", html=True), name="static")
|
|
|
6 |
import numpy as np
|
7 |
import logging
|
8 |
from io import BytesIO
|
9 |
+
import tempfile
|
10 |
+
import os
|
11 |
|
12 |
app = FastAPI()
|
13 |
|
|
|
34 |
|
35 |
# Convert the uploaded image to numpy array
|
36 |
nparr = np.frombuffer(contents, np.uint8)
|
37 |
+
frame_bgr = cv2.imdecode(nparr, cv2.IMREAD_COLOR) # Read as BGR format by default
|
38 |
|
39 |
+
if frame_bgr is None:
|
40 |
logging.error("Failed to decode the image.")
|
41 |
return {"error": "Failed to decode the image. Please ensure the file is a valid image format."}
|
42 |
|
43 |
+
logging.info(f"Uploaded image shape: {frame_bgr.shape}")
|
44 |
|
45 |
+
# Save the uploaded image temporarily to pass the file path to the model
|
46 |
+
with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
|
47 |
+
cv2.imwrite(temp_file.name, frame_bgr)
|
48 |
+
temp_file_path = temp_file.name
|
|
|
49 |
|
50 |
+
try:
|
51 |
+
# Process the uploaded image using the file path
|
52 |
+
aligned_face, instyle, message = model.detect_and_align_image(temp_file_path, top, bottom, left, right)
|
53 |
+
if aligned_face is None or instyle is None:
|
54 |
+
logging.error("Failed to process the image: No face detected or alignment failed.")
|
55 |
+
return {"error": message}
|
56 |
|
57 |
+
processed_image, message = model.image_toonify(aligned_face, instyle, model.exstyle, style_degree=0.5, style_type='cartoon1')
|
58 |
+
if processed_image is None:
|
59 |
+
logging.error("Failed to toonify the image.")
|
60 |
+
return {"error": message}
|
61 |
|
62 |
+
# Convert the processed image to RGB before returning
|
63 |
+
processed_image_rgb = cv2.cvtColor(processed_image, cv2.COLOR_BGR2RGB)
|
64 |
+
|
65 |
+
# Convert processed image to bytes
|
66 |
+
_, encoded_image = cv2.imencode('.jpg', processed_image_rgb)
|
67 |
+
|
68 |
+
# Return the processed image as a streaming response
|
69 |
+
return StreamingResponse(BytesIO(encoded_image.tobytes()), media_type="image/jpeg")
|
70 |
+
|
71 |
+
finally:
|
72 |
+
# Clean up the temporary file
|
73 |
+
os.remove(temp_file_path)
|
74 |
|
75 |
# Mount static files directory
|
76 |
app.mount("/", StaticFiles(directory="AB", html=True), name="static")
|