CodeJackR commited on
Commit
b0c81ad
·
1 Parent(s): 0e71822

Remove image resizing

Browse files
Files changed (1) hide show
  1. handler.py +4 -21
handler.py CHANGED
@@ -64,9 +64,6 @@ class EndpointHandler():
64
 
65
  # 4. Process and select the best mask
66
  try:
67
- # Get original image dimensions
68
- original_height, original_width = img.size[1], img.size[0]
69
-
70
  # Get predicted masks and scores
71
  predicted_masks = outputs.pred_masks.cpu()
72
  iou_scores = outputs.iou_scores.cpu()[0]
@@ -75,31 +72,17 @@ class EndpointHandler():
75
  if predicted_masks.ndim == 5:
76
  predicted_masks = predicted_masks.squeeze(1)
77
 
78
- # Resize masks to standard size first
79
- predicted_masks = torch.nn.functional.interpolate(
80
- predicted_masks,
81
- size=(1024, 1024),
82
- mode='bilinear',
83
- align_corners=False
84
- )
85
-
86
  # Select the best mask
87
  best_mask_idx = torch.argmax(iou_scores)
88
  best_mask = predicted_masks[0, best_mask_idx, :, :]
89
 
90
- # Resize to original image dimensions
91
- final_mask = torch.nn.functional.interpolate(
92
- best_mask.unsqueeze(0).unsqueeze(0),
93
- size=(original_height, original_width),
94
- mode='bilinear',
95
- align_corners=False
96
- ).squeeze()
97
-
98
- # Convert to binary mask
99
- mask_binary = (final_mask > 0.0).numpy().astype(np.uint8) * 255
100
 
101
  except Exception as e:
102
  print("Error processing masks: {}".format(e))
 
 
103
  mask_binary = np.zeros((height, width), dtype=np.uint8)
104
  center_x, center_y = width // 2, height // 2
105
  size = min(width, height) // 8
 
64
 
65
  # 4. Process and select the best mask
66
  try:
 
 
 
67
  # Get predicted masks and scores
68
  predicted_masks = outputs.pred_masks.cpu()
69
  iou_scores = outputs.iou_scores.cpu()[0]
 
72
  if predicted_masks.ndim == 5:
73
  predicted_masks = predicted_masks.squeeze(1)
74
 
 
 
 
 
 
 
 
 
75
  # Select the best mask
76
  best_mask_idx = torch.argmax(iou_scores)
77
  best_mask = predicted_masks[0, best_mask_idx, :, :]
78
 
79
+ # Convert to binary mask (no resizing needed since processor doesn't resize)
80
+ mask_binary = (best_mask > 0.0).numpy().astype(np.uint8) * 255
 
 
 
 
 
 
 
 
81
 
82
  except Exception as e:
83
  print("Error processing masks: {}".format(e))
84
+ # Fallback: create a simple mask
85
+ height, width = img.size[1], img.size[0]
86
  mask_binary = np.zeros((height, width), dtype=np.uint8)
87
  center_x, center_y = width // 2, height // 2
88
  size = min(width, height) // 8