data_tools.py

import os
import subprocess
from distutils import dir_util
from pathlib import Path

import cv2 as cv
import numpy as np
from PIL import Image
from matplotlib import pyplot as plt
from tqdm import tqdm

from features import NucleiFeatures

Image.MAX_IMAGE_PIXELS = None


def get_x_and_y(name):
    x, y = os.path.splitext(name)[0].split('_')[-2:]
    return int(x), int(y)


def split_image(img, x_tiles_cnt=None, y_tiles_cnt=None, x_tile_size=None, y_tile_size=None, base='img'):
    """
    Splits an image array to smaller tiles for further segmentation.
    Specify tiles count OR tiles size.
    X axis means the arr.shape[1] coordinate, be careful!
    Tile names are used to restore the initial image after segmentation.

    Parameters
    ----------
    img : numpy ndarray
        The input image.
    x_tiles_cnt : integer
        Number of tiles along the x axis of img.
    y_tiles_cnt : integer
        Number of tiles along the y axis of img.
    x_tile_size : integer
        Size of tile along x axis.
    y_tile_size : integer
        Size of tile along y axis.
    base : str
        Base for tile names.

    Returns
    -------
    tiles : list
        List of tiles.
    tile_names : list
        List of tile names.
    """
    if (x_tile_size is not None) and (y_tile_size is not None):
        x_tiles_cnt = img.shape[1] // x_tile_size
        y_tiles_cnt = img.shape[0] // y_tile_size

    if (x_tiles_cnt is not None) and (y_tiles_cnt is not None):
        x_ticks = np.linspace(0, img.shape[1], x_tiles_cnt + 1).astype(int)
        y_ticks = np.linspace(0, img.shape[0], y_tiles_cnt + 1).astype(int)

    else:
        raise Exception('Specify tiles count OR tiles size.')

    tiles = []
    tile_names = []

    for x_num, x in enumerate(zip(x_ticks[:-1], x_ticks[1:])):
        for y_num, y in enumerate(zip(y_ticks[:-1], y_ticks[1:])):
            tiles.append(img[y[0]:y[1], x[0]:x[1]])
            tile_names.append(f'{base}_{x_num}_{y_num}')
    return tiles, tile_names


def prepare_test_data(tiles, tile_names, base_dir, force=False):
    """
    Saves data in the proper way.

    Parameters
    ----------
    tiles : list
        List of tiles.
    tile_names : list
        List of tile names.
    base_dir : str
        Full path to base directory, 'full/path/../data_test' in normal case.
    force : bool
        Rewrite existing files

    Returns
    -------
    None
    """

    base_dir = Path(base_dir)
    if not os.path.exists(base_dir):
        os.makedirs(base_dir)

    if not force and len(os.listdir(base_dir)) > 0:
        raise ValueError(f'base_dir {base_dir} is not empty, use force=True option if you want to rewrite files')
    elif len(os.listdir(base_dir)):
        dir_util.remove_tree(str(base_dir))
        os.makedirs(base_dir)

    for tile, name in zip(tiles, tile_names):
        if tile.max() <= 1:
            tile = (tile * 255).astype(np.uint8)
        os.mkdir(base_dir / name)
        os.mkdir(base_dir / name / 'images')
        cv.imwrite(str(base_dir / name / 'images' / f'{name}.png'), tile)


def restore_image(work_dir, tiff=False):
    """
    Restores the initial image.

        Parameters
    ----------
    work_dir : str
        Full path to directory with files.
    tiff : bool
        Is the target image a multilayer tiff or not

    Returns
    -------
    img : numpy ndarray
        Initial image

    """

    work_dir = Path(work_dir)

    file_names = sorted(os.listdir(work_dir), key=lambda x: get_x_and_y(x)[::-1])
    coords = np.array([get_x_and_y(n) for n in file_names])
    x_max, y_max = coords.max(axis=0)

    if tiff:
        tiles = {}
        max_number = int(0)
        for n in file_names:
            tmp = cv.imread(str(work_dir / n), -1)
            tmp = (tmp + max_number) * (tmp > 0)
            max_number = tmp.max()
            tiles[get_x_and_y(n)] = tmp.copy()
    else:
        tiles = {get_x_and_y(n): cv.imread(str(work_dir / n), -1) for n in file_names}

    long_tiles = []

    for y in range(y_max + 1):
        long_tiles.append([])
        for x in range(x_max + 1):
            long_tiles[-1].append(tiles[(x, y)])

    long_tiles = [np.hstack(i) for i in long_tiles]

    return np.vstack(long_tiles)


def perform_segmentation(full_img_path, sample_dir, network_dir, force=False, features=None):
    network_dir = Path(network_dir)
    try:
        full_img = cv.imread(full_img_path, -1)
    except cv.error:
        full_img = plt.imread(full_img_path)
    tiles, tile_names = split_image(img=full_img, x_tile_size=1000, y_tile_size=1000)
    prepare_test_data(tiles, tile_names, sample_dir, force=force)

    try:
        dir_util.remove_tree(str(network_dir / 'data_test'))
    except:
        pass
    os.mkdir(str(network_dir / 'data_test'))
    dir_util.copy_tree(sample_dir, str(network_dir / 'data_test'));

    try:
        dir_util.remove_tree(str(network_dir / 'predictions'))
    except:
        pass

    try:
        dir_util.remove_tree(str(network_dir / 'albu/results_test'))
    except:
        pass

    result_dir = str(Path(sample_dir)) + '_segmented'

    subprocess.run(f"cd {network_dir} && bash 'predict_test.sh'", shell=True)
    dir_util.copy_tree(str(network_dir / 'predictions'), result_dir);

    if features is not None:
        NucleiFeatures(f'{result_dir}/lgbm_test_sub2', sample_dir,
                       features=features).df().to_csv(f'{result_dir}/{os.path.split(sample_dir)[1]}.csv',
                                                      index=False)


def color_tiff(img, n=60):
    img = img % n
    seg_color = np.zeros((*img.shape, 3), dtype=np.uint8)
    for i in tqdm(range(1, img.max() + 1)):
        seg_color[img == i] = np.random.randint(0, 255, 3)
    return seg_color