Visualize 3D objects in a video manner

visual_3d.py

""" Visualize 3D objects in a video manner.
Current config is for Cirrus OCT only.

Author: Jian Shi
Email: sj8716643@126.com
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os

import numpy as np
import matplotlib
import matplotlib.pyplot as plt
try:
    from videofig import videofig
except ModuleNotFoundError:
    import pip
    if hasattr(pip, 'main'):
        pip.main(['install', 'videofig'])
    else:
        pip._internal.main(['install', 'videofig'])
    from videofig import videofig


class Redraw(object):
    def __init__(self, file, dim=0):
        self.initialized = False
        if isinstance(file, str):
            cube = self.load_img(file)
        elif isinstance(file, np.ndarray):
            cube = file
        else:
            raise ValueError(f"Expect 'file' to be either path-like string or ndarray. Got {file}.")
        cube = np.moveaxis(cube, dim, 0)
        if cube.shape[2] > cube.shape[1]:
            cube = np.transpose(cube, (0, 2, 1))
        cube = np.flip(cube, 1)
        self.cube = cube

    def load_img(self, file):
        if file.endswith('.npy'):
            cube = np.load(file)
        else:
            cube = np.fromfile(file, np.uint8)
            if len(cube) == 67108864:
                cube = cube.reshape((128, 1024, 512))
            elif len(cube) == 40960000:
                cube = cube.reshape((200, 1024, 200))
            else:
                raise ValueError(f"{file}, not supportted filetype.")
        return cube

    def draw(self, f, ax):
        if f > len(self.cube) - 1:
            img = self.cube[-1]
        else:
            img = self.cube[f]
        if not self.initialized:
            self.img_handle = ax.imshow(img, cmap='gray')
            self.initialized = True
        else:
            self.img_handle.set_array(img)

    def __len__(self):
        return len(self.cube)


if __name__ == '__main__':
    """
    Sample command:
        $ python visual_3d.py --dim 0 1 2 --fp "/raid0/OCTNEW/HD001-200/PHD053 20130430/PHD053_Macular Cube 200x200_4-30-2013_14-29-20_OD_sn12436_cube_z.img"
    """
    import argparse
    parser = argparse.ArgumentParser()
    parser.add_argument("--fp", required=True, type=str, help="The file path to the file.")
    parser.add_argument("--dim", type=int, nargs='+', help="dim to visualize. Default is 0.", default=[0, 1, 2])
    args = parser.parse_args()

    def redraw_fn(f, axes):
        # if axes is not a list
        if not isinstance(axes, list):
            redraw_fn.sub1.draw(f, axes)
            return

        num = 0
        if 0 in args.dim:
            redraw_fn.sub1.draw(f, axes[num])
            num += 1
        if 1 in args.dim:
            redraw_fn.sub2.draw(f, axes[num])
            num += 1
        if 2 in args.dim:
            redraw_fn.sub3.draw(f, axes[num])
            num += 1

    length = 0
    if 0 in args.dim:
        redraw_fn.sub1 = Redraw(args.fp, 0)
        length = length if len(redraw_fn.sub1) < length else len(redraw_fn.sub1)
    if 1 in args.dim:
        redraw_fn.sub2 = Redraw(args.fp, 1)
        length = length if len(redraw_fn.sub2) < length else len(redraw_fn.sub2)
    if 2 in args.dim:
        redraw_fn.sub3 = Redraw(args.fp, 2)
        length = length if len(redraw_fn.sub3) < length else len(redraw_fn.sub3)

    videofig(length, redraw_fn, grid_specs={'nrows': 1, 'ncols': len(args.dim)},
             layout_specs=[f'[{i}]' for i in range(len(args.dim))])