Update main.py

main.py

@@ -32,9 +32,19 @@ async def process_image(file: UploadFile = File(...), top: int = Form(...), bott
     # Read the uploaded image file
     contents = await file.read()
 
-    # Create a temporary file to save the uploaded image
+    # Convert the uploaded image to numpy array
+    nparr = np.frombuffer(contents, np.uint8)
+    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.")
+        return {"error": "Failed to decode the image. Please ensure the file is a valid image format."}
+
+    logging.info(f"Uploaded image shape: {frame_bgr.shape}")
+
+    # Save the uploaded image temporarily to pass the file path to the model
     with tempfile.NamedTemporaryFile(delete=False, suffix=".jpg") as temp_file:
-        temp_file.write(contents)
+        cv2.imwrite(temp_file.name, frame_bgr)
         temp_file_path = temp_file.name
 
     try:
@@ -49,8 +59,11 @@ async def process_image(file: UploadFile = File(...), top: int = Form(...), bott
             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)
+        _, 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")