Initial commit

README.md

@@ -1 +1,23 @@
-# Drag-Drop
+# Acquiring Efficient Annotations for Tumor Detection and Localization in Temporal and Volumetric Medical Images 
+
+
+
+1. Configuring your environment (Prerequisites):
+
+    + Creating a virtual environment in terminal:       
+
+    ```bash
+    conda create -n DragDrop python=3.7
+    ```
+
+    + Installing necessary packages: 
+
+    ```bash
+    pip install -r requirements.txt
+    ```
+
+2. Usage:
+
+    ```bash
+    python main.py
+    ```

main.py

@@ -0,0 +1,79 @@
+import os
+import glob
+import argparse
+import numpy as np
+from scipy import ndimage
+from skimage.segmentation import watershed
+
+from utils.drag_drop import get_weak_label, get_marker, dilate_lesion
+from utils.utils import load_img_label, save_output, cal_img_gradient
+
+
+def watershed_core(input, label, N, M):
+
+    pseudo_label = np.zeros_like(label, dtype=np.uint8)
+    label_numeric, gt_N = ndimage.label(label)
+
+    for segid in range(1, gt_N+1):
+
+        label_binary = np.uint8(label_numeric == segid)
+        marker = np.zeros_like(label_binary)
+
+        # ******************** Sec. 2.1 Drag&Drop Initialization ********************
+        center_x, center_y, sphere_radius = get_weak_label(label_binary)
+        marker = get_marker(center_x, center_y, sphere_radius, N, marker)
+
+        # ******************** Sec. 2.2 Drag&Drop Propagation ********************
+        lesion = watershed(input, marker)-1
+
+        # ******************** Sec. 2.3 Noise Reduction ******************** 
+        dilated_lesion = dilate_lesion(lesion, sphere_radius, M)
+
+        pseudo_label = pseudo_label | dilated_lesion
+
+    pseudo_label[pseudo_label==3] = 1
+
+    return pseudo_label
+
+
+def watershed_engine(opt):
+
+    imgpath_list =  glob.glob(opt.image_root + '*.jpg')
+
+    for imgpath in imgpath_list:
+
+        labelpath = imgpath.replace(opt.image_root, opt.label_root).replace('.jpg', '.png')
+        tar_path = imgpath.replace(opt.image_root, opt.output_root).replace('.jpg', '.png')
+
+        grayimage, label, image= load_img_label(imgpath, labelpath)
+        gradient = cal_img_gradient(grayimage)
+
+        pseudo_label = watershed_core(gradient, label, opt.N, opt.M)
+        save_output(tar_path, pseudo_label, image)
+
+
+
+if __name__ == '__main__':
+    np.random.seed(1000)
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        '--image_root', type=str, help='custom your image root',
+        default='data/image/')
+    parser.add_argument(
+        '--label_root', type=str, help='custom your ground-truth root',
+        default='data/label/')
+    parser.add_argument(
+        '--output_root', type=str, help='custom your prediction root',
+        default='data/drag_drop/')
+
+    # ******************** Hyper-parameters ********************
+    parser.add_argument(
+        '--N', type=float, help='the kernel size of dilation',
+        default=0.2)
+    parser.add_argument(
+        '--M', type=float, help='the kernel size of dilation',
+        default=0.5)
+    opt = parser.parse_args()
+
+    os.makedirs(opt.output_root, exist_ok=True)
+    watershed_engine(opt)

requirements.txt

@@ -0,0 +1,4 @@
+numpy==1.21.6
+scipy==1.7.3
+scikit-image==0.19.3
+nibabel==4.0.2

utils/drag_drop.py

@@ -0,0 +1,47 @@
+import numpy as np
+from scipy import ndimage
+
+from utils.utils import sample_spherical
+
+
+def get_weak_label(label_binary):
+
+    lesion_index = np.where(label_binary==1)
+    center_x, center_y = (np.max(lesion_index[0])+np.min(lesion_index[0]))//2, (np.max(lesion_index[1])+np.min(lesion_index[1]))//2
+
+    width, height = np.max(lesion_index[0])-np.min(lesion_index[0]), np.max(lesion_index[1])-np.min(lesion_index[1])
+
+    sphere_radius = int(max(width, height)//2)
+
+    return center_x, center_y, sphere_radius
+
+
+def get_marker(center_x, center_y, sphere_radius, N, marker):
+
+    pos_marker_radius = sphere_radius*N
+
+    # ******************** Positive Marker ********************
+    noisex, noisey = np.int32(np.random.normal(0, 1, 1)), np.int32(np.random.normal(0, 1, 1))
+    marker[center_x+noisex[0]-round(pos_marker_radius): center_x+noisex[0]+round(pos_marker_radius)+1,
+        center_y+noisey[0]-round(pos_marker_radius): center_y+noisey[0]+round(pos_marker_radius)+1] = 2
+
+    # ******************** Negative Marker ********************
+    neg_marker_index = sample_spherical(center_x, center_y, sphere_radius, marker.shape)
+
+    for index in range(len(neg_marker_index[0])):
+        marker[neg_marker_index[0][index]-2: neg_marker_index[0][index]+2+1, 
+                neg_marker_index[1][index]-2: neg_marker_index[1][index]+2+1] = 1
+
+
+    return marker
+
+def dilate_lesion(lesion, sphere_radius, M):
+
+    dilate_size = sphere_radius*M
+    dilate_size = max(dilate_size, 1)
+
+    dilated_lesion = ndimage.grey_dilation(lesion, footprint=np.ones((int(2*dilate_size+1), int(2*dilate_size+1))).astype("uint8"))
+    dilated_lesion = np.uint8(dilated_lesion*2) - np.uint8(lesion)
+
+    return dilated_lesion
+

utils/utils.py

@@ -0,0 +1,60 @@
+import cv2
+import numpy as np
+from skimage.filters import  rank
+
+
+def load_img_label(imgpath, labelpath):
+
+    label = cv2.imread(labelpath)
+    label = cv2.cvtColor(label, cv2.COLOR_BGR2GRAY) // 255
+
+    image = cv2.imread(imgpath)
+    grayimage = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
+
+    return grayimage, label, image
+
+
+def cal_img_gradient(img):
+
+    img = rank.median(img, np.ones((2, 2)))
+    gradient_volume = rank.gradient(img, np.ones((2, 2)))
+
+    return gradient_volume
+
+
+def save_output(file_path, pred, img, pred_color = [244, 133, 0], mask_color = [255, 179, 255]):
+
+    pred_gt = cv2.cvtColor(np.uint8(pred == 1),cv2.COLOR_GRAY2RGB)
+    pred_mask = cv2.cvtColor(np.uint8(pred == 2),cv2.COLOR_GRAY2RGB)
+    pred = pred_gt * pred_color + pred_mask * mask_color
+
+    img_pred = cv2.addWeighted(np.uint8(img), 0.9, np.uint8(pred), 0.6, 1)
+
+    cv2.imwrite(file_path, img_pred)
+
+
+def count_integer_points_on_sphere(radius):
+
+    x, y= np.meshgrid(range(-radius, radius + 1), range(-radius, radius + 1))
+    distances_squared = x**2 + y**2
+    count = np.count_nonzero(abs(distances_squared - radius**2) <= 1)
+
+    return count
+
+
+def sample_spherical(cx, cy, radius, shape):
+
+    num_points = count_integer_points_on_sphere(radius)
+
+    t = np.linspace(0, 2*np.pi, num_points, endpoint=False)
+    x = np.int32(cx + radius * np.cos(t))
+    y = np.int32(cy + radius * np.sin(t))
+
+    spherical_index = np.array([x,y])
+
+    spherical_index[0][spherical_index[0]+2+1>=shape[0]] = shape[0]-1-2
+    spherical_index[1][spherical_index[1]+2+1>=shape[1]] = shape[1]-1-2
+    spherical_index[spherical_index<0] = 0
+
+    return spherical_index
+