app.py

from PIL import Image
import torch
import numpy as np
import gradio as gr
from pathlib import Path

from busam import Busam

resize_to = 512
checkpoint = "weights.pth"
device = "cpu"
print("Loading model...")
busam = Busam(checkpoint=checkpoint, device=device, side=resize_to)
minmaxnorm = lambda x: (x - x.min()) / (x.max() - x.min())

def edge_inference(img, algorithm, th_low=None, th_high=None):
    algorithm = algorithm.lower()
    print("Loading image...")
    img = np.array(img[:, :, :3])
    print("Getting features...")
    pred, size = busam.process_image(img, do_activate=True)
    print("Computing sobel...")
    if algorithm == "sobel":
        edge = busam.sobel_from_pred(pred, size)
    elif algorithm == "canny":
        th_low, th_high = th_low or 5000, th_high or 10000
        edge = busam.canny_from_pred(pred, size, th_low=th_low, th_high=th_high)
    else:
        raise ValueError("algorithm should be sobel or canny")
    edge = edge.cpu().numpy() if isinstance(edge, torch.Tensor) else edge

    print("Done")
    return Image.fromarray(
        (minmaxnorm(edge) * 255).astype(np.uint8)
    ).resize(size[::-1])

def dimred_inference(
    img,
    algorithm,
    resample_pct,
):
    algorithm = algorithm.lower()
    img = np.array(img[:, :, :3])
    print("Getting features...")
    pred, size = busam.process_image(img, do_activate=True)
    # pred is 1, F, S, S
    assert pred.shape[1] >= 3, "should have at least 3 channels"
    if algorithm == 'pca':
        from sklearn.decomposition import PCA
        reducer = PCA(n_components=3)
    elif algorithm == 'tsne':
        from sklearn.manifold import TSNE
        reducer = TSNE(n_components=3)
    elif algorithm == 'umap':
        from umap import UMAP
        reducer = UMAP(n_components=3)
    else:
        raise ValueError('algorithm should be pca, tsne or umap')
    np_y_hat = pred.detach().cpu().permute(1, 0, 2, 3).numpy()  # F, B, H, W
    np_y_hat = np_y_hat.reshape(np_y_hat.shape[0], -1)  # F, BHW
    np_y_hat = np_y_hat.T  # BHW, F
    resample_pct = 10**resample_pct
    resample_size = int(resample_pct * np_y_hat.shape[0])
    sampled_pixels = np_y_hat[:: np_y_hat.shape[0] // resample_size]
    print("dim reduction fit..." + " " * 30, end="\r")
    reducer = reducer.fit(sampled_pixels)
    print("dim reduction transform..." + " " * 30, end="\r")
    reducer.transform(np_y_hat[:10])  # to numba compile the function
    np_y_hat = reducer.transform(np_y_hat)  # BHW, 3
    print()
    print('Done. Saving...')
    # revert back to original shape
    colors = np_y_hat.reshape(pred.shape[2], pred.shape[3], 3)
    return Image.fromarray((minmaxnorm(colors) * 255).astype(np.uint8)).resize(
            size[::-1]
        )

def segmentation_inference(img, algorithm, scale):
    algorithm = algorithm.lower()
    img = np.array(img[:, :, :3])
    print("Getting features...")
    pred, size = busam.process_image(img, do_activate=True)
    print("Computing segmentation...")
    if algorithm == "kmeans":
        from sklearn.cluster import KMeans
        n_clusters = int(100 / 100**scale)
        kmeans = KMeans(n_clusters=n_clusters, random_state=0).fit(
            pred.view(pred.shape[1], -1).T
        )
        labels = kmeans.labels_
        labels = labels.reshape(pred.shape[2], pred.shape[3])
    elif algorithm == "felzenszwalb":
        from skimage.segmentation import felzenszwalb
        labels = felzenszwalb(
            (minmaxnorm(pred[0].cpu().numpy()) * 255).astype(np.uint8).transpose(1, 2, 0),
            scale=10**(8*scale-3),
            sigma=0,
            min_size=50,
        )
    elif algorithm == "slic":
        from skimage.segmentation import slic
        labels = slic(
            (minmaxnorm(pred[0].cpu().numpy()) * 255).astype(np.uint8).transpose(1, 2, 0),
            n_segments = int(100 / 100**scale),
            compactness=0.00001,
            sigma=1,
        )
    elif algorithm == 'watershed':
        from skimage.segmentation import watershed
        from skimage.feature import peak_local_max
        from scipy import ndimage as ndi
        sobel = busam.sobel_from_pred(pred, size)
        sobel = sobel.cpu().numpy() if isinstance(sobel, torch.Tensor) else sobel
        # contrast stretch sobel with 5% largest
        sobel = np.clip(sobel / np.percentile(sobel, 95), 0, 1)
        distance = ndi.distance_transform_edt(sobel < 1)  # distance to the borders
        coords = peak_local_max(distance, min_distance=int(1+100*scale), labels=sobel<1)
        mask = np.zeros(sobel.shape, dtype=bool)
        mask[tuple(coords.T)] = True
        markers, _ = ndi.label(mask)
        labels = watershed(sobel, markers)
    else:
        raise ValueError("algorithm should be kmeans, felzenszwalb or slic")
    print("Done")
    # the labels have values that are usually close to each other in the image and in magnitude, which complicates visualization
    # shuffle the labels to make them more visually distinct
    out = labels.copy()
    out[labels % 4 == 0] = labels[labels % 4 == 0] * 1 / 4
    out[labels % 4 == 1] = labels[labels % 4 == 1] * 4 // 4 + 1
    out[labels % 4 == 2] = labels[labels % 4 == 2] * 2 // 4 + 2
    out[labels % 4 == 3] = labels[labels % 4 == 3] * 3 // 4 + 3
    return Image.fromarray(
        (minmaxnorm(out) * 255).astype(np.uint8)
    ).resize(size[::-1])

def one_click_segmentation(img, row, col, threshold):
    row, col = int(row), int(col)
    img = np.array(img[:, :, :3])
    click_map = np.zeros(img.shape[:2], dtype=bool)
    side = min(img.shape[:2]) // 100
    click_map[max(0, row-side):min(img.shape[0], row+side), max(0, col-side//5):min(img.shape[0], col+side//5)] = True
    click_map[max(0, row-side//5):min(img.shape[0], row+side//5), max(0, col-side):min(img.shape[0], col+side)] = True
    print("Getting features...")
    pred, size = busam.process_image(img, do_activate=True)
    print("Getting mask...")
    mask = busam.get_mask((pred, size), (row, col))
    print("Done")
    print('shapes=', img.shape, mask.shape, click_map.shape)
    return (img, [(mask, 'Prediction'), (click_map, 'Click')])

with gr.Blocks() as demo:
    with gr.Tab('Edge detection'):
        algorithm = "canny"
        with gr.Row():
            def enable_sliders(algorithm):
                algorithm = algorithm.lower()
                return gr.Slider(visible=algorithm == "canny"), gr.Slider(visible=algorithm == "canny")

            with gr.Column():
                image_input = gr.Image(label="Input Image")
                run_button = gr.Button("Run")
                algorithm = gr.Radio(["Sobel", "Canny"], label="Algorithm", value="Sobel")
                # add sliders for th_low, th_high
                th_low_slider = gr.Slider(0, 32768, 10000, label="Canny's low threshold", visible=False)
                th_high_slider = gr.Slider(0, 32768, 20000, label="Canny's high threshold", visible=False)
            algorithm.change(enable_sliders, inputs=[algorithm], outputs=[th_low_slider, th_high_slider])
            with gr.Column():
                output_image = gr.Image(label="Output Image")
            run_button.click(edge_inference, inputs=[image_input, algorithm, th_low_slider, th_high_slider], outputs=output_image)
        gr.Examples([str(p) for p in Path('demoimgs').glob('*')], inputs=image_input)

    with gr.Tab('Reduction to 3D'):
        with gr.Row():
            with gr.Column():
                image_input = gr.Image(label="Input Image")
                algorithm = gr.Radio(["PCA", "TSNE", "UMAP"], label="Algorithm", value="PCA")
                run_button = gr.Button("Run")
                gr.Markdown("⚠️ UMAP is slow, TSNE is ULTRA-slow. They won't run on time. ⚠️")
                resample_pct = gr.Slider(-5, 0, -3, label="Resample (10^x)*100%")
            with gr.Column():
                output_image = gr.Image(label="Output Image")
            run_button.click(dimred_inference, inputs=[image_input, algorithm, resample_pct], outputs=output_image)
        gr.Examples([str(p) for p in Path('demoimgs').glob('*')], inputs=image_input)
    
    with gr.Tab('Classical Segmentation'):
        with gr.Row():
            with gr.Column():
                image_input = gr.Image(label="Input Image")
                algorithm = gr.Radio(['KMeans', 'Felzenszwalb', 'SLIC', 'Watershed'], label="Algorithm", value="SLIC")
                scale = gr.Slider(0.1, 1.0, 0.5, label="Scale")
                run_button = gr.Button("Run")
            with gr.Column():
                output_image = gr.Image(label="Output Image")
            run_button.click(segmentation_inference, inputs=[image_input, algorithm, scale], outputs=output_image)
        gr.Examples([str(p) for p in Path('demoimgs').glob('*')], inputs=image_input)

    with gr.Tab('One-click segmentation'):
        with gr.Row():
            with gr.Column():
                image_input = gr.Image(label="Input Image")
                threshold = gr.Slider(0, 1, 0.5, label="Threshold")
                with gr.Row():
                    row = gr.Textbox(10, label="Click's row")
                    col = gr.Textbox(10, label="Click's column")
                run_button = gr.Button("Run")
            with gr.Column():
                output_image = gr.AnnotatedImage(label="Output")
            run_button.click(one_click_segmentation, inputs=[image_input, row, col, threshold], outputs=output_image)
        gr.Examples([str(p) for p in Path('demoimgs').glob('*')], inputs=image_input)


demo.launch(share=False)