Implemented a basic CLI

commit-id:b2136058

Author: integraledelebesgue

packages/child-lab-cli/hello.py

@@ -4,4 +4,24 @@ version = "0.1.0"
 description = "Add your description here"
 readme = "README.md"
 requires-python = ">=3.12.7"
-dependencies = []
+dependencies = [
+    "click>=8.1.8",
+    "pyserde>=0.23.0",
+    "torch>=2.6.0",
+    "tqdm>=4.67.1",
+    "transformation-buffer",
+    "video-io",
+    "depth-pro>=0.1.0",
+]
+
+[tool.uv.sources]
+depth-pro = { git = "https://github.com/child-lab-uj/depth-pro.git" }
+transformation-buffer = { workspace = true }
+video-io = { workspace = true }
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[project.scripts]
+child-lab = "child_lab_cli:cli"

packages/child-lab-cli/pyproject.toml

@@ -0,0 +1,20 @@
+import click
+
+from .commands import (
+    calibrate,
+    estimate_transformations,
+    generate_pointcloud,
+    process,
+    visualize,
+)
+
+
+@click.group('child-lab')
+def cli() -> None: ...
+
+
+cli.add_command(calibrate)
+cli.add_command(estimate_transformations)
+cli.add_command(generate_pointcloud)
+cli.add_command(process)
+cli.add_command(visualize)

packages/child-lab-cli/src/child_lab_cli/__init__.py

@@ -0,0 +1,13 @@
+from .calibrate import calibrate
+from .estimate_transformations import estimate_transformations
+from .generate_pointcloud import generate_pointcloud
+from .process import process
+from .visualize import visualize
+
+__all__ = [
+    'calibrate',
+    'estimate_transformations',
+    'generate_pointcloud',
+    'process',
+    'visualize',
+]

packages/child-lab-cli/src/child_lab_cli/commands/__init__.py

@@ -0,0 +1,43 @@
+from pathlib import Path
+
+import click
+
+
+@click.command('calibrate')
+@click.argument('workspace', type=Path)
+@click.argument('videos', type=Path, nargs=-1)
+@click.option('--square-size', type=float, help='Board square size in centimeters')
+@click.option(
+    '--inner-board-corners',
+    nargs=2,
+    type=int,
+    help="Number of chessboard's inner corners in rows and columns that calibration algorithm should locate",
+)
+@click.option(
+    '--max-samples',
+    type=int,
+    required=False,
+    help='Maximal number of board samples to collect',
+)
+@click.option(
+    '--max-speed',
+    type=float,
+    required=False,
+    help='Maximal speed the board can move with to be captured, in pixels per second',
+)
+@click.option(
+    '--min-distance',
+    type=float,
+    required=False,
+    help='Minimal distance between new observation and the previous observations to be captured',
+)
+def calibrate(
+    workspace: Path,
+    videos: list[Path],
+    square_size: float,
+    inner_board_corners: tuple[int, int],
+    max_samples: int | None,
+    max_speed: float | None,
+    min_distance: float | None,
+) -> None:
+    pass

packages/child-lab-cli/src/child_lab_cli/commands/calibrate.py

@@ -0,0 +1,38 @@
+from pathlib import Path
+
+import click
+
+
+@click.command('estimate-transformations')
+@click.argument('workspace', type=Path)
+@click.argument('videos', type=Path, nargs=-1)
+@click.option('--marker-dictionary', type=str, help='Dictionary to detect markers from')
+@click.option('--marker-size', type=float, help='Marker size in centimeters')
+@click.option(
+    '--device',
+    type=str,
+    required=False,
+    help='Torch device to use for tensor computations',
+)
+@click.option(
+    '--checkpoint',
+    type=Path,
+    required=False,
+    help='File containing serialized Buffer to load and place new transformations in',
+)
+@click.option(
+    '--skip',
+    type=int,
+    required=False,
+    help='Seconds of videos to skip at the beginning',
+)
+def estimate_transformations(
+    workspace: Path,
+    videos: list[Path],
+    marker_dictionary: str,
+    marker_size: float,
+    device: str | None,
+    checkpoint: Path | None,
+    skip: int | None,
+) -> None:
+    pass

packages/child-lab-cli/src/child_lab_cli/commands/estimate_transformations.py

@@ -0,0 +1,120 @@
+from collections.abc import Generator
+from pathlib import Path
+
+import click
+import torch
+import tqdm
+from depth_pro import Config, DepthPro
+from jaxtyping import UInt8
+from serde.yaml import from_yaml
+from torchvision.transforms import Compose, ConvertImageDtype, Normalize, Resize
+from video_io.calibration import Calibration
+from video_io.reader import Reader
+
+from ..workspace import Workspace
+
+
+@click.command('generate-pointcloud')
+@click.argument('workspace_root', type=Path)
+@click.argument('video_name', type=str)
+@click.option(
+    '--checkpoint',
+    type=Path,
+    required=True,
+    help='DepthPro checkpoint to build the depth estimator from',
+)
+@click.option(
+    '--batch-size',
+    type=int,
+    default=16,
+    required=False,
+    help='Number of frames to process as a single batch',
+)
+@click.option(
+    '--device',
+    type=str,
+    required=False,
+    help='Torch device to use for tensor computations',
+)
+def generate_pointcloud(
+    workspace_root: Path,
+    video_name: str,
+    checkpoint: Path,
+    batch_size: int,
+    device: str | None,
+) -> None:
+    workspace = Workspace.in_directory(workspace_root)
+    output = workspace.output / 'points'
+
+    for video in workspace.calibrated_videos():
+        if video.name == video_name:
+            break
+    else:
+        raise FileNotFoundError(f'Video {video_name} not found in {workspace.input}.')
+
+    calibration = from_yaml(Calibration, video.calibration.read_text())
+
+    main_device = torch.device(device or 'cpu')
+
+    reader = Reader(video.location, torch.device('cpu'))
+    depth_estimator = DepthEstimator(checkpoint, main_device)
+
+    click.echo('Depth estimator created!')
+
+    def batched_frames(
+        reader: Reader,
+        batch_size: int,
+    ) -> Generator[torch.Tensor, None, None]:
+        while (frames := reader.read_batch(batch_size)) is not None:
+            yield frames
+
+    progress_bar = tqdm.tqdm(
+        range(0, reader.metadata.frames, batch_size),
+        desc='Processing batches of frames',
+    )
+
+    for i, frames in enumerate(batched_frames(reader, batch_size)):
+        depths = depth_estimator.predict(frames.to(main_device))
+        perspective_points = calibration.unproject_depth(depths)
+
+        torch.save(perspective_points, output / f'points_{i}.pt')
+
+        progress_bar.update()
+
+    click.echo('Done!')
+
+
+# TODO: Delete this from here
+class DepthEstimator:
+    device: torch.device
+    model: DepthPro
+    model_config: Config
+    to_model: Compose
+
+    def __init__(self, checkpoint: Path, device: torch.device) -> None:
+        self.device = device
+
+        config = Config(checkpoint=checkpoint)
+        self.model_config = config
+        self.model = DepthPro(config, device, torch.half)
+
+        self.to_model = Compose(  # type: ignore[no-untyped-call]
+            [
+                ConvertImageDtype(torch.half),
+                Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5]),  # type: ignore[no-untyped-call]
+            ]
+        )
+
+    def predict(
+        self,
+        frame_batch: UInt8[torch.Tensor, 'batch 3 height width'],
+    ) -> torch.Tensor:
+        *_, height, width = frame_batch.shape
+
+        result = self.model.predict(self.to_model(frame_batch))
+
+        return (
+            Resize((height, width))  # type: ignore[no-untyped-call]
+            .forward(result.depth)
+            .to(torch.float32)
+        )

packages/child-lab-cli/src/child_lab_cli/commands/generate_pointcloud.py

@@ -0,0 +1,36 @@
+from pathlib import Path
+
+import click
+
+
+@click.command('process')
+@click.argument('workspace', type=Path)
+@click.argument('videos', type=Path, nargs=-1)
+@click.option(
+    '--device',
+    type=str,
+    required=False,
+    help='Torch device to use for tensor computations',
+)
+@click.option(
+    '--skip',
+    type=int,
+    required=False,
+    help='Seconds of videos to skip at the beginning',
+)
+@click.option(
+    '--dynamic-transformations',
+    type=bool,
+    is_flag=True,
+    default=False,
+    help='Compute camera transformations on the fly, using heuristic algorithms',
+)
+# @click_trap()
+def process(
+    workspace: Path,
+    videos: list[Path],
+    device: str | None,
+    skip: int | None,
+    dynamic_transformations: bool,
+) -> None:
+    pass

packages/child-lab-cli/src/child_lab_cli/commands/process.py

@@ -0,0 +1,37 @@
+from pathlib import Path
+
+import click
+
+from ..workspace import Workspace
+
+# 1. Load a workspace
+# 3. Load a single video
+# 2. Estimate or load the depth
+# 3. Launch a `viser` server
+# 4. Draw the point clouds
+# 5. Load the transformation buffer
+# 6. Draw those frames of reference which are reachable
+#    from the video's frame of reference
+
+
+@click.command('visualize')
+@click.argument('workspace_root', type=Path)
+@click.argument('video_name', type=str)
+@click.option(
+    '--device',
+    type=str,
+    required=False,
+    help='Torch device to use for tensor computations',
+)
+def visualize(
+    workspace_root: Path,
+    video_name: str,
+    device: str | None,
+) -> None:
+    workspace = Workspace.in_directory(workspace_root)
+
+    for video in workspace.calibrated_videos():
+        if video.name == video_name:
+            break
+    else:
+        raise FileNotFoundError(f'Video {video_name} not found in {workspace.input}.')

packages/child-lab-cli/src/child_lab_cli/commands/visualize.py

@@ -0,0 +1,3 @@
+from .model import CalibratedVideo, NonCalibratedVideo, Workspace, WorkspaceModelError
+
+__all__ = ['CalibratedVideo', 'NonCalibratedVideo', 'Workspace', 'WorkspaceModelError']