Update README.md

README.md

@@ -48,12 +48,13 @@ Here is how to use this model:
 ```python
 
 from skimage import io, segmentation, morphology, measure, exposure
-from cell_sribd_model import MyModel
+from sribd_cellseg_models import MultiStreamCellSegModel,ModelConfig
 import numpy as np
 import tifffile as tif
 import requests
+import torch
 
-img_name = 'cell_00010.png'
+img_name = 'cell_00023.tiff'
 def normalize_channel(img, lower=1, upper=99):
     non_zero_vals = img[np.nonzero(img)]
     percentiles = np.percentile(non_zero_vals, [lower, upper])
@@ -78,12 +79,10 @@ for i in range(3):
     img_channel_i = img_data[:,:,i]
     if len(img_channel_i[np.nonzero(img_channel_i)])>0:
         pre_img_data[:,:,i] = normalize_channel(img_channel_i, lower=1, upper=99)
-
-
-#config = ModelConfig()
-#print(config)
-my_model = MyModel.from_pretrained("Lewislou/cellseg_sribd")
+#dummy_input = np.zeros((512,512,3)).astype(np.uint8)
+my_model = MultiStreamCellSegModel.from_pretrained("Lewislou/cellseg_sribd")
 checkpoints = torch.load('model.pt')
+my_model.__init__(ModelConfig())
 my_model.load_checkpoints(checkpoints)
 with torch.no_grad():
     output = my_model(pre_img_data)