CodeJackR
commited on
Commit
·
b0c81ad
1
Parent(s):
0e71822
Remove image resizing
Browse files- 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 |
-
#
|
91 |
-
|
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
|