| 142 | |
| 143 | |
| 144 | def build_mask(self, data, is_bound): |
| 145 | _, T, H, W = data.shape |
| 146 | t = repeat(torch.arange(T), "T -> T H W", T=T, H=H, W=W) |
| 147 | h = repeat(torch.arange(H), "H -> T H W", T=T, H=H, W=W) |
| 148 | w = repeat(torch.arange(W), "W -> T H W", T=T, H=H, W=W) |
| 149 | border_width = (T + H + W) // 6 |
| 150 | pad = torch.ones_like(t) * border_width |
| 151 | mask = torch.stack([ |
| 152 | pad if is_bound[0] else t + 1, |
| 153 | pad if is_bound[1] else T - t, |
| 154 | pad if is_bound[2] else h + 1, |
| 155 | pad if is_bound[3] else H - h, |
| 156 | pad if is_bound[4] else w + 1, |
| 157 | pad if is_bound[5] else W - w |
| 158 | ]).min(dim=0).values |
| 159 | mask = mask.clip(1, border_width) |
| 160 | mask = (mask / border_width).to(dtype=data.dtype, device=data.device) |
| 161 | mask = rearrange(mask, "T H W -> 1 T H W") |
| 162 | return mask |
| 163 | |
| 164 | |
| 165 | def decode_video( |