fix: resize image before prediction

app.py

@@ -40,33 +40,37 @@ LEARNER = load_learner(MODEL_PATH / 'car-segmentation_v1.pkl')
 
 
 def segment_image(image):
-    image = PILImage.create(image)
-    prediction, _, _ = LEARNER.predict(image)
-
-    print("Prediction shape:", prediction.shape)
-    print("Unique values:", numpy.unique(prediction))
-
-    # Convert prediction to NumPy array
-    prediction_array = numpy.asarray(prediction, dtype=numpy.uint8)
-
-    # Resize the mask to match the original image size
+    # Store original size
     original_size = image.size  # (width, height)
+    
+    # Resize the input image to 400x400 for the model
+    resized_image = image.resize((400, 400))
+    resized_image = PILImage.create(resized_image)
+    
+    # Get the prediction from the model
+    prediction, _, _ = LEARNER.predict(resized_image)
+    
+    # Convert prediction to a NumPy array
+    prediction_array = numpy.asarray(prediction, dtype=numpy.uint8)
+    
+    # Resize the mask back to the original image size
     prediction_resized = Image.fromarray(prediction_array).resize(original_size, Image.NEAREST)
     prediction_resized = numpy.array(prediction_resized)
-
-    # Apply a colormap for visualization
-    colormap = cm._colormaps['jet']
-    colored_mask = colormap(prediction_resized / numpy.max(prediction_resized))[:, :, :3]  # Normalize & remove alpha
-
-    # Convert PIL image to NumPy array
-    image_array = numpy.array(image).astype(numpy.float32) / 255.0  # Normalize to [0,1]
-
-    # Blend the original image and the mask
-    overlay = (image_array * 0.7) + (colored_mask * 0.3)  
-
-    # Convert back to [0,255] uint8
+    
+    # Apply a colormap for visualization (using the public API)
+    colormap = cm.colormaps['jet']
+    # Normalize the mask and apply the colormap (result is in float [0,1])
+    colored_mask = colormap(prediction_resized / numpy.max(prediction_resized))[:, :, :3]
+    
+    # Convert the original image to a NumPy array and normalize to [0,1]
+    image_array = numpy.array(image).astype(numpy.float32) / 255.0
+    
+    # Blend the original image and the colored mask
+    overlay = (image_array * 0.7) + (colored_mask * 0.3)
+    
+    # Convert the blended image back to uint8
     overlay = (overlay * 255).astype(numpy.uint8)
-
+    
     return overlay