Ashrafb commited on
Commit
daff9a2
·
verified ·
1 Parent(s): 4cc1f53

Update main.py

Browse files
Files changed (1) hide show
  1. main.py +53 -3
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
- # 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.")
 
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.")