util.py 文件源码

def _transform_state_color_space_np(s):
    # s: [h, w, c*num_frame]
    num_splits = int(s.shape[-1] / 3)    
    return np.concatenate([cv2.cvtColor(x, cv2.COLOR_RGB2GRAY)[:,:,np.newaxis] for x in np.split(s, num_splits, axis=2)], axis=2)