OpenMotionLab/MotionGPT

1import numpy as np2from PIL import Image3 4 5def format_color_vector(value, length):6    """Format a color vector.7    """8    if isinstance(value, int):9        value = value / 255.010    if isinstance(value, float):11        value = np.repeat(value, length)12    if isinstance(value, list) or isinstance(value, tuple):13        value = np.array(value)14    if isinstance(value, np.ndarray):15        value = value.squeeze()16        if np.issubdtype(value.dtype, np.integer):17            value = (value / 255.0).astype(np.float32)18        if value.ndim != 1:19            raise ValueError('Format vector takes only 1-D vectors')20        if length > value.shape[0]:21            value = np.hstack((value, np.ones(length - value.shape[0])))22        elif length < value.shape[0]:23            value = value[:length]24    else:25        raise ValueError('Invalid vector data type')26 27    return value.squeeze().astype(np.float32)28 29 30def format_color_array(value, shape):31    """Format an array of colors.32    """33    # Convert uint8 to floating34    value = np.asanyarray(value)35    if np.issubdtype(value.dtype, np.integer):36        value = (value / 255.0).astype(np.float32)37 38    # Match up shapes39    if value.ndim == 1:40        value = np.tile(value, (shape[0],1))41    if value.shape[1] < shape[1]:42        nc = shape[1] - value.shape[1]43        value = np.column_stack((value, np.ones((value.shape[0], nc))))44    elif value.shape[1] > shape[1]:45        value = value[:,:shape[1]]46    return value.astype(np.float32)47 48 49def format_texture_source(texture, target_channels='RGB'):50    """Format a texture as a float32 np array.51    """52 53    # Pass through None54    if texture is None:55        return None56 57    # Convert PIL images into numpy arrays58    if isinstance(texture, Image.Image):59        if texture.mode == 'P' and target_channels in ('RGB', 'RGBA'):60            texture = np.array(texture.convert(target_channels))61        else:62            texture = np.array(texture)63 64    # Format numpy arrays65    if isinstance(texture, np.ndarray):66        if np.issubdtype(texture.dtype, np.floating):67            texture = np.array(texture * 255.0, dtype=np.uint8)68        elif np.issubdtype(texture.dtype, np.integer):69            texture = texture.astype(np.uint8)70        else:71            raise TypeError('Invalid type {} for texture'.format(72                type(texture)73            ))74 75        # Format array by picking out correct texture channels or padding76        if texture.ndim == 2:77            texture = texture[:,:,np.newaxis]78        if target_channels == 'R':79            texture = texture[:,:,0]80            texture = texture.squeeze()81        elif target_channels == 'RG':82            if texture.shape[2] == 1:83                texture = np.repeat(texture, 2, axis=2)84            else:85                texture = texture[:,:,(0,1)]86        elif target_channels == 'GB':87            if texture.shape[2] == 1:88                texture = np.repeat(texture, 2, axis=2)89            elif texture.shape[2] > 2:90                texture = texture[:,:,(1,2)]91        elif target_channels == 'RGB':92            if texture.shape[2] == 1:93                texture = np.repeat(texture, 3, axis=2)94            elif texture.shape[2] == 2:95                raise ValueError('Cannot reformat 2-channel texture into RGB')96            else:97                texture = texture[:,:,(0,1,2)]98        elif target_channels == 'RGBA':99            if texture.shape[2] == 1:100                texture = np.repeat(texture, 4, axis=2)101                texture[:,:,3] = 255102            elif texture.shape[2] == 2:103                raise ValueError('Cannot reformat 2-channel texture into RGBA')104            elif texture.shape[2] == 3:105                tx = np.empty((texture.shape[0], texture.shape[1], 4), dtype=np.uint8)106                tx[:,:,:3] = texture107                tx[:,:,3] = 255108                texture = tx109        else:110            raise ValueError('Invalid texture channel specification: {}'111                             .format(target_channels))112    else:113        raise TypeError('Invalid type {} for texture'.format(type(texture)))114 115    return texture116