protparam.py

import os
import numpy as np
from scipy.ndimage import rotate
import matplotlib.pyplot as plt
import json
from skimage.transform import resize
from PIL import Image

from lib import coefs, plotting, parameterize, image_warp, image_utils, helpers


def plot_interpolation3(
        protein_img,
        nuc_mask,
        cell_mask,
        shift_dict,
        save_path,
        ori_fft,
        reduced_fft,
        n_coef,
        n_isos,
        binarize
):
    protein_ch = rotate(protein_img, shift_dict["theta"])
    nuclei = rotate(nuc_mask, shift_dict["theta"])
    cell = rotate(cell_mask, shift_dict["theta"])
    center_cell = plotting.center_of_mass(cell)
    center_nuclei = plotting.center_of_mass(nuclei)
    if (
            center_cell[1] > center_nuclei[1]
    ):  # Move 1 quadrant counter-clockwise
        protein_ch = rotate(protein_ch.copy(), 180)
        cell = rotate(cell.copy(), 180)
        nuclei = rotate(nuclei.copy(), 180)

    if binarize:
        thresh = plotting.threshold_otsu(protein_ch)
        protein_ch[protein_ch < thresh] = 0
        protein_ch[protein_ch >= thresh] = 1
    fig, ax = plt.subplots(1, (4 if reduced_fft != None else 3), figsize=(25, 30))
    fig.patch.set_facecolor("#191919")
    # fig.patch.set_alpha(1)
    # ax[0].imshow(shapes, origin="lower")
    ax[0].imshow(cell, alpha=0.5)
    ax[0].imshow(nuclei, alpha=0.5)
    # ax[0].set_facecolor('#191919')
    # ax[0].tight_axis()
    ax[1].imshow(protein_ch, origin="lower")
    ax[1].set_title(f"theta = {np.round(shift_dict['theta'], 1)}°", fontdict={'fontsize': 20, 'color': 'white'})
    ax[1].set_facecolor("#191919")
    contours = []
    for fcoef in [ori_fft[: n_coef * 2], ori_fft[n_coef * 2:]]:
        ix__, iy__ = coefs.inverse_fft(fcoef[:n_coef], fcoef[n_coef:])
        contours += [np.concatenate([ix__, iy__])]
    x_, y_ = parameterize.get_coordinates(
        contours[0].real, contours[1].real, [0, 0], n_isos=n_isos, plot=False
    )
    for i, (xi, yi) in enumerate(zip(x_, y_)):
        ax[0].scatter(
            yi + shift_dict["shift_c"][1],
            xi + shift_dict["shift_c"][0],
            s=0.5,
            alpha=0.3,
            color="w",
        )
        ax[1].scatter(
            yi + shift_dict["shift_c"][1],
            xi + shift_dict["shift_c"][0],
            s=0.5,
            alpha=0.3,
            color="w",
        )
        # ax[1].text(yi[i]+shift_dict["shift_c"][1], xi[i]+shift_dict["shift_c"][0], str(i))

    # Get intensity
    x = np.array(x_) + shift_dict["shift_c"][0]
    y = np.array(y_) + shift_dict["shift_c"][1]
    m = parameterize.get_intensity(protein_ch, x, y, k=5)
    m_normed = m / m.max()
    norm = plt.Normalize(vmin=0, vmax=1)
    # print(m_normed.min(), m_normed.max(), m_normed[1, :])
    for i, (xi, yi) in enumerate(zip(x_, y_)):
        ax[2].scatter(
            yi + shift_dict["shift_c"][1],
            xi + shift_dict["shift_c"][0],
            c=m_normed[i, :],
            norm=norm,
            s=2,
        )
        # ax[2].text(yi[i]+shift_dict["shift_c"][1], xi[i]+shift_dict["shift_c"][0], str(i))
    ax[2].axis("scaled")
    # ax[2].set_facecolor("#541352FF")
    ax[2].set_facecolor("#191919")
    ax[2].set_title(shift_dict['sc_label'], fontdict={'fontsize': 20, 'color': 'white'})
    if reduced_fft != None:
        fcoef_n = reduced_fft[0: n_coef * 2]
        fcoef_c = reduced_fft[n_coef * 2:]
        contours_r = []
        for fcoef in [fcoef_n, fcoef_c]:
            ix_, iy_ = coefs.inverse_fft(fcoef[:n_coef], fcoef[n_coef:])
            contours_r += [np.concatenate([ix_, iy_])]
        x_, y_ = parameterize.get_coordinates(
            contours_r[0].real, contours_r[1].real, [0, 0], n_isos=n_isos, plot=False
        )
        for i, (xi, yi) in enumerate(zip(x_, y_)):
            ax[3].plot(xi, yi, "--", alpha=0.3)
            ax[3].scatter(xi, yi, c=m_normed[i, :], norm=norm)
        ax[3].axis("scaled")
        ax[3].set_facecolor("#191919")
    plt.tight_layout()
    plt.savefig(save_path, bbox_inches="tight")
    plt.close()


def get_protein_intensity(
        protein_img,
        nuc_mask,
        cell_mask,
        shift_dict,
        ori_fft,
        n_coef,
        binarize=False,
        n_isos=[10, 10]
):
    protein_ch = rotate(protein_img, shift_dict["theta"])
    nuclei = rotate(nuc_mask, shift_dict["theta"])
    cell = rotate(cell_mask, shift_dict["theta"])
    center_cell = plotting.center_of_mass(cell)
    center_nuclei = plotting.center_of_mass(nuclei)
    if (
            center_cell[1] > center_nuclei[1]
    ):  # Move 1 quadrant counter-clockwise
        protein_ch = rotate(protein_ch, 180)

    if binarize:
        # protein_ch = exposure.equalize_hist(protein_ch)
        thresh = plotting.threshold_otsu(protein_ch)
        protein_ch[protein_ch < thresh] = 0
        protein_ch[protein_ch >= thresh] = 1

    contours = []
    n_coef_ = n_coef

    for fcoef in [ori_fft[: n_coef_ * 2], ori_fft[n_coef_ * 2:]]:
        ix__, iy__ = coefs.inverse_fft(fcoef[:n_coef_], fcoef[n_coef_:])
        contours += [np.concatenate([ix__, iy__])]
    x_, y_ = parameterize.get_coordinates(
        contours[0].real, contours[1].real, [0, 0], n_isos=n_isos, plot=False
    )
    # Get intensity
    x = np.array(x_) + shift_dict["shift_c"][0]
    y = np.array(y_) + shift_dict["shift_c"][1]
    m = parameterize.get_intensity(protein_ch, x, y, k=5)
    m_normed = m  # /m.max()
    return m_normed


def protparam(config, image_id, nuc_path, cell_path, protein_path, location, df_fft):
    fft_data = df_fft[df_fft['image'] == image_id].values[0].tolist()
    ori_fft = [complex(s.replace("i", "j")) for s in fft_data[10:]]
    shifts = dict()
    shifts["theta"] = float(fft_data[5])
    shifts["shift_c"] = (
        float(fft_data[6]),
        float(fft_data[7]),
    )
    shifts["sc_label"] = location

    nuc_mask = image_utils.read_grayscale_image(nuc_path)
    cell_mask = image_utils.read_grayscale_image(cell_path)
    prot_img = image_utils.read_grayscale_image(protein_path)

    if config["protparam_mode"] == "rings":
        intensity = get_protein_intensity(
            prot_img,
            nuc_mask,
            cell_mask,
            shift_dict=shifts,
            ori_fft=ori_fft,
            n_coef=config['n_coeffs'],
            binarize=True,
            n_isos=[10, 10]
        )
        np.save(os.path.join(config["output_dir"], "protparam", image_id + "_protein.npy"), intensity)

        if config["plot"]:
            plot_interpolation3(
                prot_img,
                nuc_mask,
                cell_mask,
                shifts,
                os.path.join(config["output_dir"], "protparam", image_id + "_protein_interpolation.png"),
                ori_fft,
                None,
                config['n_coeffs'],
                [10, 10],
                True
            )
    else:
        n_landmarks = 32  # number of landmark points for each ring, so final n_points to compute dx, dy will be 2*n_landmarks+1

        # Load average cell
        nu_centroid = [0, 0]
        avg_cell = np.load(os.path.join(config["output_dir"], "shapemode", "Avg_cell.npz"))
        ix_n = avg_cell["ix_n"]
        iy_n = avg_cell["iy_n"]
        ix_c = avg_cell["ix_c"]
        iy_c = avg_cell["iy_c"]

        # Move average shape from zero-centered coords to min=[0,0]
        min_x = np.min(ix_c)
        min_y = np.min(iy_c)
        nu_centroid[0] -= min_x
        nu_centroid[1] -= min_y
        ix_n = ix_n - min_x
        iy_n = iy_n - min_y
        ix_c = ix_n - min_x
        iy_c = iy_n - min_y

        if len(ix_n) != n_landmarks:
            ix_n, iy_n = helpers.equidistance(ix_n, iy_n, n_points=n_landmarks)
            ix_c, iy_c = helpers.equidistance(ix_c, iy_c, n_points=n_landmarks)
        nu_contour = np.stack([ix_n, iy_n]).T
        cell_contour = np.stack([ix_c, iy_c]).T

        pts_avg = np.vstack(
            [
                np.asarray(nu_centroid),
                helpers.realign_contour_startpoint(nu_contour),
                helpers.realign_contour_startpoint(cell_contour),
            ]
        )
        shape_x, shape_y = (
            np.round(cell_contour[:, 0].max()).astype("int"),
            np.round(cell_contour[:, 1].max()).astype("int"),
        )

        theta = float(fft_data[5])
        prot_img = rotate(prot_img, theta)
        nu_ = rotate(nuc_mask, theta)
        cell_ = rotate(cell_mask, theta)
        img_resized = resize(prot_img, (shape_x, shape_y), mode="constant")
        nu_resized = resize(nu_, (shape_x, shape_y), mode="constant")
        nu_resized[nu_resized > 0] = 255
        cell_resized = resize(cell_, (shape_x, shape_y), mode="constant")
        cell_resized[cell_resized > 0] = 255

        pts_ori = image_warp.find_landmarks(
            nu_resized, cell_resized, n_points=n_landmarks, border_points=False
        )
        pts_convex = (pts_avg + pts_ori) / 2

        warped1 = image_warp.warp_image(pts_ori, pts_convex, img_resized, plot=False, save_dir="")
        warped = image_warp.warp_image(pts_convex, pts_avg, warped1, plot=False, save_dir="")
        Image.fromarray((warped * 255).astype(np.uint8)).save(os.path.join(config["output_dir"], "protparam", image_id + "_warp.png"))

        if config["plot"]:
            fig, ax = plt.subplots(1, 5, figsize=(15, 30), sharex=True, sharey=True)
            ax[0].imshow(nu_, alpha=0.3)
            ax[0].imshow(cell_, alpha=0.3)
            ax[0].set_title('original shape')
            ax[1].imshow(nu_resized, alpha=0.3)
            ax[1].imshow(cell_resized, alpha=0.3)
            ax[1].set_title('resized shape+protein')
            ax[2].imshow(img_resized)
            ax[2].scatter(
                pts_ori[:, 1], pts_ori[:, 0], c=np.arange(len(pts_ori)), cmap="Reds"
            )
            ax[2].set_title("resized protein channel")
            ax[3].imshow(warped1)
            ax[3].scatter(
                pts_convex[:, 1],
                pts_convex[:, 0],
                c=np.arange(len(pts_ori)),
                cmap="Reds",
            )
            ax[3].set_title("ori_shape to midpoint")
            ax[4].imshow(warped)
            ax[4].scatter(
                pts_avg[:, 1], pts_avg[:, 0], c=np.arange(len(pts_ori)), cmap="Reds"
            )
            ax[4].set_title("midpoint to avg_shape")
            fig.savefig(os.path.join(config["output_dir"], "protparam", image_id + "_warp_plot.png"), bbox_inches="tight")
            plt.close()


def avg_protparam(config):
    cells_assigned_path = os.path.join(config["output_dir"], "shapemode", "cells_assigned_to_pc_bins.json")
    with open(cells_assigned_path, "r") as f:
        cells_assigned = json.load(f)

    avg_dir = os.path.join(config["output_dir"], "protparam", "avg")
    os.makedirs(avg_dir, exist_ok=True)

    for pc, bin_list in cells_assigned.items():
        for bin_idx, image_ids in enumerate(bin_list):
            if config["protparam_mode"] == "rings":
                intensities = []
                for image_id in image_ids:
                    path = os.path.join(config["output_dir"], "protparam", image_id + "_protein.npy")
                    if os.path.exists(path):
                        intensities.append(np.load(path))
                if len(intensities) == 0:
                    continue
                avg = np.mean(intensities, axis=0)
                np.save(os.path.join(avg_dir, f"{pc}_bin{bin_idx}_protein.npy"), avg)
                if config["plot"]:
                    fig, ax = plt.subplots(figsize=(6, 4))
                    ax.imshow(avg, aspect="auto", cmap="hot")
                    ax.set_xlabel("Points along contour")
                    ax.set_ylabel("Isocontour (nucleus → cell)")
                    ax.set_title(f"{pc} bin {bin_idx} — n={len(intensities)}")
                    plt.tight_layout()
                    plt.savefig(os.path.join(avg_dir, f"{pc}_bin{bin_idx}_protein.png"), bbox_inches="tight")
                    plt.close()

            elif config["protparam_mode"] == "warp":
                imgs = []
                for image_id in image_ids:
                    path = os.path.join(config["output_dir"], "protparam", image_id + "_warp.png")
                    if os.path.exists(path):
                        imgs.append(plt.imread(path))
                if len(imgs) == 0:
                    continue
                avg = np.mean(imgs, axis=0)
                Image.fromarray((avg * 255).astype(np.uint8)).save(os.path.join(avg_dir, f"{pc}_bin{bin_idx}_warp.png"))
            plt.close()