| 94 | |
| 95 | |
| 96 | def hstack_array(arrs, pad=5): |
| 97 | # each arr is of shape [..., H, W, 3] |
| 98 | arr_shape = arrs[0].shape |
| 99 | assert all(arr.shape == arr_shape for arr in arrs) |
| 100 | W = arr_shape[-2] |
| 101 | num_arrs = len(arrs) |
| 102 | arr_shape = list(arr_shape) |
| 103 | arr_shape[-2] = pad * (num_arrs - 1) + W * num_arrs |
| 104 | stack_arr = np.zeros(arr_shape, dtype=arrs[0].dtype) |
| 105 | for i, arr in enumerate(arrs): |
| 106 | start_idx = i * (W + pad) |
| 107 | stack_arr[..., start_idx:start_idx + W, :] = arr |
| 108 | return stack_arr |
| 109 | |
| 110 | |
| 111 | def vstack_array(arrs, pad=5): |