Diff of /predict.py [000000] .. [dce3d9]

Switch to side-by-side view

--- a
+++ b/predict.py
@@ -0,0 +1,68 @@
+import os
+import cv2
+import numpy as np
+import pandas as pd
+from glob import glob
+from tqdm import tqdm
+import pydicom as dicom
+
+import tensorflow as tf
+from tensorflow.keras.utils import CustomObjectScope
+from sklearn.metrics import (
+    accuracy_score,
+    f1_score,
+    jaccard_score,
+    precision_score,
+    recall_score,
+)
+from metrics import dice_loss, dice_coef, iou
+
+""" Creating a directory """
+
+
+def create_dir(path):
+    if not os.path.exists(path):
+        os.makedirs(path)
+
+
+if __name__ == "__main__":
+    """Seeding"""
+    SEEDS = 42
+    np.random.seed(SEEDS)
+    tf.random.set_seed(SEEDS)
+
+    """ Directory for storing files """
+    create_dir("test")
+
+    """ Loading model """
+    with CustomObjectScope(
+        {"iou": iou, "dice_coef": dice_coef, "dice_loss": dice_loss}
+    ):
+        model = tf.keras.models.load_model("files/model.h5")
+
+    """ Load the dataset """
+    test_x = glob("data/test/*/*/*.dcm")
+    print(f"Test: {len(test_x)}")
+
+    """ Loop over the data """
+    for x in tqdm(test_x):
+        """Extract the Names"""
+        dir_name = x.split("/")[-3]
+        name = dir_name + "_" + x.split("/")[-1].split(".")[0]
+
+        """ Read the Image """
+        image = dicom.dcmread(x).pixel_array
+        image = np.expand_dims(image, axis=-1)
+        image = image / np.max(image) * 255.0
+        x = image / 255.0
+        x = np.concatenate([x, x, x], axis=-1)
+        x = np.expand_dims(x, axis=0)
+
+        """ Prediction """
+        mask = model.predict(x)[0]
+        mask = mask > 0.5
+        mask = mask.astype(np.int32)
+        mask = mask * 255
+
+        cat_images = np.concatenate([image, mask], axis=1)
+        cv2.imwrite(f"test/{name}.png", cat_images)