DavidFM43 commited on
Commit
d38fb36
·
1 Parent(s): 3ae192a

small bug fix

Browse files
Files changed (1) hide show
  1. app.py +1 -1
app.py CHANGED
@@ -66,7 +66,7 @@ transform = transforms.Compose(
66
  def get_overlay(image: torch.Tensor, preds: torch.Tensor, alpha: float) -> torch.Tensor:
67
  """Generates the segmentation ovelay for an satellite image"""
68
 
69
- masks = preds.squeeze().label_to_onehot(masks, 7)
70
  overlay = draw_segmentation_masks(
71
  image, masks=masks, alpha=alpha, colors=class_rgb_colors
72
  )
 
66
  def get_overlay(image: torch.Tensor, preds: torch.Tensor, alpha: float) -> torch.Tensor:
67
  """Generates the segmentation ovelay for an satellite image"""
68
 
69
+ masks = label_to_onehot(preds.squeeze(), 7)
70
  overlay = draw_segmentation_masks(
71
  image, masks=masks, alpha=alpha, colors=class_rgb_colors
72
  )