# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/00_vision_plot.ipynb.

# %% auto 0
__all__ = ['show_med_img']

# %% ../nbs/00_vision_plot.ipynb 1
from fastai.data.all import *
from torchio.visualization import rotate

# %% ../nbs/00_vision_plot.ipynb 3
def _get_slice(image, channel:int, indices:(int, list), anatomical_plane:int, voxel_size:(int, list)):
    '''A private method to get a 2D tensor and aspect ratio for plotting. Modified code from torchio function `plot_volume`.'''

    if voxel_size is None:
        warnings.warn('Voxel size not defined. Aspect ratio of the plot might not be correct.')
        voxel_size = 1

    if not is_listy(voxel_size):
        voxel_size = (voxel_size, voxel_size, voxel_size)

    data = image.data[channel]

    sr, sa, ss = voxel_size[0], voxel_size[1], voxel_size[2]
    sliced_img, aspect = None, None

    if indices is None: indices = np.array(data.shape) // 2

    i, j, k = indices

    if anatomical_plane == 1:
        sliced_img = rotate(data[:, j, :], radiological=True, n=1)
        aspect = ss / sr
    elif anatomical_plane == 2:
        sliced_img = rotate(data[:, :, k], radiological=True, n=1)
        aspect = sa / sr
    else:
        sliced_img = rotate(data[i, :, :], radiological=True, n=1)
        aspect = ss / sa

    return sliced_img, aspect

# %% ../nbs/00_vision_plot.ipynb 4
@delegates(plt.Axes.imshow, keep=True, but=['shape', 'imlim'])
def show_med_img(im, ctx, channel:int, indices:(int, list), anatomical_plane:int, voxel_size:(int, list), ax=None, figsize=None, title=None, **kwargs):
    '''Show a PyTorch image on `ax`. Modified code from fastai function `show_image`.'''

    if hasattrs(im, ('data','cpu','permute')):
        im = im.data.cpu()
        im, aspect = _get_slice(im, channel=channel, anatomical_plane=anatomical_plane, voxel_size=voxel_size, indices=indices)

    ax = ifnone(ax,ctx)
    if ax is None: _,ax = plt.subplots(figsize=figsize) #ax er kun none naar en bruker .show().

    ax.imshow(im, aspect=aspect, **kwargs)
    if title is not None: ax.set_title(title)
    ax.axis('off')
    return ax
