s6.app.cog package¶

Submodules¶

s6.app.cog.augmented_dataset module¶

Augmented keypoint dataset and image transforms.

Wraps s6.utils.datapipe.StructuredDatasetTorch to produce square patches around a single keypoint with optional augmentation (mirror, crop, rotation, color jitter, blur, occlusion) and an optional distance-based sampling scheme. Helpers provide deterministic transformations used by the dataset implementation.

Examples

Instantiate from a JSON config and fetch a sample:

>>> from s6.app.cog.augmented_dataset import AugmentedKeypointDataset
>>> ds = AugmentedKeypointDataset.from_json("dataset.cfg.json")
>>> image, keypoint = ds[0]  # image: (3, H, W) in [0,1), keypoint: (2,)

See also

s6.app.cog.keypoint

class s6.app.cog.augmented_dataset.AugmentedKeypointDataset(config: Config)¶

Bases: Dataset

Dataset that applies keypoint-centric augmentations.

Wraps s6.utils.datapipe.StructuredDatasetTorch and returns RGB tensors and a single keypoint in normalized coordinates. Supports square center crop, random crop+resize, rotation, color jitter, Gaussian blur, and occlusion. Optionally balances sampling by keypoint distance to image center using bucketized bins.

Parameters:: config (AugmentedKeypointDataset.Config) – Dataset roots, datakeys, and augmentation/sampling options.

class Config(*, base_dir: str | List[str] = [], datakeys: List[str] = [], mirror: MirrorConfig = MirrorConfig(enabled=False), crop: CropConfig = CropConfig(enabled=True, crop_factor=0.9, output_size=256), rotation: RotationConfig = RotationConfig(enabled=True, max_rotation=360.0), blur: BlurConfig = BlurConfig(enabled=False, sigma_range=(0.0, 2.0)), occlusion: OcclusionConfig = OcclusionConfig(enabled=False, size_range=(0.1, 0.3), intensity_range=(0, 255)), color_jitter: ColorJitterConfig = ColorJitterConfig(brightness=0.2, contrast=0.2, gamma=0.2), sd_kwargs: dict = {}, sampling: SamplingConfig = SamplingConfig(enabled=False, bins=10, seed=0))¶

Bases: BaseModel

base_dir: str | List[str]¶

blur: BlurConfig¶

color_jitter: ColorJitterConfig¶

crop: CropConfig¶

datakeys: List[str]¶

classmethod default()¶

mirror: MirrorConfig¶

occlusion: OcclusionConfig¶

rotation: RotationConfig¶

sampling: SamplingConfig¶

sd_kwargs: dict¶

classmethod from_json(cfg_path: str) → AugmentedKeypointDataset¶

Instantiate dataset from a JSON config file.

Parameters:: cfg_path (str) – Path to a JSON file accepted by AugmentedKeypointDataset.Config.
Returns:: Dataset instance configured as specified in the JSON.
Return type:: AugmentedKeypointDataset

property output_size: int¶

Square output size in pixels for augmented patches.

Returns the configured crop output size so callers can allocate tensors like (N, C, output_size, output_size) without reaching into config.

class s6.app.cog.augmented_dataset.BlurConfig(*, enabled: bool = False, sigma_range: Tuple[float, float] = (0.0, 2.0))¶

Bases: BaseModel

enabled: bool¶

sigma_range: Tuple[float, float]¶

class s6.app.cog.augmented_dataset.ColorJitterConfig(*, brightness: float = 0.2, contrast: float = 0.2, gamma: float = 0.2)¶

Bases: BaseModel

brightness: float¶

contrast: float¶

gamma: float¶

class s6.app.cog.augmented_dataset.CropConfig(*, enabled: bool = True, crop_factor: float = 0.9, output_size: int = 256)¶

Bases: BaseModel

crop_factor: float¶

enabled: bool¶

output_size: int¶

class s6.app.cog.augmented_dataset.MirrorConfig(*, enabled: bool = False)¶

Bases: BaseModel

enabled: bool¶

class s6.app.cog.augmented_dataset.OcclusionConfig(*, enabled: bool = False, size_range: Tuple[float, float] = (0.1, 0.3), intensity_range: Tuple[int, int] = (0, 255))¶

Bases: BaseModel

enabled: bool¶

intensity_range: Tuple[int, int]¶

size_range: Tuple[float, float]¶

class s6.app.cog.augmented_dataset.RotationConfig(*, enabled: bool = True, max_rotation: float = 360.0)¶

Bases: BaseModel

enabled: bool¶

max_rotation: float¶

class s6.app.cog.augmented_dataset.SamplingConfig(*, enabled: bool = False, bins: int = 10, seed: int = 0)¶

Bases: BaseModel

Configure distance-based bucket sampling of keypoints.

bins: int¶

enabled: bool¶

seed: int¶

s6.app.cog.augmented_dataset.center_square_crop(image: ndarray, keypoint: Tuple[int, int]) → Tuple[ndarray, ndarray]¶

Crop the largest centered square and remap keypoint.

Parameters:

image (ndarray) – Source image (H, W[, C]) in BGR order.
keypoint (tuple of int) – Keypoint pixel coordinates (x, y).

Returns:

Cropped image and keypoint in the cropped coordinate frame.

Return type:

tuple[ndarray, ndarray]

s6.app.cog.augmented_dataset.horizontal_mirror(image: ndarray, keypoint: Tuple[int, int]) → Tuple[ndarray, ndarray]¶

Mirror image and keypoint horizontally around the vertical axis.

Parameters:

image (ndarray) – Input image in BGR order.
keypoint (tuple of int) – Keypoint pixel coordinates (x, y).

Returns:

Mirrored image and transformed keypoint with x' = (w - 1) - x.

Return type:

tuple[ndarray, ndarray]

s6.app.cog.augmented_dataset.random_blur(image: ndarray, sigma_range: Tuple[float, float]) → ndarray¶

Apply random Gaussian blur to the image.

Parameters:

image (ndarray) – Input image in BGR order.
sigma_range (tuple of float) – Lower/upper bounds for the Gaussian sigma.

Returns:

Blurred image with the same shape and dtype as input.

Return type:

ndarray

s6.app.cog.augmented_dataset.random_color_jitter(image: ndarray, brightness: float = 0.2, contrast: float = 0.2, gamma: float = 0.2) → ndarray¶

Apply brightness/contrast/gamma jitter in-place.

Parameters:

image (ndarray) – Input image in BGR order.
brightness (float, optional) – Multiplicative brightness jitter amplitude (±). Defaults to 0.2.
contrast (float, optional) – Contrast scaling around 0.5 gray level. Defaults to 0.2.
gamma (float, optional) – Exponent applied to normalized intensities. Defaults to 0.2.

Returns:

Jittered image in the same shape and dtype uint8.

Return type:

ndarray

s6.app.cog.augmented_dataset.random_crop_and_resize(image: ndarray, keypoint: Tuple[int, int], crop_factor: float = 0.8, output_size: int = 256) → Tuple[ndarray, ndarray]¶

Randomly crop around keypoint and resize to a square.

Parameters:

image (ndarray) – Input image in BGR order.
keypoint (tuple of int) – Keypoint location in pixels.
crop_factor (float, optional) – Fraction of the shorter side to use for the crop. Defaults to 0.8.
output_size (int, optional) – Target square size in pixels. Defaults to 256.

Returns:

(patch, keypoint) in resized coordinates.

Return type:

tuple[ndarray, ndarray]

s6.app.cog.augmented_dataset.random_occlusion(image: ndarray, keypoint: ndarray, size_range: Tuple[float, float], intensity_range: Tuple[int, int]) → ndarray¶

Overlay a rotated rectangle occlusion while avoiding the keypoint.

Parameters:

image (ndarray) – Input image in BGR order.
keypoint (ndarray) – Keypoint in pixel coordinates.
size_range (tuple of float) – Fractional side length relative to the shorter image side.
intensity_range (tuple of int) – Intensity range for the occluder patch (0–255).

Returns:

Image with an occlusion applied or unchanged if placement failed.

Return type:

ndarray

s6.app.cog.augmented_dataset.random_rotate(image: ndarray, keypoint: ndarray, max_angle: float = 30.0) → Tuple[ndarray, ndarray]¶

Rotate image and keypoint by a random angle.

Parameters:

image (ndarray) – Input image in BGR order.
keypoint (ndarray) – Keypoint (2,) in pixel coordinates.
max_angle (float, optional) – Maximum absolute rotation angle in degrees. Defaults to 30.

Returns:

Rotated image and transformed keypoint.

Return type:

tuple[ndarray, ndarray]

s6.app.cog.benchmark module¶

Benchmark ONNX model inference and compare results.

Loads one or more ONNX models with onnxruntime, generates synthetic inputs matching model signatures, runs N timed inferences, and prints a side‑by‑side comparison of latency statistics.

Examples

Benchmark two models for 200 runs on CPU:

>>> python -m s6.app.cog.benchmark ./m1.onnx ./m2.onnx -n 200

s6.app.cog.benchmark.generate_inputs(session)¶

Generate random inputs that match the model’s input signatures.

Replaces dynamic axes with 1 to keep inputs concrete.

Parameters:: session (onnxruntime.InferenceSession) – ONNXRuntime session to inspect for input metadata.
Returns:: Mapping from input names to random arrays with compatible dtypes and shapes.
Return type:: dict[str, np.ndarray]

s6.app.cog.benchmark.main()¶: Run the benchmark and print side‑by‑side latency statistics.

s6.app.cog.benchmark.parse_args()¶

Parse CLI arguments for the ONNX benchmark script.

Returns:: Contains models (list of paths) and num_runs.
Return type:: argparse.Namespace

s6.app.cog.dataloader module¶

Visualize augmented keypoint patches with a DataLoader.

Creates an AugmentedKeypointDataset, wraps it in a torch.utils.data.DataLoader, pulls one batch, and renders the augmented patch with its keypoint overlaid. Useful for sanity‑checking dataset configuration and augmentations.

Examples

Preview a batch of size 8 from a StructuredDataset directory:

>>> python -m s6.app.cog.dataloader ./data/my_dataset ...   --datakeys B.image B.tip_point --batch-size 8 --output-size 256

s6.app.cog.dataloader.main()¶: Parse arguments, build dataset/loader, and display one batch.

s6.app.cog.keypoint module¶

Keypoint training, preview, and ONNX export utilities.

Provides a lightweight training loop, preview utilities, and ONNX export for a generic keypoint model operating on single‑channel input. Dataset augmentation is provided by s6.app.cog.augmented_dataset.AugmentedKeypointDataset.

Key Concepts¶

Input tensors have shape (N, 1, H, W); keypoints are normalized to [0, 1) with shape (N, K, 2).
Loss uses Smooth L1 on normalized keypoints.
Checkpoints are stored under checkpoints/ and can be restored.
ONNX export exposes outputs named keypoints and heatmaps.

Examples

Preview a batch of augmented patches:

>>> python -m s6.app.cog.keypoint --config ./cfg/ds.json --preview-data

Run a dry‑run (graph + one step):

>>> python -m s6.app.cog.keypoint --config ./cfg/ds.json --dry-run

Train for 50 epochs and export to ONNX:

>>> python -m s6.app.cog.keypoint --config ./cfg/ds.json --train ...   -e 50 -b 16 -lr 1e-3 --deploy model.onnx

s6.app.cog.keypoint.batch_draw_keypoints(images, keypoints)¶

Render keypoints on a batch of images for visualization.

Parameters:

images (torch.Tensor) – Image batch (B, C, H, W) in [0, 1).
keypoints (torch.Tensor) – Normalized keypoints (B, 2) or (B, K, 2).

Returns:

RGB images with circles drawn, shape (B, 3, H, W).

Return type:

torch.Tensor

s6.app.cog.keypoint.deploy_model(args)¶

Load a checkpoint and export the model to ONNX.

Parameters:: args (argparse.Namespace) – Includes deploy output path and restore checkpoint id or path. Uses dataset output_size to determine spatial dims.

s6.app.cog.keypoint.log_graph(writer, model, dataloader, device)¶

Log the model graph to TensorBoard.

Parameters:

writer (SummaryWriter) – TensorBoard writer instance.
model (torch.nn.Module) – Model to trace with a sample batch.
dataloader (DataLoader) – Loader providing sample images.
device (torch.device) – Device where the model and sample are placed.

s6.app.cog.keypoint.main()¶: Parse arguments and execute selected keypoint workflow.

s6.app.cog.keypoint.make_train_dataloader(args)¶

Create dataset and DataLoader for training.

Parameters:: args (argparse.Namespace) – Parsed CLI arguments containing config, batch_size and num_workers.
Returns:: (dataset, dataloader) ready for training.
Return type:: tuple

s6.app.cog.keypoint.preview_dataset(args)¶

Preview a batch with ground‑truth keypoints overlaid.

Parameters:: args (argparse.Namespace) – Must include config, batch_size, and num_workers.

s6.app.cog.keypoint.train(args)¶

Train the keypoint detection model.

Parameters:: args (argparse.Namespace) – Training configuration including epochs, batch_size, learning_rate, logging options, and optional restore.

s6.app.cog.keypoint.train_step(images, keypoints, model, optimizer, criterion, device)¶

Run a single optimization step.

Parameters:

images (torch.Tensor) – Batch of images (B, C, H, W).
keypoints (torch.Tensor) – Ground‑truth keypoints (B, K, 2) normalized to [0, 1).
model (torch.nn.Module) – Keypoint model.
optimizer (torch.optim.Optimizer) – Optimizer for model parameters.
criterion (Callable) – Loss between predicted and ground‑truth keypoints.
device (torch.device) – Compute device.

Returns:

(loss, outputs, images, keypoints) with outputs shaped like keypoints.

Return type:

tuple

s6.app.cog.yolo module¶

YOLOv11 quickstart: image inference or fine‑tuning.

Thin wrapper around ultralytics to either run single‑image inference or fine‑tune a YOLOv11 model on a custom dataset.

Examples

Run inference and save an annotated image:

>>> python -m s6.app.cog.yolo -m yolov11n.pt -s input.jpg -o output.jpg

Fine‑tune on a dataset config for 50 epochs:

>>> python -m s6.app.cog.yolo -m yolov11n.pt --train -d data.yaml --epochs 50

s6.app.cog.yolo.main()¶

Load YOLO model and run the selected mode.

Runs inference when --source is provided without --train; otherwise launches fine‑tuning using --data and training hyperparameters.

s6.app.cog.yolo.parse_args()¶

Parse CLI arguments for YOLO inference/training.

Returns:: Parsed arguments including model path and mode flags.
Return type:: argparse.Namespace

s6.app.cog package¶

Submodules¶

s6.app.cog.augmented_dataset module¶

s6.app.cog.benchmark module¶

s6.app.cog.dataloader module¶

s6.app.cog.keypoint module¶

Key Concepts¶

s6.app.cog.yolo module¶

Module contents¶