(flow_map_np)
| 27 | return (3e-4 - 3e-6) * mul + 3e-6 |
| 28 | |
| 29 | def flow2rgb(flow_map_np): |
| 30 | h, w, _ = flow_map_np.shape |
| 31 | rgb_map = np.ones((h, w, 3)).astype(np.float32) |
| 32 | normalized_flow_map = flow_map_np / (np.abs(flow_map_np).max()) |
| 33 | |
| 34 | rgb_map[:, :, 0] += normalized_flow_map[:, :, 0] |
| 35 | rgb_map[:, :, 1] -= 0.5 * (normalized_flow_map[:, :, 0] + normalized_flow_map[:, :, 1]) |
| 36 | rgb_map[:, :, 2] += normalized_flow_map[:, :, 1] |
| 37 | return rgb_map.clip(0, 1) |
| 38 | |
| 39 | def train(model, local_rank): |
| 40 | if local_rank == 0: |