pengsida/NeuralBody
1
1import torch2from matplotlib import cm3import matplotlib.pyplot as plt4import matplotlib.patches as patches5import numpy as np6import cv27 8 9def unnormalize_img(img, mean, std):10 """11 img: [3, h, w]12 """13 img = img.detach().cpu().clone()14 # img = img / 255.15 img *= torch.tensor(std).view(3, 1, 1)16 img += torch.tensor(mean).view(3, 1, 1)17 min_v = torch.min(img)18 img = (img - min_v) / (torch.max(img) - min_v)19 return img20 21 22def bgr_to_rgb(img):23 return img[:, :, [2, 1, 0]]24 25 26def horizon_concate(inp0, inp1):27 h0, w0 = inp0.shape[:2]28 h1, w1 = inp1.shape[:2]29 if inp0.ndim == 3:30 inp = np.zeros((max(h0, h1), w0 + w1, 3), dtype=inp0.dtype)31 inp[:h0, :w0, :] = inp032 inp[:h1, w0:(w0 + w1), :] = inp133 else:34 inp = np.zeros((max(h0, h1), w0 + w1), dtype=inp0.dtype)35 inp[:h0, :w0] = inp036 inp[:h1, w0:(w0 + w1)] = inp137 return inp38 39 40def vertical_concate(inp0, inp1):41 h0, w0 = inp0.shape[:2]42 h1, w1 = inp1.shape[:2]43 if inp0.ndim == 3:44 inp = np.zeros((h0 + h1, max(w0, w1), 3), dtype=inp0.dtype)45 inp[:h0, :w0, :] = inp046 inp[h0:(h0 + h1), :w1, :] = inp147 else:48 inp = np.zeros((h0 + h1, max(w0, w1)), dtype=inp0.dtype)49 inp[:h0, :w0] = inp050 inp[h0:(h0 + h1), :w1] = inp151 return inp52 53 54def transparent_cmap(cmap):55 """Copy colormap and set alpha values"""56 mycmap = cmap57 mycmap._init()58 mycmap._lut[:,-1] = 0.359 return mycmap60 61cmap = transparent_cmap(plt.get_cmap('jet'))62 63 64def set_grid(ax, h, w, interval=8):65 ax.set_xticks(np.arange(0, w, interval))66 ax.set_yticks(np.arange(0, h, interval))67 ax.grid()68 ax.set_yticklabels([])69 ax.set_xticklabels([])70 71 72color_list = np.array(73 [74 0.000, 0.447, 0.741,75 0.850, 0.325, 0.098,76 0.929, 0.694, 0.125,77 0.494, 0.184, 0.556,78 0.466, 0.674, 0.188,79 0.301, 0.745, 0.933,80 0.635, 0.078, 0.184,81 0.300, 0.300, 0.300,82 0.600, 0.600, 0.600,83 1.000, 0.000, 0.000,84 1.000, 0.500, 0.000,85 0.749, 0.749, 0.000,86 0.000, 1.000, 0.000,87 0.000, 0.000, 1.000,88 0.667, 0.000, 1.000,89 0.333, 0.333, 0.000,90 0.333, 0.667, 0.000,91 0.333, 1.000, 0.000,92 0.667, 0.333, 0.000,93 0.667, 0.667, 0.000,94 0.667, 1.000, 0.000,95 1.000, 0.333, 0.000,96 1.000, 0.667, 0.000,97 1.000, 1.000, 0.000,98 0.000, 0.333, 0.500,99 0.000, 0.667, 0.500,100 0.000, 1.000, 0.500,101 0.333, 0.000, 0.500,102 0.333, 0.333, 0.500,103 0.333, 0.667, 0.500,104 0.333, 1.000, 0.500,105 0.667, 0.000, 0.500,106 0.667, 0.333, 0.500,107 0.667, 0.667, 0.500,108 0.667, 1.000, 0.500,109 1.000, 0.000, 0.500,110 1.000, 0.333, 0.500,111 1.000, 0.667, 0.500,112 1.000, 1.000, 0.500,113 0.000, 0.333, 1.000,114 0.000, 0.667, 1.000,115 0.000, 1.000, 1.000,116 0.333, 0.000, 1.000,117 0.333, 0.333, 1.000,118 0.333, 0.667, 1.000,119 0.333, 1.000, 1.000,120 0.667, 0.000, 1.000,121 0.667, 0.333, 1.000,122 0.667, 0.667, 1.000,123 0.667, 1.000, 1.000,124 1.000, 0.000, 1.000,125 1.000, 0.333, 1.000,126 1.000, 0.667, 1.000,127 0.167, 0.000, 0.000,128 0.333, 0.000, 0.000,129 0.500, 0.000, 0.000,130 0.667, 0.000, 0.000,131 0.833, 0.000, 0.000,132 1.000, 0.000, 0.000,133 0.000, 0.167, 0.000,134 0.000, 0.333, 0.000,135 0.000, 0.500, 0.000,136 0.000, 0.667, 0.000,137 0.000, 0.833, 0.000,138 0.000, 1.000, 0.000,139 0.000, 0.000, 0.167,140 0.000, 0.000, 0.333,141 0.000, 0.000, 0.500,142 0.000, 0.000, 0.667,143 0.000, 0.000, 0.833,144 0.000, 0.000, 1.000,145 0.000, 0.000, 0.000,146 0.143, 0.143, 0.143,147 0.286, 0.286, 0.286,148 0.429, 0.429, 0.429,149 0.571, 0.571, 0.571,150 0.714, 0.714, 0.714,151 0.857, 0.857, 0.857,152 1.000, 1.000, 1.000,153 0.50, 0.5, 0154 ]155).astype(np.float32)156colors = color_list.reshape((-1, 3)) * 255157colors = np.array(colors, dtype=np.uint8).reshape(len(colors), 1, 1, 3)158