Update main.py
Browse files
main.py
CHANGED
@@ -8,6 +8,7 @@ import logging
|
|
8 |
from io import BytesIO
|
9 |
import tempfile
|
10 |
import os
|
|
|
11 |
|
12 |
app = FastAPI()
|
13 |
|
@@ -20,9 +21,53 @@ def load_model():
|
|
20 |
model = Model(device='cuda' if torch.cuda.is_available() else 'cpu')
|
21 |
model.load_model('cartoon4')
|
22 |
|
|
|
|
|
|
|
|
|
23 |
# Configure logging
|
24 |
logging.basicConfig(level=logging.INFO)
|
25 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
26 |
@app.post("/upload/")
|
27 |
async def process_image(file: UploadFile = File(...), top: int = Form(...), bottom: int = Form(...), left: int = Form(...), right: int = Form(...)):
|
28 |
global model
|
@@ -42,13 +87,18 @@ async def process_image(file: UploadFile = File(...), top: int = Form(...), bott
|
|
42 |
|
43 |
logging.info(f"Uploaded image shape: {frame_bgr.shape}")
|
44 |
|
45 |
-
#
|
|
|
|
|
|
|
|
|
|
|
46 |
with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
|
47 |
-
cv2.imwrite(temp_file.name,
|
48 |
temp_file_path = temp_file.name
|
49 |
|
50 |
try:
|
51 |
-
# Process the
|
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.")
|
|
|
8 |
from io import BytesIO
|
9 |
import tempfile
|
10 |
import os
|
11 |
+
from insightface.app import FaceAnalysis
|
12 |
|
13 |
app = FastAPI()
|
14 |
|
|
|
21 |
model = Model(device='cuda' if torch.cuda.is_available() else 'cpu')
|
22 |
model.load_model('cartoon4')
|
23 |
|
24 |
+
# Initialize the InsightFace model for face detection
|
25 |
+
face_detector = FaceAnalysis(allowed_modules=['detection'])
|
26 |
+
face_detector.prepare(ctx_id=0 if torch.cuda.is_available() else -1, det_size=(640, 640))
|
27 |
+
|
28 |
# Configure logging
|
29 |
logging.basicConfig(level=logging.INFO)
|
30 |
|
31 |
+
def detect_and_crop_face(image, padding=0.6):
|
32 |
+
# Get original dimensions
|
33 |
+
orig_h, orig_w = image.shape[:2]
|
34 |
+
|
35 |
+
# Resize the image for detection
|
36 |
+
resized_image = cv2.resize(image, (640, 640))
|
37 |
+
|
38 |
+
# Detect faces on the resized image
|
39 |
+
faces = face_detector.get(resized_image)
|
40 |
+
|
41 |
+
# If faces are detected, sort by x-coordinate and select the leftmost face
|
42 |
+
if faces:
|
43 |
+
faces = sorted(faces, key=lambda face: face.bbox[0])
|
44 |
+
face = faces[0] # Select the leftmost face
|
45 |
+
bbox = face.bbox.astype(int)
|
46 |
+
|
47 |
+
# Calculate scaling factors
|
48 |
+
h_scale = orig_h / 640
|
49 |
+
w_scale = orig_w / 640
|
50 |
+
|
51 |
+
# Map the bounding box to the original image size
|
52 |
+
x1, y1, x2, y2 = bbox
|
53 |
+
x1 = int(x1 * w_scale)
|
54 |
+
y1 = int(y1 * h_scale)
|
55 |
+
x2 = int(x2 * w_scale)
|
56 |
+
y2 = int(y2 * h_scale)
|
57 |
+
|
58 |
+
# Calculate padding
|
59 |
+
width = x2 - x1
|
60 |
+
height = y2 - y1
|
61 |
+
x1 = max(0, x1 - int(padding * width))
|
62 |
+
y1 = max(0, y1 - int(padding * height))
|
63 |
+
x2 = min(orig_w, x2 + int(padding * width))
|
64 |
+
y2 = min(orig_h, y2 + int(padding * height))
|
65 |
+
|
66 |
+
cropped_face = image[y1:y2, x1:x2]
|
67 |
+
return cropped_face
|
68 |
+
|
69 |
+
return None
|
70 |
+
|
71 |
@app.post("/upload/")
|
72 |
async def process_image(file: UploadFile = File(...), top: int = Form(...), bottom: int = Form(...), left: int = Form(...), right: int = Form(...)):
|
73 |
global model
|
|
|
87 |
|
88 |
logging.info(f"Uploaded image shape: {frame_bgr.shape}")
|
89 |
|
90 |
+
# Detect and crop face
|
91 |
+
cropped_face = detect_and_crop_face(frame_bgr)
|
92 |
+
if cropped_face is None:
|
93 |
+
return {"error": "No face detected or alignment failed."}
|
94 |
+
|
95 |
+
# Save the cropped face temporarily
|
96 |
with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
|
97 |
+
cv2.imwrite(temp_file.name, cropped_face)
|
98 |
temp_file_path = temp_file.name
|
99 |
|
100 |
try:
|
101 |
+
# Process the cropped face using the file path
|
102 |
aligned_face, instyle, message = model.detect_and_align_image(temp_file_path, top, bottom, left, right)
|
103 |
if aligned_face is None or instyle is None:
|
104 |
logging.error("Failed to process the image: No face detected or alignment failed.")
|