| 87 | return output.astype(imtype) |
| 88 | |
| 89 | def tensor2flow(output, imtype=np.uint8): |
| 90 | if isinstance(output, torch.autograd.Variable): |
| 91 | output = output.data |
| 92 | if len(output.size()) == 5: |
| 93 | output = output[0, -1] |
| 94 | if len(output.size()) == 4: |
| 95 | output = output[0] |
| 96 | output = output.cpu().float().numpy() |
| 97 | output = np.transpose(output, (1, 2, 0)) |
| 98 | #mag = np.max(np.sqrt(output[:,:,0]**2 + output[:,:,1]**2)) |
| 99 | #print(mag) |
| 100 | hsv = np.zeros((output.shape[0], output.shape[1], 3), dtype=np.uint8) |
| 101 | hsv[:, :, 0] = 255 |
| 102 | hsv[:, :, 1] = 255 |
| 103 | mag, ang = cv2.cartToPolar(output[..., 0], output[..., 1]) |
| 104 | hsv[..., 0] = ang * 180 / np.pi / 2 |
| 105 | hsv[..., 2] = cv2.normalize(mag, None, 0, 255, cv2.NORM_MINMAX) |
| 106 | rgb = cv2.cvtColor(hsv, cv2.COLOR_HSV2RGB) |
| 107 | return rgb |
| 108 | |
| 109 | def add_dummy_to_tensor(tensors, add_size=0): |
| 110 | if add_size == 0 or tensors is None: return tensors |