Numpy Image Hacks

grid.py

"""
Creating a spherical gaussian heatmap
 - in 2D
 - in 3D
"""
import pdb
import numpy as np

import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D

def get_heatmap3D(pcd_data, pcd_mean, hmap_sigma=1.7, hmap_sup = [10,10,10], show=False, verbose=False):
    # Step 1 - Derived Params
    hmap_mean = [hmap_sup[0]/2, hmap_sup[1]/2, hmap_sup[2]/2]
    (pcd_W,pcd_H,pcd_D) = pcd_data.shape
    pcd_mean_x, pcd_mean_y, pcd_mean_z = pcd_mean

    # Step 2 - Grid for heatmap
    lin_x = np.linspace(0,hmap_sup[0]-1,hmap_sup[0])
    lin_y = np.linspace(0,hmap_sup[1]-1,hmap_sup[1])
    lin_z = np.linspace(0,hmap_sup[2]-1,hmap_sup[2])
    gridX, gridY, gridZ = np.meshgrid(lin_y, lin_x, lin_z)
    dimensions = gridX.shape[0] * gridY.shape[1] * gridY.shape[2]
    gridx, gridy, gridz = np.reshape(gridX, (dimensions, -1)), np.reshape( gridY, (dimensions, -1)), np.reshape(gridZ, (dimensions, -1))
    grid = np.concatenate((gridx, gridy, gridz), axis=1) # grid.shape = [support^dims, dims]
    
    # Step 3 - Calculate heatmap
    gaussian = lambda coord: np.exp(-1*((coord[0]-hmap_mean[0])**2 
                                    + (coord[1]-hmap_mean[1])**2 
                                    + (coord[2]-hmap_mean[2])**2
                                ) / (2 * hmap_sigma * hmap_sigma))
    heatmap = np.apply_along_axis(gaussian, 1, grid).reshape((hmap_sup[0],hmap_sup[1],hmap_sup[2]))

    # Step 4 - Broadcast heatmap onto image data structure 
    border_sidex_left = int(pcd_mean_x)
    border_sidex_right = int(pcd_W - pcd_mean_x)
    border_sidey_left = int(pcd_mean_y)
    border_sidey_right = int(pcd_H - pcd_mean_y)
    border_top = int(pcd_D - pcd_mean_z)
    border_bottom = int(pcd_mean_z)
    
    pcd_stretch_x_left = int(pcd_mean_x - hmap_sup[0]/2)
    pcd_stretch_x_right = int(pcd_mean_x + hmap_sup[0]/2)
    pcd_stretch_y_left = int(pcd_mean_y - hmap_sup[1]/2)
    pcd_stretch_y_right = int(pcd_mean_y + hmap_sup[1]/2)
    pcd_stretch_z_bottom = int(pcd_mean_z - hmap_sup[2]/2)
    pcd_stretch_z_top = int(pcd_mean_z + hmap_sup[2]/2)

    heatmap_stretch_x_left = int(hmap_mean[0] - hmap_sup[0]/2)
    heatmap_stretch_x_right = int(hmap_mean[0] + hmap_sup[0]/2)
    heatmap_stretch_y_left = int(hmap_mean[1] - hmap_sup[1]/2)
    heatmap_stretch_y_right = int(hmap_mean[1] + hmap_sup[1]/2)
    heatmap_stretch_z_bottom = int(hmap_mean[2] - hmap_sup[2]/2)
    heatmap_stretch_z_top = int(hmap_mean[2] + hmap_sup[2]/2)

    if border_sidex_left < hmap_sup[0]//2:
        pcd_stretch_x_left = int(pcd_mean_x - border_sidex_left)
        heatmap_stretch_x_left = int(hmap_mean[0] - border_sidex_left)
    if border_sidex_right < hmap_sup[0]//2:
        pcd_stretch_x_right = int(pcd_mean_x + border_sidex_left)
        heatmap_stretch_x_right = int(hmap_mean[0] + border_sidex_right)
    if border_sidey_left < hmap_sup[1]//2:
        pcd_stretch_y_left = int(pcd_mean_y - border_sidey_left)
        heatmap_stretch_y_left = int(hmap_mean[1] - border_sidey_left)
    if border_sidey_right < hmap_sup[1]//2:
        pcd_stretch_y_right = int(pcd_mean_y + border_sidey_right)
        heatmap_stretch_y_right = int(hmap_mean[1] + border_sidey_right)
    if border_bottom < hmap_sup[2]//2:
        pcd_stretch_z_bottom = int(pcd_mean_z - border_bottom)
        heatmap_stretch_z_bottom = int(hmap_mean[2] - border_bottom)
    if border_top < hmap_sup[2]//2:
        pcd_stretch_z_top = int(pcd_mean_z + border_top)
        heatmap_stretch_z_top = int(hmap_mean[2] + border_top)

    if verbose:
        print (' - data:', pcd_stretch_x_left,pcd_stretch_x_right, ' || '
            , pcd_stretch_y_left,pcd_stretch_y_right, ' || '
            , pcd_stretch_z_bottom, pcd_stretch_z_top)
        print (' - hmap:', heatmap_stretch_x_left, heatmap_stretch_x_right, ' || '
                        , heatmap_stretch_y_left,heatmap_stretch_y_right, ' || '
                    , heatmap_stretch_z_bottom,heatmap_stretch_z_top)

    # pdb.set_trace()
    pcd_data[pcd_stretch_x_left:pcd_stretch_x_right
            , pcd_stretch_y_left:pcd_stretch_y_right
            , pcd_stretch_z_bottom:pcd_stretch_z_top] = heatmap[
                                                          heatmap_stretch_x_left:heatmap_stretch_x_right
                                                        , heatmap_stretch_y_left:heatmap_stretch_y_right
                                                        , heatmap_stretch_z_bottom:heatmap_stretch_z_top      
                                                        ]

    if show:
        fig = plt.figure()
        ax = fig.add_subplot(111, projection='3d')
        ax.set_xlabel('X axis')
        ax.set_ylabel('Y axis')
        ax.set_zlabel('Z axis')

        lin_x = np.linspace(0,pcd_data.shape[0]-1,pcd_data.shape[0])
        lin_y = np.linspace(0,pcd_data.shape[1]-1,pcd_data.shape[1])
        lin_z = np.linspace(0,pcd_data.shape[2]-1,pcd_data.shape[2])
        gridX, gridY, gridZ = np.meshgrid(lin_x, lin_y, lin_z)
        dimensions = gridX.shape[0] * gridY.shape[1] * gridZ.shape[2]
        gridx, gridy, gridz = np.reshape(gridX, (dimensions, -1)), np.reshape( gridY, (dimensions, -1)), np.reshape(gridZ, (dimensions, -1))
        grid = np.concatenate((gridx, gridy, gridz), axis=1)
        grid = grid.astype(np.int64)
        grid_data = pcd_data[grid[:,0], grid[:,1], grid[:,2]]
        pdb.set_trace()
        pnt3d=ax.scatter(grid[:,0],grid[:,1],grid[:,2], s=grid_data.flatten()*1, c=grid_data.flatten())

        cbar=plt.colorbar(pnt3d)
        cbar.set_label("Heatmap probability")
        plt.show()

    return pcd_data

def get_heatmap2D(img_data, img_mean, hmap_sigma=1.7, hmap_sup = [14,14], show=False):
    # Step 1 - Derived Params
    hmap_mean = [hmap_sup[0]/2, hmap_sup[1]/2] # =[y,x]
    (img_H,img_W) = img_data.shape
    img_mean_x, img_mean_y = img_mean

    # Step 2 - Grid for heatmap
    lin_x = np.linspace(0,hmap_sup[0]-1,hmap_sup[0])
    lin_y = np.linspace(0,hmap_sup[1]-1,hmap_sup[1])
    gridX, gridY = np.meshgrid(lin_y, lin_x)
    dimensions = gridX.shape[0] * gridY.shape[1]
    gridx, gridy = np.reshape(gridX, (dimensions, -1)), np.reshape( gridY, (dimensions, -1))
    grid = np.concatenate((gridx, gridy), axis=1) # grid.shape = [support^dims, dims]

    # Step 3 - Calculate heatmap
    gaussian = lambda coord: np.exp(-1*((coord[0]-hmap_mean[0])**2 
                                    + (coord[1]-hmap_mean[1])**2
                                ) / (2 * hmap_sigma * hmap_sigma))
    heatmap = np.apply_along_axis(gaussian, 1, grid).reshape((hmap_sup[0],hmap_sup[1]))

    # Step 4 - Broadcast heatmap onto image data structure 
    border_top = int(img_mean_y)
    border_bottom = int(img_H - img_mean_y)
    border_left = int(img_mean_x)
    border_right = int(img_W - img_mean_x)

    img_stretch_x_left = int(img_mean_x - hmap_sup[0]/2)
    img_stretch_x_right = int(img_mean_x + hmap_sup[0]//2)
    img_stretch_y_top = int(img_mean_y - hmap_sup[1]/2)
    img_stretch_y_bottom = int(img_mean_y + hmap_sup[1]//2)
    heatmap_stretch_x_left = int(hmap_mean[0] - hmap_sup[0]/2)
    heatmap_stretch_x_right = int(hmap_mean[0] + hmap_sup[0]/2)
    heatmap_stretch_y_top = int(hmap_mean[1] - hmap_sup[1]/2)
    heatmap_stretch_y_bottom = int(hmap_mean[1] + hmap_sup[1]/2)

    if border_right < hmap_sup[0]//2:
        img_stretch_x_right = int(img_mean_x + border_right)
        heatmap_stretch_x_right = int(hmap_mean[0] + border_right)
    if border_left < hmap_sup[0]//2:
        img_stretch_x_left = int(img_mean_x - border_left)
        heatmap_stretch_x_left = int(hmap_mean[0] - border_left)
    if border_top < hmap_sup[1]//2:
        img_stretch_y_top = int(img_mean_y - border_top)
        heatmap_stretch_y_top = int(hmap_mean[1] - border_top)
    if border_bottom < hmap_sup[1]//2:
        img_stretch_y_bottom = int(img_mean_y + border_bottom)
        heatmap_stretch_y_bottom = int(hmap_mean[1] + border_bottom)
    
    img_data[img_stretch_y_top:img_stretch_y_bottom
            , img_stretch_x_left:img_stretch_x_right] = heatmap[heatmap_stretch_y_top:heatmap_stretch_y_bottom
                                                            , heatmap_stretch_x_left:heatmap_stretch_x_right]

    if show:
        plt.imshow(img_data)
        plt.show()
    
    return img_data

if __name__ == "__main__":
    if (1):
        img_data = np.zeros((424, 512)) #(H,W)
        # img_mean = np.array([135,411]) #(x,y)
        img_mean = np.array([123, 423]) # border case
        img_data = get_heatmap2D(img_data, img_mean, show=True)
    else:
        if (0):
            pcd_data = np.zeros((50,40,35))
            pcd_mean = np.array([30,25,19])
            pcd_data = get_heatmap3D(pcd_data, pcd_mean, show=True, verbose=True)
        elif 1:
            pcd_data = np.zeros((50,40,35))
            pcd_mean = np.array([49,39,19])
            pcd_data = get_heatmap3D(pcd_data, pcd_mean, show=True, verbose=True)
        else:
            pcd_data = np.zeros((12,12,10))
            pcd_mean = np.array([5,5,6])
            pcd_data = get_heatmap3D(pcd_data, pcd_mean, hmap_sup = [5,5,5]
                                    , show=True, verbose=True)
        

"""
# Step5 - Open3D
import open3d as o3d
import numpy as np

pcd = o3d.geometry.PointCloud()
np_points = grid
pcd.points = o3d.utility.Vector3dVector(np_points)
pcd.colors = o3d.utility.Vector3dVector(np.repeat(heatmap.reshape(support*support*support,1), 3, axis=1))
o3d.visualization.draw_geometries([pcd])
"""

json_for_numpy.py

import json
import numpy as np

class NpEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, np.integer):
            return int(obj)
        elif isinstance(obj, np.floating):
            return float(obj)
        elif isinstance(obj, np.ndarray):
            return obj.tolist()
        else:
            return super(NpEncoder, self).default(obj)
    
json.dumps(data, cls=NpEncoder, sort_keys=True, indent=4)

padding.py

import matplotlib.pyplot as plt
def image_pad(ip, pad_x_l=1, pad_x_r=1, pad_y_l=0, pad_y_r=0):
    if (ip.shape[0] == 3):
        r = ip[0,:,:]
        g = ip[1,:,:]
        b = ip[2,:,:]
        r = np.pad(r, pad_width=((pad_x_l, pad_x_r), (pad_y_l, pad_y_r)), mode='constant', constant_values=1)
        g = np.pad(g, pad_width=((pad_x_l, pad_x_r), (pad_y_l, pad_y_r)), mode='constant', constant_values=1)
        b = np.pad(b, pad_width=((pad_x_l, pad_x_r), (pad_y_l, pad_y_r)), mode='constant', constant_values=1)
        
        op = np.zeros((3, r.shape[0], r.shape[1]))
        op[0,:,:] = r
        op[1,:,:] = g
        op[2,:,:] = b
        return op
    
    else:
        print (np.shape)
        print ('Not an image')
        return ip
    
if __name__ == "__main__":

    ip = np.random.rand(3,758, 758)
    print (ip.shape)
    op = image_pad(ip, pad_x_l=0, pad_x_r=15, pad_y_l=15, pad_y_r=15)
    print (op.shape)
    
    f, axarr = plt.subplots(1,2)
    print (op.transpose(-1, 1, 0).shape)
    axarr[0].imshow(ip.transpose(-1, 1, 0))
    axarr[1].imshow(op.transpose(-1, 1, 0))
    plt.show()

rotation_along_z_axis.py

import tensorflow as tf
x = tf.constant([[1, 2, 3], [4, 5, 6], [7,8,9]])
tf.reverse(tf.transpose(x), [1])

import tensorflow as tf
x = tf.constant([[[1.0,1.1],[1.2,1.3]], [[2.1,2.2],[2.3,2.4]], [[3.1,3.2],[3.3,3.4]]]) # shape = (3,2,2)
tf.transpose(tf.reverse(x, [1]), [0,2,1])

# rot = 270 (anti-clockwise)
import tensorflow as tf
x = tf.constant([[[1.0, 2.0, 3.0],[1.1, 2.1, 3.1]], [[1.2, 2.2, 3.2],[1.3, 2.3, 3.3]]]) # shape = (2,2,3)
print (x[:,:,0])
x_ = tf.transpose(tf.reverse(x, [0]), [1,0,2])
print (x_[:,:,0])

# rot = 270 (anti-clockwise)
x = tf.constant([[[[1.0], [2.0], [3.0]],[[1.1], [2.1], [3.1]]], [[[1.2], [2.2], [3.2]],[[1.3], [2.3], [3.3]]]]) # shape = (2,2,3,1)
print (x[:,:,0,0])
x_ = tf.transpose(tf.reverse(x, [0]), [1,0,2,3])
print (x_[:,:,0,0])

# rot = 270 (anti-clockwise)
x = tf.constant([[[[1.0,10], [2.0,20], [3.0,30]],[[1.1,11], [2.1,21], [3.1,31]]], [[[1.2,12], [2.2,22], [3.2,32]],[[1.3,13], [2.3,23], [3.3,33]]]]) # shape = (2,2,3,2)
print (x[:,:,0,0])
x_ = tf.transpose(tf.reverse(x, [0]), [1,0,2,3])
print (x_[:,:,0,0])

# rot = 90 (anti-clockwise)
x = tf.constant([[[[1.0,10], [2.0,20], [3.0,30]],[[1.1,11], [2.1,21], [3.1,31]]], [[[1.2,12], [2.2,22], [3.2,32]],[[1.3,13], [2.3,23], [3.3,33]]]]) # shape = (2,2,3,2)
print (x[:,:,0,0])
x_ = tf.reverse(tf.transpose(x, [1,0,2,3]), [0])
print (x_[:,:,0,0])