Update main.py

main.py

@@ -1,4 +1,4 @@
-from fastapi import FastAPI, File, UploadFile
+from fastapi import FastAPI, File, UploadFile, Form
 from fastapi.responses import StreamingResponse, FileResponse
 from fastapi.staticfiles import StaticFiles
 import torch
@@ -23,28 +23,8 @@ def load_model():
 # Configure logging
 logging.basicConfig(level=logging.INFO)
 
-def detect_and_crop_face(image):
-    # Load the pre-trained face detector
-    face_cascade = cv2.CascadeClassifier(cv2.data.haarcascades + 'haarcascade_frontalface_default.xml')
-    
-    # Convert to grayscale
-    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-    
-    # Detect faces
-    faces = face_cascade.detectMultiScale(gray, scaleFactor=1.1, minNeighbors=5, minSize=(30, 30))
-    
-    # If faces are detected, crop the first face
-    if len(faces) > 0:
-        x, y, w, h = faces[0]
-        cropped_face = image[y:y+h, x:x+w]
-        return cropped_face
-    
-    # Log if no face is detected
-    logging.warning("No face detected.")
-    return None
-
 @app.post("/upload/")
-async def process_image(file: UploadFile = File(...)):
+async def process_image(file: UploadFile = File(...), top: int = Form(...), bottom: int = Form(...), left: int = Form(...), right: int = Form(...)):
     global model
     if model is None:
         load_model()
@@ -54,7 +34,7 @@ async def process_image(file: UploadFile = File(...)):
 
     # Convert the uploaded image to numpy array
     nparr = np.frombuffer(contents, np.uint8)
-    frame_bgr = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
+    frame_bgr = cv2.imdecode(nparr, cv2.IMREAD_COLOR)  # Read as BGR format by default
     
     if frame_bgr is None:
         logging.error("Failed to decode the image.")
@@ -62,19 +42,14 @@ async def process_image(file: UploadFile = File(...)):
 
     logging.info(f"Uploaded image shape: {frame_bgr.shape}")
 
-    # Detect and crop face
-    cropped_face = detect_and_crop_face(frame_bgr)
-    if cropped_face is None:
-        return {"error": "No face detected or alignment failed."}
-
-    # Save the cropped face temporarily
+    # Save the uploaded image temporarily to pass the file path to the model
     with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
-        cv2.imwrite(temp_file.name, cropped_face)
+        cv2.imwrite(temp_file.name, frame_bgr)
         temp_file_path = temp_file.name
 
     try:
-        # Process the cropped face
-        aligned_face, instyle, message = model.detect_and_align_image(temp_file_path, 0, 0, 0, 0)
+        # Process the uploaded image using the file path
+        aligned_face, instyle, message = model.detect_and_align_image(temp_file_path, top, bottom, left, right)
         if aligned_face is None or instyle is None:
             logging.error("Failed to process the image: No face detected or alignment failed.")
             return {"error": message}
@@ -84,12 +59,17 @@ async def process_image(file: UploadFile = File(...)):
             logging.error("Failed to toonify the image.")
             return {"error": message}
 
+        # Convert the processed image to RGB before returning
         processed_image_rgb = cv2.cvtColor(processed_image, cv2.COLOR_BGR2RGB)
+
+        # Convert processed image to bytes
         _, encoded_image = cv2.imencode('.jpg', processed_image_rgb)
 
+        # Return the processed image as a streaming response
         return StreamingResponse(BytesIO(encoded_image.tobytes()), media_type="image/jpeg")
     
     finally:
+        # Clean up the temporary file
         os.remove(temp_file_path)
 
 # Mount static files directory