TheOneReborn commited on
Commit
6d3ae22
·
1 Parent(s): d518e91

fix: resize image before prediction

Browse files
Files changed (1) hide show
  1. app.py +27 -23
app.py CHANGED
@@ -40,33 +40,37 @@ LEARNER = load_learner(MODEL_PATH / 'car-segmentation_v1.pkl')
40
 
41
 
42
  def segment_image(image):
43
- image = PILImage.create(image)
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._colormaps['jet']
59
- colored_mask = colormap(prediction_resized / numpy.max(prediction_resized))[:, :, :3] # Normalize & remove alpha
60
-
61
- # Convert PIL image to NumPy array
62
- image_array = numpy.array(image).astype(numpy.float32) / 255.0 # Normalize to [0,1]
63
-
64
- # Blend the original image and the mask
65
- overlay = (image_array * 0.7) + (colored_mask * 0.3)
66
-
67
- # Convert back to [0,255] uint8
 
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