Output writer for video output.
| 180 | |
| 181 | |
| 182 | class VideoWriter(OutputWriter): |
| 183 | """Output writer for video output.""" |
| 184 | |
| 185 | def __init__(self, output_path: Path, fps: float = 30.0, render_depth: bool = True) -> None: |
| 186 | """Initialize VideoWriter.""" |
| 187 | output_path.parent.mkdir(exist_ok=True, parents=True) |
| 188 | self.output_path = output_path |
| 189 | self.image_writer = iio.get_writer(output_path, fps=fps) |
| 190 | |
| 191 | self.max_depth_estimate = None |
| 192 | if render_depth: |
| 193 | self.depth_writer = iio.get_writer(output_path.with_suffix(".depth.mp4"), fps=fps) |
| 194 | |
| 195 | def add_frame(self, image: torch.Tensor, depth: torch.Tensor) -> None: |
| 196 | """Add a single frame to output.""" |
| 197 | image_np = image.detach().cpu().numpy() |
| 198 | self.image_writer.append_data(image_np) |
| 199 | |
| 200 | if self.depth_writer is not None: |
| 201 | if self.max_depth_estimate is None: |
| 202 | self.max_depth_estimate = depth.max().item() |
| 203 | |
| 204 | colored_depth_pt = colorize_depth( |
| 205 | depth, |
| 206 | min(self.max_depth_estimate, METRIC_DEPTH_MAX_CLAMP_METER), # type: ignore[call-overload] |
| 207 | ) |
| 208 | colored_depth_np = colored_depth_pt.squeeze(0).permute(1, 2, 0).cpu().numpy() |
| 209 | self.depth_writer.append_data(colored_depth_np) |
| 210 | |
| 211 | def close(self): |
| 212 | """Finish writing.""" |
| 213 | self.image_writer.close() |
nothing calls this directly
no outgoing calls
no test coverage detected