Update README.md

README.md

@@ -59,14 +59,32 @@ Example usage:
 import cv2
 from transformers import AutoModel
 
+def crop_mammo(img, crop_model, device):
+  img_shape = torch.tensor([img.shape[:2]])
+  x = crop_model.preprocess(img)
+  x = torch.from_numpy(x).expand(1, 1, -1, -1).float().to(device)
+  with torch.inference_mode():
+    coords = model(x, img_shape)
+  coords = coords[0].numpy()
+  x, y, w, h = coords
+  return img[y: y + h, x: x + w]
+
+device = "cuda:0"
+
+crop_model = AutoModel.from_pretrained("ianpan/mammo-crop", trust_remote_code=True)
+crop_model = crop_model.eval().to(device)
+
 model = AutoModel.from_pretrained("ianpan/mammoscreen", trust_remote_code=True)
-model = model.eval().to("cuda:0")
+model = model.eval().to(device)
 
 cc_img = cv2.imread("mammo_cc.png", cv2.IMREAD_GRAYSCALE)
 mlo_img = cv2.imread("mammo_mlo.png", cv2.IMREAD_GRAYSCALE)
 
+cc_img = crop_mammo(cc_img, crop_model, device)
+mlo_img = crop_mammo(mlo_img, crop_model, device)
+
 with torch.inference_mode():
-  output = model({"cc": cc_img, "mlo": mlo_img}, device="cuda:0")
+  output = model({"cc": cc_img, "mlo": mlo_img}, device=device)
 ```
 
 Note that the model preprocesses the data within the `forward` function into the necessary format.
@@ -75,7 +93,7 @@ If you want the predicted density class, take the argmax: `output['density'].arg
 
 You can also access each neural net separately using `model.net{i}`. However, you must apply the preprocessing outside of the `forward` function.
 ```
-input_dict = model.net0.preprocess({"cc": cc_img, "mlo": mlo_img}, device="cuda:0")
+input_dict = model.net0.preprocess({"cc": cc_img, "mlo": mlo_img}, device=device)
 with torch.inference_mode():
   out = model.net0(input_dict) 
 ```
@@ -90,7 +108,10 @@ mlo_images = ["rt_pt1_mlo.png", lt_pt1_mlo.png", "rt_pt2_mlo.png", "lt_pt2_mlo.p
 cc_images = [cv2.imread(_, cv2.IMREAD_GRAYSCALE) for _ in cc_images]
 mlo_images = [cv2.imread(_, cv2.IMREAD_GRAYSCALE) for _ in mlo_images]
 
+cc_images = [crop_mammo(_, crop_model, device) for _ in cc_images]
+mlo_images = [crop_mammo(_, crop_model, device), for _ in mlo_images]
+
 input_dict = [{"cc": cc_img, "mlo": mlo_img} for cc_img, mlo_img in zip(cc_images, mlo_images)]
 with torch.inference_mode():
-  output = model(input_dict, device="cuda:0")
+  output = model(input_dict, device=device)
 ```