replace preprocess specific to medsam

utils/utils.py

@@ -1,23 +1,46 @@
+import types
+
 import numpy as np
 import streamlit as st
 import torch
 from distinctipy import distinctipy
 from segment_anything import (SamAutomaticMaskGenerator, SamPredictor,
                               sam_model_registry)
+from torch.nn import functional as F
 
 
 def get_color():
     return distinctipy.get_colors(200)
 
 
+def medsam_preprocess(self, x: torch.Tensor) -> torch.Tensor:
+    """Normalize pixel values and pad to a square input."""
+    # Normalize colors
+    x = (x - x.min()) / torch.clip(
+        x.max() - x.min(), min=1e-8, max=None)  # normalize to [0, 1], (H, W, 3)
+
+    # Pad
+    h, w = x.shape[-2:]
+    padh = self.image_encoder.img_size - h
+    padw = self.image_encoder.img_size - w
+    x = F.pad(x, (0, padw, 0, padh))
+    return x
+
+
 @st.cache_resource
 def get_model(checkpoint='checkpoint/sam_vit_b_01ec64.pth'):
     device         = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    build_sam      = sam_model_registry['vit_b']
-    model          = build_sam(checkpoint=checkpoint)
+    model          = sam_model_registry['vit_b'](checkpoint=checkpoint)
+
+    # Replace preprocess function
+    funcType = types.MethodType
+    model.preprocess = funcType(medsam_preprocess, model)
+    model.mask_threshold = 0.5
+
     model = model.to(device)
     if torch.cuda.is_available():
         torch.cuda.empty_cache()
+
     predictor      = SamPredictor(model)
     mask_generator = SamAutomaticMaskGenerator(model)
     return predictor, mask_generator