use huggingface models

app.py

@@ -7,20 +7,20 @@ import torch
 import torch.nn as nn
 
 from einops import rearrange
-from importlib import import_module
 from pytorch_grad_cam import GradCAM
 from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
 from skimage.exposure import match_histograms
-from skp.utils import load_model_from_config, load_kfold_ensemble_as_list
+from transformers import AutoModel
 
 
 class ModelForGradCAM(nn.Module):
-    def __init__(self, model):
+    def __init__(self, model, female):
         super().__init__()
         self.model = model
+        self.female = female
 
     def forward(self, x):
-        return self.model({"x": x})["logits1"]
+        return self.model(x, self.female, return_logits=True)
 
 
 def convert_bone_age_to_string(bone_age: float):
@@ -47,67 +47,29 @@ def convert_bone_age_to_string(bone_age: float):
     return str_output
 
 
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-cfg_crop = import_module("skp.configs.boneage.cfg_crop_simple_resize").cfg
-crop_model = load_model_from_config(
-    cfg_crop, weights_path="crop.pt", device=device, eval_mode=True
-)
-
-cfg = import_module("skp.configs.boneage.cfg_female_channel_reg_cls_match_hist").cfg
-cfg.backbone = "convnextv2_tiny"
-
-model_list = load_kfold_ensemble_as_list(
-    cfg, [f"net{i}.pt" for i in range(3)], device=device, eval_mode=True
-)
-
-ref_img = rearrange(cv2.imread("ref_img.png", 0), "h w -> h w 1 ")
-
-with open("greulich_and_pyle_ages.json", "r") as f:
-    greulich_and_pyle_ages = json.load(f)["bone_ages"]
-
-greulich_and_pyle_ages = {k: np.asarray(v) for k, v in greulich_and_pyle_ages.items()}
-
-model_grad_cam = ModelForGradCAM(model_list[0])
-target_layers = [model_grad_cam.model.backbone.stages[-1]]
-
-
 @spaces.GPU
 def predict_bone_age(Radiograph, Sex, Heatmap):
-    x0 = rearrange(Radiograph, "h w -> h w 1")
-    x = cfg_crop.val_transforms(image=x0)["image"]
-    x = torch.from_numpy(x)
-    x = rearrange(x, "h w c -> 1 c h w")
+    x = crop_model.preprocess(Radiograph)
+    x = torch.from_numpy(x).float().to(device)
+    x = rearrange(x, "h w -> 1 1 h w")
     # crop
+    img_shape = torch.tensor([Radiograph.shape[:2]]).to(device)
     with torch.inference_mode():
-        box = crop_model({"x": x.to(device).float()}, return_loss=False)["logits"][
-            0
-        ].cpu()
-        box[[0, 2]] = box[[0, 2]] * x0.shape[1]
-        box[[1, 3]] = box[[1, 3]] * x0.shape[0]
-    box = box.numpy().astype("int")
-    x, y, w, h = box
-    x0 = x0[y : y + h, x : x + w]
+        box = crop_model(x, img_shape=img_shape).to("cpu").numpy()
+    x, y, w, h = box[0]
+    cropped = Radiograph[y : y + h, x : x + w]
     # histogram matching
-    x0 = match_histograms(x0, ref_img)
-    x = cfg.val_transforms(image=x0)["image"]
-    # create image channel for female/male
-    ch = np.zeros_like(x)
-    if Sex:  # 0- male, 1- female
-        ch[...] = 255
-    x = np.concatenate([x, ch], axis=-1)
-    x = torch.from_numpy(x)
-    x = rearrange(x, "h w c -> 1 c h w")
+    x = match_histograms(cropped, ref_img)
+
+    x = model.preprocess(x)
+    x = torch.from_numpy(x).float().to(device)
+    x = rearrange(x, "h w -> 1 1 h w")
+    female = torch.tensor([Sex]).to(device)
     with torch.inference_mode():
-        bone_age = []
-        for each_model in model_list:
-            pred = each_model({"x": x.to(device).float()}, return_loss=False)[
-                "logits1"
-            ][0].cpu()
-            pred = (pred.softmax(0) * torch.arange(240)).sum().numpy()
-            bone_age.append(pred)
-        bone_age = np.mean(bone_age)
+        bone_age = model(x, female)[0].item()
 
+    # get closest G&P ages
+    # from: https://rad.esmil.com/Reference/G_P_BoneAge/
     gp_ages = greulich_and_pyle_ages["female" if Sex else "male"]
     diffs_gp = np.abs(bone_age - gp_ages)
     diffs_gp = np.argsort(diffs_gp)
@@ -119,29 +81,33 @@ def predict_bone_age(Radiograph, Sex, Heatmap):
     closest2 = convert_bone_age_to_string(closest2)
 
     if Heatmap:
+        # net1 and net2 to give good GradCAMs
+        # net0 is bad for some reason
+        # because GradCAM expects 1 input tensor, need to
+        # pass female during class instantiation
+        model_grad_cam = ModelForGradCAM(model.net1, female)
+        target_layers = [model_grad_cam.model.backbone.stages[-1]]
         targets = [ClassifierOutputTarget(round(bone_age))]
         with GradCAM(model=model_grad_cam, target_layers=target_layers) as cam:
-            grayscale_cam = cam(
-                input_tensor=x.to(device).float(), targets=targets, eigen_smooth=True
-            )
+            grayscale_cam = cam(input_tensor=x, targets=targets, eigen_smooth=True)
 
         heatmap = cv2.applyColorMap(
             (grayscale_cam[0] * 255).astype("uint8"), cv2.COLORMAP_JET
         )
-        image = cv2.cvtColor(x[0, 0].cpu().numpy().astype("uint8"), cv2.COLOR_GRAY2RGB)
+        image = cv2.cvtColor(
+            x[0, 0].to("cpu").numpy().astype("uint8"), cv2.COLOR_GRAY2RGB
+        )
         image_weight = 0.6
         grad_cam_image = (1 - image_weight) * heatmap[..., ::-1] + image_weight * image
-        grad_cam_image = grad_cam_image.astype("uint8")
+        grad_cam_image = grad_cam_image
     else:
         # if no heatmap desired, just show image
-        grad_cam_image = cv2.cvtColor(
-            x[0, 0].cpu().numpy().astype("uint8"), cv2.COLOR_GRAY2RGB
-        )
+        grad_cam_image = cv2.cvtColor(x[0, 0].to("cpu").numpy(), cv2.COLOR_GRAY2RGB)
 
     return (
         bone_age_str,
         f"The closest Greulich & Pyle bone ages are:\n  1) {closest1}\n  2) {closest2}",
-        grad_cam_image,
+        grad_cam_image.astype("uint8"),
     )
 
 
@@ -157,11 +123,8 @@ with gr.Blocks() as demo:
         """
     # Deep Learning Model for Pediatric Bone Age
 
-    This model predicts the bone age from a single frontal view hand radiograph. 
-    The model was trained on the publicly available 
-    [RSNA Pediatric Bone Age Challenge](https://www.rsna.org/rsnai/ai-image-challenge/rsna-pediatric-bone-age-challenge-2017) dataset. 
-    The model achieves a mean absolute error of 4.26 months on the original test set comprising 200 multi-annotated hand radiographs, 
-    which is competitive with [top solutions](https://pubs.rsna.org/doi/10.1148/radiol.2018180736) from the original challenge.
+    This model predicts the bone age from a single frontal view hand radiograph. Read more about the model here:
+    <https://huggingface.co/ianpan/bone-age>
 
     There is also an option to output a heatmap over the radiograph to show regions where the model is focusing on
     to make its prediction. However, this takes extra computation and will increase the runtime.
@@ -172,7 +135,7 @@ with gr.Blocks() as demo:
 
     Created by: Ian Pan, <https://ianpan.me>
     
-    Last updated: December 15, 2024
+    Last updated: December 16, 2024
     """
     )
     gr.Interface(
@@ -184,8 +147,27 @@ with gr.Blocks() as demo:
             ["examples/10043.png", "Female", "No"],
             ["examples/8888.png", "Female", "Yes"],
         ],
-        cache_examples=False,
+        cache_examples="lazy",
     )
 
 if __name__ == "__main__":
+    device = "cuda" if torch.cuda.is_available() else "cpu"
+    print(f"Using device `{device}` ...")
+
+    crop_model = AutoModel.from_pretrained(
+        "ianpan/bone-age-crop", trust_remote_code=True
+    )
+    model = AutoModel.from_pretrained("ianpan/bone-age", trust_remote_code=True)
+
+    crop_model, model = crop_model.eval().to(device), model.eval().to(device)
+
+    ref_img = cv2.imread("ref_img.png", 0)
+
+    with open("greulich_and_pyle_ages.json", "r") as f:
+        greulich_and_pyle_ages = json.load(f)["bone_ages"]
+
+    greulich_and_pyle_ages = {
+        k: np.asarray(v) for k, v in greulich_and_pyle_ages.items()
+    }
+
     demo.launch(share=True)

requirements.txt

@@ -5,4 +5,5 @@ gradio
 scikit-image
 spaces
 timm
-torch
+torch
+transformers