Update main.py

main.py

@@ -20,6 +20,15 @@ net.getLayer(conv8).blobs = [np.full([1, 313], 2.606, dtype='float32')]
 # Initialize FastAPI app
 app = FastAPI()
 
+def read_image(file):
+    # Open the image with Pillow
+    pil_image = Image.open(file)
+
+    # Convert the Pillow image to a NumPy array
+    image_array = np.array(pil_image)
+
+    return image_array
+
 def colorize_image(image):
     # Convert the PIL image to OpenCV format
     image = np.array(image)
@@ -38,23 +47,25 @@ def colorize_image(image):
 
     L = cv2.split(lab)[0]
     colorized = np.concatenate((L[:, :, np.newaxis], ab), axis=2)
-    colorized = cv2.cvtColor(colorized, cv2.COLOR_LAB2BGR)
+    colorized = cv2.cvtColor(colorized, cv2.COLOR_LAB2RGB)
     colorized = np.clip(colorized, 0, 1)
     colorized = (255 * colorized).astype("uint8")
 
-    # Convert the colorized image from BGR to RGB
-    colorized = cv2.cvtColor(colorized, cv2.COLOR_BGR2RGB)
-
     return colorized
 
 @app.post("/upload/")
 async def upload(file: UploadFile = File(...)):
-    input_image = Image.open(file.file).convert("RGB")
-    colorized_image = colorize_image(input_image)
+    # Read the image using Pillow and convert to NumPy array
+    image_array = read_image(file.file)
+    
+    # Process the image
+    colorized_image = colorize_image(image_array)
 
     # Convert the colorized image to bytes
-    _, buffer = cv2.imencode('.png', colorized_image)
-    io_buf = io.BytesIO(buffer)
+    pil_image = Image.fromarray(colorized_image)
+    buf = io.BytesIO()
+    pil_image.save(buf, format="PNG")
+    buf.seek(0)
 
-    return StreamingResponse(io_buf, media_type="image/png")
+    return StreamingResponse(buf, media_type="image/png")