Prepare the input of model

def set_data(self, frames):
        """
        Prepare the input of model
        """
        data_frames = []
        for frame in frames:
            #frame H x W x C
            frame = frame.swapaxes(0, 1) # swap width and height to form format W x H x C
            if len(frame.shape) < 3:
                frame = np.array([frame]).swapaxes(0, 2).swapaxes(0, 1) # Add grayscale channel
            data_frames.append(frame)
        frames_n = len(data_frames)
        data_frames = np.array(data_frames) # T x W x H x C
        data_frames = np.rollaxis(data_frames, 3) # C x T x W x H
        data_frames = data_frames.swapaxes(2, 3) # C x T x H x W  = NCDHW

        self.data = data_frames
        self.length = frames_n