Mdai deploy

.mdai/helper.py

@@ -4,48 +4,53 @@
 import util
 import torch
 from io import BytesIO
-from ct.ct_pe_constants import W_CENTER_DEFAULT, W_WIDTH_DEFAULT
-
-
-def format_img(img, input_slice_number, normalize_images=False):
-    num_slices = img.shape[0]
-    num_windows = num_slices - input_slice_number + 1
-
-    # crop
-    row = (img.shape[-2] - 192) // 2
-    col = (img.shape[-1] - 192) // 2
-    img = img[:, row : row + 192, col : col + 192]
-    img = img.astype(np.float32)
-
-    # noramlize Hounsfield Units
-    if normalize_images:
-        img = util.normalize(img)
-
-    # expand dimention for tensor
-    img_split = np.array([img[i : i + input_slice_number] for i in range(num_windows)])
-    img_expand = [
-        np.expand_dims(np.expand_dims(split, axis=0), axis=0) for split in img_split
-    ]
-    return img_expand
+from ct.ct_pe_constants import (
+    W_CENTER_DEFAULT,
+    W_WIDTH_DEFAULT,
+    CONTRAST_HU_MEAN,
+    CONTRAST_HU_MIN,
+    CONTRAST_HU_MAX,
+)
 
 
 def resize_array(np_array, slice_shape=(208, 208), interpolation=cv2.INTER_AREA):
     return cv2.resize(np_array, slice_shape, interpolation=interpolation)
 
 
-def preprocess(x_arrays, input_slice_number):
+def preprocess(x_arrays):
     x_arrays = sorted(x_arrays, key=lambda dcm: int(dcm.ImagePositionPatient[-1]))
-    x_arrays = [resize_array(dcm.pixel_array) for dcm in x_arrays]
+    x_arrays = np.array([resize_array(dcm.pixel_array) for dcm in x_arrays])
+    # rescale
+    interpolation = cv2.INTER_AREA
+    x_arrays = util.resize_slice_wise(x_arrays, (208, 208), interpolation)
+
+    # crop
+    row = (x_arrays.shape[-2] - 192) // 2
+    col = (x_arrays.shape[-1] - 192) // 2
+    x_arrays = x_arrays[:, row : row + 192, col : col + 192]
+    return x_arrays
 
-    x_stacked = np.stack(x_arrays, 0)
-    x_stacked = util.apply_window(x_stacked, W_CENTER_DEFAULT, W_WIDTH_DEFAULT)
 
-    x = format_img(x_stacked, input_slice_number, normalize_images=True)
-    x = [torch.from_numpy(np.array(window)) for window in x]
+def get_windows(x_stacked, input_slice_number):
+    num_slices = x_stacked.shape[0]
+    num_windows = num_slices - input_slice_number + 1
 
-    x_un_normalized = format_img(x_stacked, input_slice_number, normalize_images=False)
+    x_stacked = util.normalize(x_stacked)
+    for i in range(num_windows):
+        img_split = np.array(x_stacked[i : i + input_slice_number])
+        img_expand = np.expand_dims(np.expand_dims(img_split, axis=0), axis=0)
+        yield torch.from_numpy(np.array(img_expand))
 
-    return x, x_un_normalized
+
+def get_best_window(x_stacked, input_slice_number, best_window):
+
+    x_unnorm_best = np.array(x_stacked[best_window : best_window + input_slice_number])
+    # noramlize Hounsfield Units
+    x_stacked = util.normalize(x_stacked)
+    x_best = np.array(x_stacked[best_window : best_window + input_slice_number])
+    x_unnorm_best = np.expand_dims(np.expand_dims(x_unnorm_best, axis=0), axis=0)
+    x_best = np.expand_dims(np.expand_dims(x_best, axis=0), axis=0)
+    return torch.from_numpy(np.array(x_best)), x_unnorm_best
 
 
 def compute_gradcam_gif(cam, x, x_un_normalized):
@@ -62,6 +67,8 @@ def compute_gradcam_gif(cam, x, x_un_normalized):
         save_all=True,
         append_images=cam_frames[1:] if len(cam_frames) > 1 else [],
         format="GIF",
+        loop=0,
+        optimize=True,
     )
 
     return gradcam_output_buffer

.mdai/mdai_deploy.py

@@ -6,12 +6,11 @@
 from custom_gradcam import CustomGradCAM
 from saver import ModelSaver
 import util
-from helper import preprocess, compute_gradcam_gif
+from helper import compute_gradcam_gif, preprocess, get_best_window, get_windows
 
 DEFAULT_PROBABILITY_THRESHOLD = "0.5"
 DEFAULT_INPUT_SLICE_NUMBER = "24"
-GRADCAM_OFF = "0"
-GRADCAM_MAX_PROBABILITY = "1"
+GRADCAM = "0"
 
 
 class MDAIModel:
@@ -100,7 +99,11 @@ def predict(self, data):
         x_arrays = []
         for instance in input_instances:
             tags = instance["tags"]
-            ds = pydicom.dcmread(BytesIO(instance["file"]))
+            try:
+                ds = pydicom.dcmread(BytesIO(instance["file"]))
+                arr = ds.pixel_array
+            except:
+                continue
             x_orig = ds
             x_arrays.append(x_orig)
 
@@ -111,30 +114,31 @@ def predict(self, data):
         if len(x_arrays) < input_slice_number:
             input_slice_number = len(x_arrays)
 
-        x, x_un_normalized = preprocess(x_arrays, input_slice_number)
+        x_stacked = preprocess(x_arrays)
         self.model.eval()
 
         best_window = 0
         probability = 0.0
+        i = 0
         with torch.no_grad():
-            for i, window in enumerate(x):
+            for window in get_windows(x_stacked, input_slice_number):
                 cls_logits = self.model.forward(
                     window.to(self.device, dtype=torch.float)
                 )
                 cls_probs = torch.sigmoid(cls_logits).to("cpu").numpy()
                 if cls_probs[0][0] > probability:
                     probability = cls_probs[0][0]
                     best_window = i
-
+                i += 1
         if not probability >= float(
             input_args.get("probability_threshold", DEFAULT_PROBABILITY_THRESHOLD)
         ):
             result = {
                 "type": "NONE",
                 "study_uid": tags["StudyInstanceUID"],
                 "series_uid": tags["SeriesInstanceUID"],
-                "instance_uid": tags["SOPInstanceUID"],
                 "frame_number": None,
+                "probability": float(probability),
             }
         else:
             result = {
@@ -147,23 +151,23 @@ def predict(self, data):
                 "probability": float(probability),
             }
 
-            if input_args.get("gradcam", GRADCAM_OFF) == GRADCAM_MAX_PROBABILITY:
+            if input_args.get("gradcam", GRADCAM) == "1":
+                x_best, x_unnorm_best = get_best_window(
+                    x_stacked, input_slice_number, best_window
+                )
                 self.grad_cam.register_hooks()
 
                 with torch.set_grad_enabled(True):
-                    probs, idx = self.grad_cam.forward(x[best_window])
+                    probs, idx = self.grad_cam.forward(x_best)
                     self.grad_cam.backward(idx=idx[0])
                     cam = self.grad_cam.get_cam("module.encoders.3")
 
                 self.grad_cam.remove_hooks()
 
-                gradcam_output_buffer = compute_gradcam_gif(
-                    cam, x[best_window], x_un_normalized[best_window]
-                )
+                gradcam_output_buffer = compute_gradcam_gif(cam, x_best, x_unnorm_best)
                 gradcam_explanation = [
                     {
                         "name": "Grad-CAM",
-                        "description": "Visualize how parts of the image affects neural network’s output by looking into the activation maps. From _Grad-CAM: Visual Explanations from Deep Networks via Gradient-based Localization_ (https://arxiv.org/abs/1610.02391)",
                         "content": gradcam_output_buffer.getvalue(),
                         "content_type": "image/gif",
                     }