Upload lrp_pipeline_2.py

lrp_pipeline_2.py

@@ -0,0 +1,417 @@
+import cv2
+import torch
+import torch.nn as nn
+import numpy as np
+import torchvision
+import os
+import copy
+from sklearn.mixture import GaussianMixture as GMM
+from sklearn.cluster import KMeans
+from simple_lama_inpainting import SimpleLama
+from PIL import Image
+from matplotlib.colors import ListedColormap
+import matplotlib.pyplot as plt
+import matplotlib
+import csv
+
+matplotlib.use("Agg")
+
+import base64
+
+from utils import (
+    select_sample_images,
+    create_cell_descriptors_table,
+    calculate_cell_descriptors,
+)
+
+preprocessed_folder = "uploads/"
+intermediate_folder = "heatmaps/"
+segmentation_folder = "segmentations/"
+tables_folder = "tables/"
+cell_descriptors_path = "cell_descriptors/cell_descriptors.csv"
+imgclasses = {0: "abnormal", 1: "normal"}
+
+
+def toconv(layers):
+    newlayers = []
+    for i, layer in enumerate(layers):
+        if isinstance(layer, nn.Linear):
+            newlayer = None
+            if i == 0:
+                m, n = 512, layer.weight.shape[0]
+                newlayer = nn.Conv2d(m, n, 4)
+                newlayer.weight = nn.Parameter(layer.weight.reshape(n, m, 4, 4))
+            else:
+                m, n = layer.weight.shape[1], layer.weight.shape[0]
+                newlayer = nn.Conv2d(m, n, 1)
+                newlayer.weight = nn.Parameter(layer.weight.reshape(n, m, 1, 1))
+            newlayer.bias = nn.Parameter(layer.bias)
+            newlayers += [newlayer]
+        else:
+            newlayers += [layer]
+    return newlayers
+
+
+def newlayer(layer, g):
+    layer = copy.deepcopy(layer)
+    try:
+        layer.weight = nn.Parameter(g(layer.weight))
+    except AttributeError:
+        pass
+    try:
+        layer.bias = nn.Parameter(g(layer.bias))
+    except AttributeError:
+        pass
+    return layer
+
+
+def heatmap(R, sx, sy, intermediate_path):
+    b = 10 * ((np.abs(R) ** 3.0).mean() ** (1.0 / 3))
+    my_cmap = plt.cm.seismic(np.arange(plt.cm.seismic.N))
+    my_cmap[:, 0:3] *= 0.85
+    my_cmap = ListedColormap(my_cmap)
+    plt.figure(figsize=(sx, sy))
+    plt.subplots_adjust(left=0, right=1, bottom=0, top=1)
+    plt.axis("off")
+    plt.imshow(R, cmap=my_cmap, vmin=-b, vmax=b, interpolation="nearest")
+    # plt.show()
+    plt.savefig(intermediate_path, bbox_inches="tight", pad_inches=0)
+    plt.close()
+
+
+def get_LRP_heatmap(image, L, layers, imgclasses, intermediate_path):
+    img = np.array(image)[..., ::-1] / 255.0
+    mean = torch.FloatTensor([0.485, 0.456, 0.406]).reshape(1, -1, 1, 1)  # torch.cuda
+    std = torch.FloatTensor([0.229, 0.224, 0.225]).reshape(1, -1, 1, 1)  # torch.cuda
+    X = (torch.FloatTensor(img[np.newaxis].transpose([0, 3, 1, 2]) * 1) - mean) / std
+
+    A = [X] + [None] * L
+    for l in range(L):
+        A[l + 1] = layers[l].forward(A[l])
+
+    scores = np.array(A[-1].cpu().data.view(-1))
+    ind = np.argsort(-scores)
+    for i in ind[:2]:
+        print("%20s (%3d): %6.3f" % (imgclasses[i], i, scores[i]))
+
+    T = torch.FloatTensor(
+        (1.0 * (np.arange(2) == ind[0]).reshape([1, 2, 1, 1]))
+    )  # SET FOR THE HIGHEST SCORE CLASS
+    R = [None] * L + [(A[-1] * T).data]
+    for l in range(1, L)[::-1]:
+        A[l] = (A[l].data).requires_grad_(True)
+        if isinstance(layers[l], torch.nn.MaxPool2d):
+            layers[l] = torch.nn.AvgPool2d(2)
+        if isinstance(layers[l], torch.nn.Conv2d) or isinstance(
+            layers[l], torch.nn.AvgPool2d
+        ):
+            rho = lambda p: p + 0.25 * p.clamp(min=0)
+            incr = lambda z: z + 1e-9  # USE ONLY THE GAMMA RULE FOR ALL LAYERS
+
+            z = incr(newlayer(layers[l], rho).forward(A[l]))  # step 1
+            # adding epsilon
+            epsilon = 1e-9
+            z_nonzero = torch.where(z == 0, torch.tensor(epsilon, device=z.device), z)
+            s = (R[l + 1] / z_nonzero).data
+            # s = (R[l+1]/z).data                                    # step 2
+            (z * s).sum().backward()
+            c = A[l].grad  # step 3
+            R[l] = (A[l] * c).data  # step 4
+        else:
+            R[l] = R[l + 1]
+
+    A[0] = (A[0].data).requires_grad_(True)
+    lb = (A[0].data * 0 + (0 - mean) / std).requires_grad_(True)
+    hb = (A[0].data * 0 + (1 - mean) / std).requires_grad_(True)
+
+    z = layers[0].forward(A[0]) + 1e-9  # step 1 (a)
+    z -= newlayer(layers[0], lambda p: p.clamp(min=0)).forward(lb)  # step 1 (b)
+    z -= newlayer(layers[0], lambda p: p.clamp(max=0)).forward(hb)  # step 1 (c)
+
+    # adding epsilon
+    epsilon = 1e-9
+    z_nonzero = torch.where(z == 0, torch.tensor(epsilon, device=z.device), z)
+    s = (R[1] / z_nonzero).data  # step 2
+
+    (z * s).sum().backward()
+    c, cp, cm = A[0].grad, lb.grad, hb.grad  # step 3
+    R[0] = (A[0] * c + lb * cp + hb * cm).data  # step 4
+    heatmap(
+        np.array(R[0][0].cpu()).sum(axis=0), 2, 2, intermediate_path
+    )  # HEATMAPPING TO SEE LRP MAPS WITH NEW RULE
+    return R[0][0].cpu()
+
+
+def get_nucleus_mask_for_graphcut(R):
+    res = np.array(R).sum(axis=0)
+    # Reshape the data to a 1D array
+    data_1d = res.flatten().reshape(-1, 1)
+    n_clusters = 2
+    kmeans = KMeans(n_clusters=n_clusters, random_state=0)
+    # kmeans.fit(data_1d)
+    kmeans.fit(data_1d)
+    # Step 4: Assign data points to clusters
+    cluster_assignments = kmeans.labels_
+    # Step 5: Reshape cluster assignments into a 2D binary matrix
+    binary_matrix = cluster_assignments.reshape(128, 128)
+    # Now, binary_matrix contains 0s and 1s, separating the data into two classes using K-Means clustering
+    rel_grouping = np.zeros((128, 128, 3), dtype=np.uint8)
+    rel_grouping[binary_matrix == 1] = [255, 0, 0]  # Main object (Blue)
+    rel_grouping[binary_matrix == 2] = [128, 0, 0]  # Second label (Dark Blue)
+    rel_grouping[binary_matrix == 0] = [0, 0, 255]  # Background (Red)
+    return rel_grouping
+
+
+def segment_nucleus(image, rel_grouping):  # clustered = rel_grouping
+
+    # GET THE BOUNDING BOX FROM CLUSTERED
+    blue_pixels = np.sum(np.all(rel_grouping == [255, 0, 0], axis=-1))
+    red_pixels = np.sum(np.all(rel_grouping == [0, 0, 255], axis=-1))
+    if red_pixels > blue_pixels:
+        color = np.array([255, 0, 0])
+    else:
+        color = np.array([0, 0, 255])
+    mask = cv2.inRange(rel_grouping, color, color)
+    contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
+    contour_areas = []
+    for contour in contours:
+        x, y, w, h = cv2.boundingRect(contour)
+        contour_areas.append(cv2.contourArea(contour))
+    contour_areas.sort()
+    contour_areas = np.array(contour_areas)
+    quartile_50 = np.percentile(contour_areas, 50)
+    selected_contours = [
+        contour for contour in contours if cv2.contourArea(contour) >= quartile_50
+    ]
+    x, y, w, h = cv2.boundingRect(np.concatenate(selected_contours))
+
+    # APPLY GRABCUT
+    fgModel = np.zeros((1, 65), dtype="float")
+    bgModel = np.zeros((1, 65), dtype="float")
+    mask = np.zeros(image.shape[:2], np.uint8)
+    rect = (x, y, x + w, y + h)
+
+    # IF BOUNDING BOX IS THE WHOLE IMAGE, THEN BOUNDING BOX METHOD WONT'T WORK -> SO USE INIT WITH MASK METHOD ITSELF
+    if (x, y, x + w, y + h) == (0, 0, 128, 128):
+
+        if (
+            red_pixels > blue_pixels
+        ):  # red is the dominant color and thus the background
+            mask[(rel_grouping == [255, 0, 0]).all(axis=2)] = (
+                cv2.GC_PR_FGD
+            )  # Probable Foreground
+            mask[(rel_grouping == [0, 0, 255]).all(axis=2)] = (
+                cv2.GC_PR_BGD
+            )  # Probable Background
+        else:  # blue is the dominant color and thus the background
+            mask[(rel_grouping == [0, 0, 255]).all(axis=2)] = (
+                cv2.GC_PR_FGD
+            )  # Probable Foreground
+            mask[(rel_grouping == [255, 0, 0]).all(axis=2)] = (
+                cv2.GC_PR_BGD
+            )  # Probable Background
+
+        (mask, bgModel, fgModel) = cv2.grabCut(
+            image,
+            mask,
+            rect,
+            bgModel,
+            fgModel,
+            iterCount=10,
+            mode=cv2.GC_INIT_WITH_MASK,
+        )
+
+    # ELSE PASS THE BOUNDING BOX FOR GRABCUT
+    else:
+        (mask, bgModel, fgModel) = cv2.grabCut(
+            image,
+            mask,
+            rect,
+            bgModel,
+            fgModel,
+            iterCount=10,
+            mode=cv2.GC_INIT_WITH_RECT,
+        )
+
+    # FORM THE COLORED SEGMENTATION MASK
+    clean_binary_mask = np.where(
+        (mask == cv2.GC_FGD) | (mask == cv2.GC_PR_FGD), 1, 0
+    ).astype("uint8")
+    nucleus_segment = np.zeros((128, 128, 3), dtype=np.uint8)
+    nucleus_segment[clean_binary_mask == 1] = [255, 0, 0]  # Main object (Blue)
+    nucleus_segment[clean_binary_mask == 0] = [0, 0, 255]  # Background (Red)
+    return nucleus_segment, clean_binary_mask
+
+
+def remove_nucleus(image1, blue_mask1, simple_lama):  # image, blue_mask, x, y
+    # expand the nucleus mask
+    # image1 = cv2.resize(image, (128,128))
+    # blue_mask1 = cv2.resize(blue_mask, (128,128))
+    kernel = np.ones((5, 5), np.uint8)  # Adjust the kernel size as needed
+    expandedmask = cv2.dilate(blue_mask1, kernel, iterations=1)
+    image_pil = Image.fromarray(cv2.cvtColor(image1, cv2.COLOR_BGR2RGB))
+    mask_pil = Image.fromarray(expandedmask)
+    result = simple_lama(image_pil, mask_pil)
+    result_cv2 = np.array(result)
+    result_cv2 = cv2.cvtColor(result_cv2, cv2.COLOR_RGB2BGR)
+    # result_cv2 = cv2.resize(result_cv2, (x,y))
+    return expandedmask, result_cv2
+
+
+def get_final_mask(nucleus_removed_img, blue_mask, expanded_mask):
+    # apply graphcut - init with rectangle (not mask approximation mask)
+    fgModel = np.zeros((1, 65), dtype="float")
+    bgModel = np.zeros((1, 65), dtype="float")
+
+    rect = (1, 1, nucleus_removed_img.shape[1], nucleus_removed_img.shape[0])
+
+    (mask, bgModel, fgModel) = cv2.grabCut(
+        nucleus_removed_img,
+        expanded_mask,
+        rect,
+        bgModel,
+        fgModel,
+        iterCount=20,
+        mode=cv2.GC_INIT_WITH_RECT,
+    )
+
+    clean_binary_mask = np.where(
+        (mask == cv2.GC_FGD) | (mask == cv2.GC_PR_FGD), 1, 0
+    ).astype("uint8")
+    colored_segmentation_mask = np.zeros((128, 128, 3), dtype=np.uint8)
+    colored_segmentation_mask[clean_binary_mask == 1] = [
+        128,
+        0,
+        0,
+    ]  # Main object (Blue)
+    colored_segmentation_mask[clean_binary_mask == 0] = [0, 0, 255]  # Background (Red)
+    colored_segmentation_mask[blue_mask > 0] = [255, 0, 0]
+    return colored_segmentation_mask
+
+
+def lrp_main(pixel_conversion):
+    i = 0
+    return_dict_count = 1
+    return_dict = {}
+    selected_indices = select_sample_images()
+    resized_shape = (128, 128)
+    cell_descriptors = [
+        ["Image Name", "Nucleus Area", "Cytoplasm Area", "Nucleus to Cytoplasm Ratio"]
+    ]
+
+    # MODEL SECTION STARTS FOR NEW MODEL
+    vgg16 = torchvision.models.vgg16(pretrained=True)
+    new_avgpool = nn.AdaptiveAvgPool2d(output_size=(4, 4))
+    vgg16.avgpool = new_avgpool
+    classifier_list = [
+        nn.Linear(8192, vgg16.classifier[0].out_features)
+    ]  # vgg16.classifier[0].out_features = 4096
+    classifier_list += list(vgg16.classifier.children())[
+        1:-1
+    ]  # Remove the first and last layers
+    classifier_list += [
+        nn.Linear(vgg16.classifier[6].in_features, 2)
+    ]  # vgg16.classifier[6].in_features = 4096
+    vgg16.classifier = nn.Sequential(
+        *classifier_list
+    )  # Replace the model classifier
+
+    PATH = "herlev_best_adam_vgg16_modified12_final.pth"
+    checkpoint = torch.load(PATH, map_location=torch.device("cpu"))
+    vgg16.load_state_dict(checkpoint)
+    # vgg16.to(torch.device('cuda'))
+    vgg16.eval()
+
+    layers = list(vgg16._modules["features"]) + toconv(
+        list(vgg16._modules["classifier"])
+    )
+    L = len(layers)
+    # MODEL SECTION ENDS
+
+    simple_lama = SimpleLama()
+
+    for imagefile in os.listdir(preprocessed_folder):
+        if (
+            "MACOSX".lower() in imagefile.lower()
+            or "." == imagefile[0]
+            or "_" == imagefile[0]
+        ):
+            print(imagefile)
+            continue
+        image_path = (
+            preprocessed_folder + os.path.splitext(imagefile)[0].lower() + ".png"
+        )
+        intermediate_path = (
+            intermediate_folder
+            + os.path.splitext(imagefile)[0].lower()
+            + "_heatmap.png"
+        )
+        save_path = (
+            segmentation_folder + os.path.splitext(imagefile)[0].lower() + "_mask.png"
+        )
+        table_path = (
+            tables_folder + os.path.splitext(imagefile)[0].lower() + "_table.png"
+        )
+
+        # print(i, imagefile)
+        image = cv2.imread(image_path)
+        original_shape = image.shape
+
+        image = cv2.resize(image, (128, 128))
+
+        layers_copy = copy.deepcopy(layers)
+        R = get_LRP_heatmap(image, L, layers_copy, imgclasses, intermediate_path)
+
+        rel_grouping = get_nucleus_mask_for_graphcut(R)
+
+        nucleus_segment, clean_binary_mask = segment_nucleus(image, rel_grouping)
+
+        expanded_mask, nucleus_removed_image = remove_nucleus(image, clean_binary_mask, simple_lama)
+
+        colored_segmentation_mask = get_final_mask(
+            nucleus_removed_image, clean_binary_mask, expanded_mask
+        )
+
+        cv2.imwrite(save_path, colored_segmentation_mask)
+
+        nucleus_area, cytoplasm_area, ratio = calculate_cell_descriptors(
+            original_shape, resized_shape, pixel_conversion, colored_segmentation_mask
+        )
+        cell_descriptors.append(
+            [
+                os.path.splitext(imagefile)[0].lower(),
+                nucleus_area,
+                cytoplasm_area,
+                ratio,
+            ]
+        )
+
+        create_cell_descriptors_table(table_path, nucleus_area, cytoplasm_area, ratio)
+
+        if i in selected_indices:
+            return_dict[f"image{return_dict_count}"] = str(
+                base64.b64encode(open(image_path, "rb").read()).decode("utf-8")
+            )
+            return_dict[f"inter{return_dict_count}"] = str(
+                base64.b64encode(open(intermediate_path, "rb").read()).decode("utf-8")
+            )
+            return_dict[f"mask{return_dict_count}"] = str(
+                base64.b64encode(open(save_path, "rb").read()).decode("utf-8")
+            )
+            return_dict[f"table{return_dict_count}"] = str(
+                base64.b64encode(open(table_path, "rb").read()).decode("utf-8")
+            )
+            return_dict_count += 1
+
+        i += 1
+
+        # Visualization
+        # for im in [image, gt2, rel_grouping, nucleus_segment, clean_binary_mask*255, nucleus_removed_image, colored_segmentation_mask]:
+        #   cv2_imshow(im)
+
+    # write cell_descriptors list to csv file
+    with open(cell_descriptors_path, "w", newline="") as csv_file:
+        writer = csv.writer(csv_file)
+        writer.writerows(cell_descriptors)
+
+    return return_dict