Manage image resizing

handler.py

@@ -64,24 +64,21 @@ class EndpointHandler():
         
         # 4. Process and select the best mask
         try:
-            original_height, original_width = inputs["original_sizes"][0].tolist()
-            pred_masks = outputs.pred_masks.cpu()
+            # Use the processor's post-processing utility to resize masks and remove padding
+            masks = self.processor.post_process_masks(
+                outputs.pred_masks.cpu(), 
+                inputs["original_sizes"].cpu(), 
+                inputs["reshaped_input_sizes"].cpu()
+            )[0]
+            
+            # The output of post_process_masks is a tensor of shape (num_masks, H, W)
+            # where H and W are the original image dimensions.
             iou_scores = outputs.iou_scores.cpu()[0]
-
-            if pred_masks.ndim == 5:
-                pred_masks = pred_masks.squeeze(1)
-
             best_mask_idx = torch.argmax(iou_scores)
-            best_mask_tensor = pred_masks[0, best_mask_idx, :, :]
-
-            upscaled_mask = F.interpolate(
-                best_mask_tensor.unsqueeze(0).unsqueeze(0).float(),
-                size=(original_height, original_width),
-                mode='bilinear',
-                align_corners=False
-            ).squeeze()
+            best_mask_tensor = masks[best_mask_idx, :, :]
 
-            mask_binary = (upscaled_mask > 0.0).numpy().astype(np.uint8) * 255
+            # Convert to binary mask (float tensor to uint8 numpy array)
+            mask_binary = (best_mask_tensor > 0.0).numpy().astype(np.uint8) * 255
         
         except Exception as e:
             print("Error processing masks: {}".format(e))
@@ -107,11 +104,15 @@ def main():
 
     # 2. Instantiate handler and get the PIL Image result
     handler = EndpointHandler(path=".")
-    result_img = handler(payload)
+    result = handler(payload)
 
-    # 3. Save the resulting image
-    result_img.save(output_path)
-    print("Wrote mask to {}".format(output_path))
+    # 3. Extract the image from the result and save it
+    if result and isinstance(result, list) and 'mask' in result[0]:
+        result_img = result[0]['mask']
+        result_img.save(output_path)
+        print("Wrote mask to {}".format(output_path))
+    else:
+        print("Failed to get a valid mask from the handler.")
 
 if __name__ == "__main__":
     main()