Add fracture inference

app.py

@@ -27,6 +27,24 @@ def rescale(x):
     return x
 
 
+def get_cervical_spine_coordinates(x, original_shape):
+    # Assumes x is torch tensor with shape (1, 8, Z, H, W)
+    x = x.squeeze(0).numpy()[:7]
+    rescale_factor = [original_shape[0] / x.shape[1], original_shape[1] / x.shape[2], original_shape[2] / x.shape[3]]
+    coords_dict = {}
+    for level in range(x.shape[0]):
+        coords = np.where(x[level] >= 0.4)
+        coords = np.vstack(coords).astype("float")
+        coords[0] = coords[0] * rescale_factor[0]
+        coords[1] = coords[1] * rescale_factor[1]
+        coords[2] = coords[2] * rescale_factor[2]
+        coords = coords.astype("int")
+        coords_dict[level] = coords[0].min(), coords[0].max(),\
+                             coords[1].min(), coords[1].max(),\
+                             coords[2].min(), coords[2].max()
+    return coords_dict
+
+
 def generate_segmentation_video(study):
     img = nibabel.load(study).get_fdata()[:, ::-1, ::-1].transpose(2, 1, 0)
     img = window(img)
@@ -38,13 +56,34 @@ def generate_segmentation_video(study):
         seg_output = seg_model(X)
 
     seg_output = torch.sigmoid(seg_output)
-    p_spine = seg_output[:, :7].sum(1)
-    seg_output = torch.argmax(seg_output, dim=1) + 1
+    c_spine_coords = get_cervical_spine_coordinates(seg_output, img.shape)
+
+    chunk_features = []
+    for level, coords in c_spine_coords.items():
+        z1, z2, h1, h2, w1, w2 = coords
+        X = torch.from_numpy(img[z1:z2+1, h1:h2+1, w1:w2+1]).float().unsqueeze(0).unsqueeze(0)
+        X = F.interpolate(X, size=(64, 288, 288), mode="nearest")
+        X = rescale(X)
+        with torch.no_grad():
+            chunk_features.append(x3d_model.extract_features(X))
+
+    chunk_features = torch.stack(chunk_features, dim=1)
+    with torch.no_grad():
+        final_output = torch.sigmoid(seq_model((chunk_features, torch.ones((chunk_features.size(1), )))))
+
+    final_output_dict = {f"C{i+1}": final_output[:, i].item() for i in range(7)}
+    final_output_dict["Overall"] = final_output[:, -1].item()
+
+    seg_output = F.interpolate(seg_output, size=img.shape, mode="nearest").squeeze(0).numpy()
+    # shape = (8, Z, H, W)
+    p_spine = seg_output[:7].sum(0)
+    # shape = (Z, H, W)
+    seg_output = np.argmax(seg_output, dim=0) + 1
+    # shape = (Z, H, W)
     seg_output[p_spine < 0.5] = 0
-    seg_output = F.interpolate(seg_output.unsqueeze(0).float(), size=img.shape, mode="nearest")
-    seg_output = seg_output.squeeze(0).squeeze(0).numpy()
     seg_output = (seg_output * 255 / 7).astype("uint8")
     seg_output = np.stack([cv2.applyColorMap(_, cv2.COLORMAP_JET) for _ in seg_output])
+    seg_output[p_spine < 0.5] = 0
 
     frames = []
     skip = 8
@@ -54,8 +93,8 @@ def generate_segmentation_video(study):
         i = cv2.cvtColor(i, cv2.COLOR_GRAY2RGB)
         frame = np.concatenate((i, o), 1)
         frames.append(frame)
-    mediapy.write_video("video.mp4", frames, fps=30)
-    return "video.mp4"
+    mediapy.write_video("video.mp4", frames, fps=10)
+    return "video.mp4", final_output_dict
 
 
 ffmpeg_path = shutil.which('ffmpeg')
@@ -65,15 +104,26 @@ config = OmegaConf.load("configs/pseudoseg000.yaml")
 config.model.load_pretrained = "seg.ckpt"
 config.model.params.encoder_params.pretrained = False
 seg_model = builder.build_model(config).eval()
+
+config = OmegaConf.load("configs/chunk000.yaml")
+config.model.load_pretrained = "x3d.ckpt"
+config.model.params.pretrained = False
+x3d_model = builder.build_model(config).eval()
+
+config = OmegaConf.load("configs/chunkseq003.yaml")
+config.model.load_pretrained = "seq.ckpt"
+seq_model = builder.build_model(config).eval()
+
 examples = glob.glob("examples/*.nii.gz")
 
 with gr.Blocks(theme="dark-peach") as demo:
     select_study = gr.Dropdown(choices=sorted(examples), type="value", label="Select a study")
     button_predict = gr.Button("Predict")
-    video_output = gr.Video()
+    video_output = gr.Video(label="Cervical Spine Segmentation")
+    label_output = gr.Label(label="Fracture Predictions", show_label=False)
     button_predict.click(fn=generate_segmentation_video,
                          inputs=select_study,
-                         outputs=video_output)
+                         outputs=[video_output, label_output])
 
 
 if __name__ == "__main__":