patch_plot.py

import random
import cv2
import os
import numpy as np
import torch
import argparse
import matplotlib
import matplotlib.pyplot as plt
import matplotlib.patches as plt_patches
import matplotlib.patheffects as pe
from matplotlib import cm, colors
from deep_dating.networks import PatchCNN
from deep_dating.preprocessing import PatchExtractor, PatchMethod
from deep_dating.datasets import MPS, ScribbleLens, CLaMM, DatasetSplitter, SetType, Himanis
from deep_dating.util import save_figure, plt_clear, to_index, DATASETS_PATH
from tqdm import tqdm


def plot_patch_prediction(patch_extractor, output, final_prediction, true_label):
    """Color in all patches and make a boxplot at the bottom

    :param patch_extractor: instance of PatchExtractor that has been run
    :param output: outputs per patch
    :param final_prediction: global prediction
    :param true_label: ground truth label
    """
    # Clear plot from patch extractor
    plt_clear()
    fig, ax = plt.subplots(figsize=(15, 8))
    ax.imshow(patch_extractor.img, cmap="gray")

    patch_drawing_info = patch_extractor.get_extra_draw_info()
    
    unique_colors = np.sort(np.unique(output))
    cmap_full = plt.get_cmap("Spectral")
    cmap = plt.get_cmap("Spectral", unique_colors.shape[0])
    min_, max_ = np.min(output), np.max(output)
    if true_label:
        min_, max_ = min(min_, true_label), max(max_, true_label)
    norm = colors.Normalize(min_ - 1, max_ + 1)

    fontsize = 18

    # Draw each patch
    for i, (x, y, w, h) in enumerate(patch_drawing_info):
        label = output[i]
        color_idx = np.where(unique_colors == label)[0][0]
        color = cmap(color_idx)
        rect = plt_patches.Rectangle((x, y), w, h, linewidth=2, edgecolor=color, 
                                     alpha=0.65, facecolor=color, linestyle="dotted")
        ax.add_patch(rect)
        ax.annotate(str(label), (x + (w / 2), y + (h / 2)), fontsize=fontsize,
                    color="white", weight='bold', ha='center', va='center',
                    path_effects=[pe.withStroke(linewidth=3, foreground="black")])

    # Colorbar
    # color_map = cm.ScalarMappable(norm=norm, cmap=cmap_full)
    # cbar = fig.colorbar(color_map, ax=ax, orientation="horizontal", shrink=0.3, pad=0.05)
    # box_plot = cbar.ax.boxplot(output, vert=False, positions=[0.5], widths=[0.5])
    # for median in box_plot['medians']:
    #     median.set_color('none')
    # cbar.ax.get_yaxis().set_visible(False)
    # #cbar.ax.axvline(final_prediction, c='black', linewidth=2)
    # if true_label:
    #     cbar.ax.axvline(true_label, c='green', linewidth=2)
    # cbar.ax.tick_params(labelsize=fontsize)
    # cbar.set_label("Distribution of Years", fontsize=fontsize)

    #label_text = f" - Label (Ground Truth): {int(true_label)}" if true_label else ""
    #plt.title(f"Prediction (Median): {int(final_prediction)}" + label_text, fontsize=15)
    plt.title(f"Image Label = Year {true_label} -- Predictions:", fontsize=16)
    
    save_figure("example_img", fig=fig, show=True)


def get_saliency_map(patch, model):
    """Generate saliency map based on gradients
    Adapted from https://github.com/sunnynevarekar/pytorch-saliency-maps/blob/master/Saliency_maps_in_pytorch.ipynb

    :param patch: single patch image
    :param model: trained model
    :return: saliency map and prediction
    """
    model.eval()

    patch = cv2.cvtColor(patch, cv2.COLOR_GRAY2RGB)
    patches = [model.apply_transforms(x) for x in [patch]]
    patches = torch.from_numpy(np.array(patches))
    patches.requires_grad = True
    
    preds = model(patches)
    score, _ = torch.max(preds, 1)

    score.backward()
    
    slc, _ = torch.max(torch.abs(patches.grad[0]), dim=0) # get max along channel axis
    slc = (slc - slc.min()) / (slc.max()-slc.min()) # normalize to [0..1]

    return slc.numpy(), score.cpu().detach().numpy()[0]


def make_map(patch_extractor, patches, model):
    """Makes saliency maps, never ended up using this

    :param patch_extractor: instance of PatchExtractor that has been run
    :param patches: list of patches
    :param model: trained model
    :return: predictions per patch
    """
    patch_drawing_info = patch_extractor.get_extra_draw_info()
    saliency_map = np.zeros(patch_extractor.img.shape)
    
    outputs = []

    for i, (x, y, w, h) in tqdm(enumerate(patch_drawing_info), total=len(patch_drawing_info)):
        patch_map, model_output = get_saliency_map(patches[i], model)
        patch_map = cv2.resize(patch_map, (h, w), interpolation=cv2.INTER_NEAREST)
        outputs.append(model_output)
        saliency_map[y:y+h, x:x+w] = np.maximum(patch_map, saliency_map[y:y+h, x:x+w])
        
        break

    #saliency_map = cv2.normalize(saliency_map, None, alpha=0, beta=255, norm_type=cv2.NORM_MINMAX, dtype=cv2.CV_32F)

    #saliency_map = saliency_map.astype(np.uint8)
    print(saliency_map[y:y+h, x:x+w])
    plt_clear()
    fig, ax = plt.subplots(1, 1)
    #ax[0].imshow(patch_extractor.img, cmap="gray")

    alpha = 0.3
    output = saliency_map #patch_extractor.img
    #output = cv2.addWeighted(output, alpha, saliency_map, 1 - alpha, 0)
    
    color = "blue"
    x, y, _, _ = patch_drawing_info[0]
    x_end, y_end, w, h = patch_drawing_info[-1]
    w = w + x_end - x
    h = h + y_end - y
    # rect = plt_patches.Rectangle((x, y), w, h, linewidth=2, edgecolor=color, facecolor="none")
    # ax[0].add_patch(rect)
    
    my_cm = matplotlib.cm.get_cmap(plt.cm.hot)
    output = my_cm(output, bytes=True)
    output = cv2.cvtColor(output, cv2.COLOR_RGBA2RGB) 
    print(output.shape)
    

    text_idxs = np.where(patch_extractor.img == 0)
    back_idxs = np.where(patch_extractor.img != 0)
    print(text_idxs)
    # img = cv2.cvtColor(patch_extractor.img, cv2.COLOR_GRAY2RGB)
    img = np.zeros(output.shape, dtype=np.uint8)
    white = np.ones(output.shape, dtype=np.uint8)
    white[:, :, :] = [220,220,220]
    img[text_idxs] = output[text_idxs]
    img[back_idxs] = white[back_idxs].copy()

    # out_black = np.where(output == [0, 0, 0])
    # img[out_black] = white[out_black]
    
    print(np.unique(img))

    ax.imshow(img)
    
    save_figure("saliency_map", fig=fig, show=True)

    return outputs


def run_patch_pipeline(img_path, agg_func=np.median, true_label=None, plot=True, 
                       make_saliency_map=True, dataset=MPS(),
                       model_path="runs_v2/MPS_P1_Crossval/model_epoch_3_split_1.pt"):
    """Run pipeline and make prediction over a patch
    Does not use batches, hence might run out of memory if image contains lots of patches (fix this)

    :param img_path: path to image
    :param agg_func: aggregation function, defaults to np.median
    :param true_label: ground truth date of image if known, defaults to None
    :param plot: make plot, defaults to True
    :param make_saliency_map: make saliency map or not, defaults to True
    :param dataset: dating dataset, defaults to MPS()
    :param model_path: path to model, defaults to "runs_v2/MPS_P1_Crossval/model_epoch_3_split_1.pt"
    :return: date prediction of image 
    """

    classes = len(np.unique(dataset.y))
    model = PatchCNN("inception_resnet_v2", verbose=False, num_classes=classes)
    model.load(model_path, continue_training=False)
    
    patch_extractor = PatchExtractor(plot=plot, method=PatchMethod.SLIDING_WINDOW_LINES, is_binary=True)
    patches = patch_extractor.extract_patches(img_path)

    if plot and make_saliency_map:
        output = make_map(patch_extractor, patches, model)
        return
    else:
        patches = [cv2.cvtColor(x, cv2.COLOR_GRAY2RGB) for x in patches]
        patches = [model.apply_transforms(x) for x in patches]
        patches = torch.from_numpy(np.array(patches))

        output = model(patches)
        output = output.cpu().detach().numpy()

        if model.classification:
            class_idxs = np.argmax(output, axis=1)
            _, labels_unique = to_index(dataset.y)
            output = labels_unique[class_idxs]
            print(output)
            print(np.unique(output))
        else:
            output = output.flatten()
    
    output = np.round(output).astype(int)
    final_prediction = int(np.round(agg_func(output)))

    if plot:
        plot_patch_prediction(patch_extractor, output, final_prediction, true_label)

    return final_prediction


def run_over_dataset(dataset, set_type=SetType.VAL):
    """Run pipeline over all images in a pre-defined dataset

    :param dataset: dating dataset
    :param set_type: train, test or val, defaults to SetType.VAL
    """
    global y
    
    x, y = DatasetSplitter(dataset).get_data(set_type)
    data = list(zip(x, y))
    
    random.shuffle(data)

    for img_file, label in data:
        run_patch_pipeline(img_file, true_label=label, make_saliency_map=True)


if __name__ =="__main__":

    parser = argparse.ArgumentParser()
    parser.add_argument("--img", help="path to binary image", type=str, required=True)
    parser.add_argument("--model_path", help="path to model weights", type=str, required=True)
    parser.add_argument("--dataset", help="dataset name", type=str, required=True, choices=["mps", "clamm", "scribble", "himanis"])
    args = parser.parse_args()

    dataset = None
    for x in [MPS(), CLaMM(), ScribbleLens(), Himanis()]:
        if x.name.value == args.dataset:
            dataset = x
            break

    if dataset is None:
        raise Exception("invalid dataset given")

    # Show the prediction per patch for a specific image
    img = args.img
    print("Running prediction ...")
    run_patch_pipeline(img, true_label=None, make_saliency_map=False, dataset=dataset, model_path=args.model_path)

    # Can also manually run over entire dataset
    # dataset = MPS(os.path.join(DATASETS_PATH, "MPS_Binet"), dir_depth=1)
    # run_over_dataset(dataset)