Spaces:
Running
Running
Commit
·
6d3ae22
1
Parent(s):
d518e91
fix: resize image before prediction
Browse files
app.py
CHANGED
@@ -40,33 +40,37 @@ LEARNER = load_learner(MODEL_PATH / 'car-segmentation_v1.pkl')
|
|
40 |
|
41 |
|
42 |
def segment_image(image):
|
43 |
-
|
44 |
-
prediction, _, _ = LEARNER.predict(image)
|
45 |
-
|
46 |
-
print("Prediction shape:", prediction.shape)
|
47 |
-
print("Unique values:", numpy.unique(prediction))
|
48 |
-
|
49 |
-
# Convert prediction to NumPy array
|
50 |
-
prediction_array = numpy.asarray(prediction, dtype=numpy.uint8)
|
51 |
-
|
52 |
-
# Resize the mask to match the original image size
|
53 |
original_size = image.size # (width, height)
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
54 |
prediction_resized = Image.fromarray(prediction_array).resize(original_size, Image.NEAREST)
|
55 |
prediction_resized = numpy.array(prediction_resized)
|
56 |
-
|
57 |
-
# Apply a colormap for visualization
|
58 |
-
colormap = cm.
|
59 |
-
|
60 |
-
|
61 |
-
|
62 |
-
|
63 |
-
|
64 |
-
|
65 |
-
|
66 |
-
|
67 |
-
|
|
|
68 |
overlay = (overlay * 255).astype(numpy.uint8)
|
69 |
-
|
70 |
return overlay
|
71 |
|
72 |
|
|
|
40 |
|
41 |
|
42 |
def segment_image(image):
|
43 |
+
# Store original size
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
44 |
original_size = image.size # (width, height)
|
45 |
+
|
46 |
+
# Resize the input image to 400x400 for the model
|
47 |
+
resized_image = image.resize((400, 400))
|
48 |
+
resized_image = PILImage.create(resized_image)
|
49 |
+
|
50 |
+
# Get the prediction from the model
|
51 |
+
prediction, _, _ = LEARNER.predict(resized_image)
|
52 |
+
|
53 |
+
# Convert prediction to a NumPy array
|
54 |
+
prediction_array = numpy.asarray(prediction, dtype=numpy.uint8)
|
55 |
+
|
56 |
+
# Resize the mask back to the original image size
|
57 |
prediction_resized = Image.fromarray(prediction_array).resize(original_size, Image.NEAREST)
|
58 |
prediction_resized = numpy.array(prediction_resized)
|
59 |
+
|
60 |
+
# Apply a colormap for visualization (using the public API)
|
61 |
+
colormap = cm.colormaps['jet']
|
62 |
+
# Normalize the mask and apply the colormap (result is in float [0,1])
|
63 |
+
colored_mask = colormap(prediction_resized / numpy.max(prediction_resized))[:, :, :3]
|
64 |
+
|
65 |
+
# Convert the original image to a NumPy array and normalize to [0,1]
|
66 |
+
image_array = numpy.array(image).astype(numpy.float32) / 255.0
|
67 |
+
|
68 |
+
# Blend the original image and the colored mask
|
69 |
+
overlay = (image_array * 0.7) + (colored_mask * 0.3)
|
70 |
+
|
71 |
+
# Convert the blended image back to uint8
|
72 |
overlay = (overlay * 255).astype(numpy.uint8)
|
73 |
+
|
74 |
return overlay
|
75 |
|
76 |
|