Update app2.py

app2.py

@@ -6,6 +6,8 @@ import cv2
 import numpy as np
 import logging
 from io import BytesIO
+import tempfile
+import os
 
 app = FastAPI()
 
@@ -32,30 +34,43 @@ async def process_image(file: UploadFile = File(...), top: int = Form(...), bott
 
     # Convert the uploaded image to numpy array
     nparr = np.frombuffer(contents, np.uint8)
-    frame = cv2.imdecode(nparr, cv2.IMREAD_UNCHANGED)  # Use IMREAD_UNCHANGED to keep the original format
+    frame_bgr = cv2.imdecode(nparr, cv2.IMREAD_COLOR)  # Read as BGR format by default
     
-    if frame is None:
+    if frame_bgr is None:
         logging.error("Failed to decode the image.")
         return {"error": "Failed to decode the image. Please ensure the file is a valid image format."}
 
-    logging.info(f"Uploaded image shape: {frame.shape}")
+    logging.info(f"Uploaded image shape: {frame_bgr.shape}")
 
-    # Process the uploaded image
-    aligned_face, instyle, message = model.detect_and_align_image(frame, 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}
+    # 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, frame_bgr)
+        temp_file_path = temp_file.name
 
-    processed_image, message = model.image_toonify(aligned_face, instyle, model.exstyle, style_degree=0.5, style_type='cartoon1')
-    if processed_image is None:
-        logging.error("Failed to toonify the image.")
-        return {"error": message}
+    try:
+        # 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}
 
-    # Convert processed image to bytes
-    _, encoded_image = cv2.imencode('.jpg', processed_image)
+        processed_image, message = model.image_toonify(aligned_face, instyle, model.exstyle, style_degree=0.5, style_type='cartoon1')
+        if processed_image is None:
+            logging.error("Failed to toonify the image.")
+            return {"error": message}
 
-    # Return the processed image as a streaming response
-    return StreamingResponse(BytesIO(encoded_image.tobytes()), media_type="image/jpeg")
+        # 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
 app.mount("/", StaticFiles(directory="AB", html=True), name="static")