733 lines
27 KiB
Python
733 lines
27 KiB
Python
"""Transform command - Apply signal transformations to recordings."""
|
|
|
|
import importlib
|
|
import importlib.util
|
|
import inspect
|
|
import os
|
|
from pathlib import Path
|
|
|
|
import click
|
|
|
|
from utils.data.recording import Recording
|
|
from utils.io.recording import load_recording
|
|
from utils.transforms import iq_augmentations, iq_channel_models, iq_impairments
|
|
from utils_cli.utils.common import (
|
|
echo_progress,
|
|
echo_verbose,
|
|
format_sample_count,
|
|
save_recording,
|
|
)
|
|
|
|
|
|
def get_available_transforms(module):
|
|
"""Get list of public transform functions from a module.
|
|
|
|
Args:
|
|
module: Python module to inspect
|
|
|
|
Returns:
|
|
dict: {name: function} for all public callables
|
|
"""
|
|
transforms = {}
|
|
for name, obj in inspect.getmembers(module, inspect.isfunction):
|
|
if not name.startswith("_"):
|
|
transforms[name] = obj
|
|
return transforms
|
|
|
|
|
|
def get_transform_help(func):
|
|
"""Extract help info from a transform function.
|
|
|
|
Args:
|
|
func: Transform function to inspect
|
|
|
|
Returns:
|
|
dict: {description, params}
|
|
"""
|
|
sig = inspect.signature(func)
|
|
doc = inspect.getdoc(func) or ""
|
|
|
|
# Get first line of docstring as description
|
|
description = doc.split("\n")[0] if doc else "No description"
|
|
|
|
# Extract parameters from signature (skip 'signal')
|
|
params = {}
|
|
for param_name, param in sig.parameters.items():
|
|
if param_name == "signal":
|
|
continue
|
|
|
|
default = param.default
|
|
param_type = "optional" if default != inspect.Parameter.empty else "required"
|
|
default_str = f" (default: {default})" if default != inspect.Parameter.empty else ""
|
|
|
|
params[param_name] = {
|
|
"type": param_type,
|
|
"default": default,
|
|
"annotation": str(param.annotation) if param.annotation != inspect.Parameter.empty else "any",
|
|
"display": f"{param_name} ({param_type}){default_str}",
|
|
}
|
|
|
|
return {"description": description, "full_doc": doc, "params": params}
|
|
|
|
|
|
def show_transform_help(transform_name, func):
|
|
"""Display compact help for a specific transform."""
|
|
info = get_transform_help(func)
|
|
|
|
click.echo(f"\n{transform_name}")
|
|
click.echo("-" * 50)
|
|
click.echo(info["description"])
|
|
|
|
if info["params"]:
|
|
click.echo("\nParameters:")
|
|
for param_name, param_info in sorted(info["params"].items()):
|
|
click.echo(f" {param_name:20} {param_info['display']}")
|
|
|
|
click.echo()
|
|
|
|
|
|
def quick_view_transform(recording, output_path, title="Transform Result"):
|
|
"""Create a quick PNG visualization of transformed recording using constellation plot."""
|
|
try:
|
|
from utils.view.view_signal_simple import view_simple_sig
|
|
|
|
# Create PNG in same directory as output
|
|
output_dir = Path(output_path).parent
|
|
base_name = Path(output_path).stem
|
|
png_path = output_dir / f"{base_name}_preview.png"
|
|
|
|
# Use simple view with constellation
|
|
view_simple_sig(recording, output_path=str(png_path), constellation_mode=True, title=title, saveplot=True)
|
|
|
|
click.echo(f"Visualization saved to: {png_path}")
|
|
except Exception as e:
|
|
click.echo(f"Warning: Could not create visualization: {e}")
|
|
|
|
|
|
def generate_transform_suffix(transform_name, params):
|
|
"""Generate a short suffix for the output filename based on transform and params.
|
|
|
|
Args:
|
|
transform_name: Name of the transform
|
|
params: Dict of parameters
|
|
|
|
Returns:
|
|
str: A short suffix like "awgn15" or "freqoffset10k"
|
|
"""
|
|
suffix = transform_name.replace("_", "")
|
|
|
|
# Add key parameter values
|
|
if "snr_db" in params:
|
|
suffix += f"{int(params['snr_db'])}"
|
|
elif "snr" in params:
|
|
suffix += f"{int(params['snr'])}"
|
|
elif "amplitude_variance" in params:
|
|
suffix += f"{int(params['amplitude_variance']*100)}av"
|
|
elif "phase_variance" in params:
|
|
suffix += f"{int(params['phase_variance']*100000)}pv"
|
|
elif "compression_gain" in params:
|
|
suffix += f"{params['compression_gain']:.2f}".rstrip("0").rstrip(".")
|
|
elif "offset_hz" in params:
|
|
hz = params["offset_hz"]
|
|
if abs(hz) >= 1e6:
|
|
suffix += f"{hz/1e6:.0f}m"
|
|
elif abs(hz) >= 1e3:
|
|
suffix += f"{hz/1e3:.0f}k"
|
|
else:
|
|
suffix += f"{hz:.0f}"
|
|
elif "offset" in params:
|
|
suffix += f"{params['offset']:.2f}".rstrip("0").rstrip(".")
|
|
elif "doppler_hz" in params:
|
|
suffix += f"{params['doppler_hz']:.0f}"
|
|
|
|
return suffix
|
|
|
|
|
|
def parse_transform_params(param_strings):
|
|
"""Parse transform parameters from CLI options.
|
|
|
|
Args:
|
|
param_strings: List of 'KEY=VALUE' strings
|
|
|
|
Returns:
|
|
dict: {key: value} with types inferred
|
|
"""
|
|
params = {}
|
|
if not param_strings:
|
|
return params
|
|
|
|
for param_str in param_strings:
|
|
if "=" not in param_str:
|
|
raise click.BadParameter(f"Parameter must be KEY=VALUE, got: {param_str}")
|
|
|
|
key, value = param_str.split("=", 1)
|
|
key = key.strip()
|
|
value = value.strip()
|
|
|
|
# Try to infer type
|
|
try:
|
|
# Try to parse scientific notation and floats
|
|
if "e" in value.lower() or "." in value:
|
|
params[key] = float(value)
|
|
else:
|
|
params[key] = int(value)
|
|
except ValueError:
|
|
# Keep as string
|
|
params[key] = value
|
|
|
|
return params
|
|
|
|
|
|
def load_custom_transforms(transform_dir):
|
|
"""Load custom transform functions from a directory.
|
|
|
|
Args:
|
|
transform_dir: Path to directory containing .py files with transform functions
|
|
|
|
Returns:
|
|
dict: {transform_name: function} for all public functions in all .py files
|
|
|
|
Raises:
|
|
click.ClickException: If directory doesn't exist or no transforms found
|
|
"""
|
|
transform_dir = Path(transform_dir)
|
|
|
|
if not transform_dir.exists():
|
|
raise click.ClickException(f"Transform directory does not exist: {transform_dir}")
|
|
|
|
if not transform_dir.is_dir():
|
|
raise click.ClickException(f"Path is not a directory: {transform_dir}")
|
|
|
|
transforms = {}
|
|
py_files = list(transform_dir.glob("*.py"))
|
|
|
|
if not py_files:
|
|
raise click.ClickException(f"No .py files found in {transform_dir}")
|
|
|
|
for py_file in py_files:
|
|
try:
|
|
# Load module dynamically
|
|
spec = importlib.util.spec_from_file_location(py_file.stem, py_file)
|
|
if spec is None or spec.loader is None:
|
|
click.echo(f"Warning: Could not load {py_file.name}")
|
|
continue
|
|
|
|
module = importlib.util.module_from_spec(spec)
|
|
spec.loader.exec_module(module)
|
|
|
|
# Extract all public functions
|
|
for name, obj in inspect.getmembers(module, inspect.isfunction):
|
|
if not name.startswith("_"):
|
|
# Store with source file info for metadata
|
|
obj._transform_source_file = py_file.name
|
|
transforms[name] = obj
|
|
except Exception as e:
|
|
raise click.ClickException(f"Failed to load {py_file.name}: {e}")
|
|
|
|
return transforms
|
|
|
|
|
|
def check_input_errors(item_name: str, item, available, input, help_transform):
|
|
if item is None:
|
|
if help_transform:
|
|
raise click.UsageError(f"{item_name.upper()} must be specified for --help-transform")
|
|
else:
|
|
raise click.UsageError(f"{item_name.upper()} must be specified (or use --list)")
|
|
if item not in available:
|
|
raise click.ClickException(f"Unknown {item_name}: {item}\n" f"Use --list to see available options")
|
|
if input is None and not help_transform:
|
|
raise click.UsageError("INPUT must be specified")
|
|
|
|
|
|
def load_input(input, verbose):
|
|
# Load input
|
|
try:
|
|
recording = load_recording(input)
|
|
except Exception as e:
|
|
raise click.ClickException(f"Failed to load input: {e}")
|
|
|
|
echo_verbose(f"Loaded {format_sample_count(recording.data.shape[-1])} samples", verbose)
|
|
return recording
|
|
|
|
@click.group()
|
|
def transform():
|
|
"""Apply signal transformations to recordings.
|
|
|
|
Transform supports three categories of operations:
|
|
- augment: Modify signal to create new ML examples
|
|
- impair: Degrade signal with noise, distortion, etc.
|
|
- apply_channel: Apply channel models (fading, Doppler, etc.)
|
|
|
|
Each operation is applied independently. Chain multiple transforms by
|
|
running this command multiple times.
|
|
|
|
\b
|
|
Examples:
|
|
# List available augmentations
|
|
utils transform augment --list
|
|
\b
|
|
# Apply channel swap
|
|
utils transform augment channel_swap input.npy
|
|
\b
|
|
# Apply AWGN impairment
|
|
utils transform impair awgn input.npy --snr-db 15
|
|
\b
|
|
# Apply Rayleigh fading channel
|
|
utils transform apply_channel rayleigh input.npy --num-paths 5
|
|
"""
|
|
pass
|
|
|
|
|
|
@transform.command(name="augment")
|
|
@click.argument("augmentation", required=False)
|
|
@click.argument("input", type=click.Path(exists=True), required=False)
|
|
@click.argument("output", type=click.Path(), required=False)
|
|
@click.option("--list", "list_transforms", is_flag=True, help="List available augmentations")
|
|
@click.option("--help-transform", is_flag=True, help="Show parameters for this augmentation")
|
|
@click.option("--params", multiple=True, help="Transform parameters as KEY=VALUE (can be repeated)")
|
|
@click.option("--view", is_flag=True, help="Save visualization PNG with constellation plot")
|
|
@click.option("--overwrite", is_flag=True, help="Overwrite output if it exists")
|
|
@click.option("--verbose", "-v", is_flag=True, help="Verbose output")
|
|
@click.option("--quiet", "-q", is_flag=True, help="Suppress output")
|
|
def augment(augmentation, input, output, list_transforms, help_transform, params, view, overwrite, verbose, quiet):
|
|
"""Apply augmentation transforms to recordings.
|
|
|
|
Augmentations modify signals to create new training examples without
|
|
degrading quality (e.g., channel swap, time reversal, quantization).
|
|
|
|
Examples:
|
|
|
|
# List all augmentations
|
|
\b
|
|
utils transform augment --list
|
|
|
|
# Show parameters for an augmentation
|
|
\b
|
|
utils transform augment channel_swap --help-transform
|
|
|
|
# Apply augmentation
|
|
\b
|
|
utils transform augment channel_swap input.npy
|
|
|
|
# Apply with parameters and save visualization
|
|
\b
|
|
utils transform augment drop_samples input.npy --params max_section_size=5 --view
|
|
"""
|
|
available = get_available_transforms(iq_augmentations)
|
|
|
|
if list_transforms:
|
|
click.echo("Available augmentations:")
|
|
for name in sorted(available.keys()):
|
|
func = available[name]
|
|
docstring = (func.__doc__ or "").split("\n")[0].strip()
|
|
click.echo(f" {name:30} {docstring}")
|
|
return
|
|
|
|
if help_transform:
|
|
check_input_errors("augmentation", augmentation, available, input, help_transform)
|
|
show_transform_help(augmentation, available[augmentation])
|
|
return
|
|
|
|
check_input_errors("augmentation", augmentation, available, input, help_transform)
|
|
|
|
# Generate output filename if not provided
|
|
if output is None:
|
|
input_path = Path(input)
|
|
input_stem = input_path.stem
|
|
ext = input_path.suffix
|
|
suffix = generate_transform_suffix(augmentation, parse_transform_params(params))
|
|
output = str(input_path.parent / f"{input_stem}_{suffix}{ext}")
|
|
echo_verbose(f"Auto-generated output: {output}", verbose)
|
|
|
|
# Check if output exists
|
|
if not overwrite and Path(output).exists():
|
|
raise click.ClickException(f"Output file '{output}' already exists\n" f"Use --overwrite to replace")
|
|
|
|
echo_progress(f"Augmenting: {os.path.basename(input)} → {os.path.basename(output)}", quiet)
|
|
echo_verbose(f"Transform: {augmentation}", verbose)
|
|
|
|
# Load input
|
|
recording = load_input(input, verbose)
|
|
|
|
# Parse and apply transform
|
|
try:
|
|
transform_func = available[augmentation]
|
|
transform_params = parse_transform_params(params)
|
|
echo_verbose(f"Parameters: {transform_params}", verbose)
|
|
|
|
result = transform_func(recording, **transform_params)
|
|
except Exception as e:
|
|
raise click.ClickException(f"Transform failed: {e}")
|
|
|
|
# Track transform in metadata (Recording.metadata is a property that returns a copy)
|
|
# So we need to work with a copy and create a new Recording with updated metadata
|
|
updated_metadata = result.metadata.copy()
|
|
if "transforms_applied" not in updated_metadata:
|
|
updated_metadata["transforms_applied"] = []
|
|
|
|
updated_metadata["transforms_applied"].append(
|
|
{"type": "augment", "name": augmentation, "params": parse_transform_params(params)}
|
|
)
|
|
|
|
# Create new recording with updated metadata
|
|
result = Recording(data=result.data, metadata=updated_metadata, annotations=result.annotations)
|
|
|
|
# Save output
|
|
try:
|
|
save_recording(result, output, overwrite=overwrite, verbose=verbose)
|
|
echo_progress(f"Saved to: {output}", quiet)
|
|
except Exception as e:
|
|
raise click.ClickException(f"Failed to save output: {e}")
|
|
|
|
# Optional: Create visualization
|
|
if view:
|
|
echo_verbose("Creating visualization...", verbose)
|
|
quick_view_transform(result, output, title=f"{augmentation.replace('_', ' ').title()} - {Path(output).name}")
|
|
|
|
|
|
@transform.command(name="impair")
|
|
@click.argument("impairment", required=False)
|
|
@click.argument("input", type=click.Path(exists=True), required=False)
|
|
@click.argument("output", type=click.Path(), required=False)
|
|
@click.option("--list", "list_transforms", is_flag=True, help="List available impairments")
|
|
@click.option("--help-transform", is_flag=True, help="Show parameters for this impairment")
|
|
@click.option("--params", multiple=True, help="Transform parameters as KEY=VALUE (can be repeated)")
|
|
@click.option("--view", is_flag=True, help="Save visualization PNG with constellation plot")
|
|
@click.option("--overwrite", is_flag=True, help="Overwrite output if it exists")
|
|
@click.option("--verbose", "-v", is_flag=True, help="Verbose output")
|
|
@click.option("--quiet", "-q", is_flag=True, help="Suppress output")
|
|
def impair(impairment, input, output, list_transforms, help_transform, params, view, overwrite, verbose, quiet):
|
|
"""Apply impairment transforms to recordings.
|
|
|
|
Impairments degrade signals by adding noise, distortion, and other
|
|
channel effects (e.g., AWGN, phase noise, IQ imbalance).
|
|
|
|
Examples:
|
|
|
|
# List all impairments
|
|
\b
|
|
utils transform impair --list
|
|
|
|
# Show parameters for an impairment
|
|
\b
|
|
utils transform impair add_awgn_to_signal --help-transform
|
|
|
|
# Apply impairment
|
|
\b
|
|
utils transform impair add_awgn_to_signal input.npy --params snr=10
|
|
|
|
# Apply with visualization
|
|
\b
|
|
utils transform impair add_phase_noise input.npy --params phase_variance=0.001 --view
|
|
"""
|
|
available = get_available_transforms(iq_impairments)
|
|
|
|
if list_transforms:
|
|
click.echo("Available impairments:")
|
|
for name in sorted(available.keys()):
|
|
func = available[name]
|
|
docstring = (func.__doc__ or "").split("\n")[0].strip()
|
|
click.echo(f" {name:30} {docstring}")
|
|
return
|
|
|
|
if help_transform:
|
|
check_input_errors("impairment", impairment, available, input, help_transform)
|
|
show_transform_help(impairment, available[impairment])
|
|
return
|
|
|
|
check_input_errors("impairment", impairment, available, input, help_transform)
|
|
|
|
# Generate output filename if not provided
|
|
if output is None:
|
|
input_path = Path(input)
|
|
input_stem = input_path.stem
|
|
ext = input_path.suffix
|
|
suffix = generate_transform_suffix(impairment, parse_transform_params(params))
|
|
output = str(input_path.parent / f"{input_stem}_{suffix}{ext}")
|
|
echo_verbose(f"Auto-generated output: {output}", verbose)
|
|
|
|
# Check if output exists
|
|
if not overwrite and Path(output).exists():
|
|
raise click.ClickException(f"Output file '{output}' already exists\n" f"Use --overwrite to replace")
|
|
|
|
echo_progress(f"Impairing: {os.path.basename(input)} → {os.path.basename(output)}", quiet)
|
|
echo_verbose(f"Transform: {impairment}", verbose)
|
|
|
|
# Load input
|
|
recording = load_input(input, verbose)
|
|
|
|
# Parse and apply transform
|
|
try:
|
|
transform_func = available[impairment]
|
|
transform_params = parse_transform_params(params)
|
|
echo_verbose(f"Parameters: {transform_params}", verbose)
|
|
|
|
result = transform_func(recording, **transform_params)
|
|
except Exception as e:
|
|
raise click.ClickException(f"Transform failed: {e}")
|
|
|
|
# Track transform in metadata (Recording.metadata is a property that returns a copy)
|
|
updated_metadata = result.metadata.copy()
|
|
if "transforms_applied" not in updated_metadata:
|
|
updated_metadata["transforms_applied"] = []
|
|
|
|
updated_metadata["transforms_applied"].append(
|
|
{"type": "impair", "name": impairment, "params": parse_transform_params(params)}
|
|
)
|
|
|
|
# Create new recording with updated metadata
|
|
result = Recording(data=result.data, metadata=updated_metadata, annotations=result.annotations)
|
|
|
|
# Save output
|
|
try:
|
|
save_recording(result, output, overwrite=overwrite, verbose=verbose)
|
|
echo_progress(f"Saved to: {output}", quiet)
|
|
except Exception as e:
|
|
raise click.ClickException(f"Failed to save output: {e}")
|
|
|
|
# Optional: Create visualization
|
|
if view:
|
|
echo_verbose("Creating visualization...", verbose)
|
|
quick_view_transform(result, output, title=f"{impairment.replace('_', ' ').title()} - {Path(output).name}")
|
|
|
|
|
|
@transform.command(name="apply_channel")
|
|
@click.argument("channel_model", required=False)
|
|
@click.argument("input", type=click.Path(exists=True), required=False)
|
|
@click.argument("output", type=click.Path(), required=False)
|
|
@click.option("--list", "list_transforms", is_flag=True, help="List available channel models")
|
|
@click.option("--help-transform", is_flag=True, help="Show parameters for this channel model")
|
|
@click.option("--params", multiple=True, help="Transform parameters as KEY=VALUE (can be repeated)")
|
|
@click.option("--view", is_flag=True, help="Save visualization PNG with constellation plot")
|
|
@click.option("--overwrite", is_flag=True, help="Overwrite output if it exists")
|
|
@click.option("--verbose", "-v", is_flag=True, help="Verbose output")
|
|
@click.option("--quiet", "-q", is_flag=True, help="Suppress output")
|
|
def apply_channel(
|
|
channel_model, input, output, list_transforms, help_transform, params, view, overwrite, verbose, quiet
|
|
):
|
|
"""Apply channel models to recordings.
|
|
|
|
Channel models simulate RF propagation effects like fading, Doppler shift,
|
|
and multipath reflections.
|
|
|
|
Use --list to see available channel models and their parameters.
|
|
|
|
\b
|
|
Examples:
|
|
utils transform apply_channel rayleigh_fading_channel input.npy --params num_paths=3 snr_db=15
|
|
|
|
\b
|
|
utils transform apply_channel doppler_channel recordings/input.npy \\
|
|
--params satellite_velocity=7500 \\
|
|
--params satellite_initial_distance=400000 \\
|
|
--params frequency=1e9 \\
|
|
--params sample_rate=2e6
|
|
"""
|
|
available = get_available_transforms(iq_channel_models)
|
|
|
|
if list_transforms:
|
|
click.echo("Available channel models:")
|
|
for name in sorted(available.keys()):
|
|
func = available[name]
|
|
docstring = (func.__doc__ or "").split("\n")[0].strip()
|
|
click.echo(f" {name:30} {docstring}")
|
|
return
|
|
|
|
if help_transform:
|
|
check_input_errors("channel_model", channel_model, available, input, help_transform)
|
|
show_transform_help(channel_model, available[channel_model])
|
|
return
|
|
|
|
check_input_errors("channel_model", channel_model, available, input, help_transform)
|
|
|
|
# Generate output filename if not provided
|
|
if output is None:
|
|
input_path = Path(input)
|
|
input_stem = input_path.stem
|
|
ext = input_path.suffix
|
|
suffix = generate_transform_suffix(channel_model, parse_transform_params(params))
|
|
output = str(input_path.parent / f"{input_stem}_{suffix}{ext}")
|
|
echo_verbose(f"Auto-generated output: {output}", verbose)
|
|
|
|
# Check if output exists
|
|
if not overwrite and Path(output).exists():
|
|
raise click.ClickException(f"Output file '{output}' already exists\n" f"Use --overwrite to replace")
|
|
|
|
echo_progress(f"Applying channel: {os.path.basename(input)} → {os.path.basename(output)}", quiet)
|
|
echo_verbose(f"Channel model: {channel_model}", verbose)
|
|
|
|
# Load input
|
|
recording = load_input(input, verbose)
|
|
|
|
# Parse and apply transform
|
|
try:
|
|
transform_func = available[channel_model]
|
|
transform_params = parse_transform_params(params)
|
|
echo_verbose(f"Parameters: {transform_params}", verbose)
|
|
|
|
result = transform_func(recording, **transform_params)
|
|
except Exception as e:
|
|
raise click.ClickException(f"Transform failed: {e}")
|
|
|
|
# Track transform in metadata (Recording.metadata is a property that returns a copy)
|
|
updated_metadata = result.metadata.copy()
|
|
if "transforms_applied" not in updated_metadata:
|
|
updated_metadata["transforms_applied"] = []
|
|
|
|
updated_metadata["transforms_applied"].append(
|
|
{"type": "channel", "name": channel_model, "params": parse_transform_params(params)}
|
|
)
|
|
|
|
# Create new recording with updated metadata
|
|
result = Recording(data=result.data, metadata=updated_metadata, annotations=result.annotations)
|
|
|
|
# Save output
|
|
try:
|
|
save_recording(result, output, overwrite=overwrite, verbose=verbose)
|
|
echo_progress(f"Saved to: {output}", quiet)
|
|
except Exception as e:
|
|
raise click.ClickException(f"Failed to save output: {e}")
|
|
|
|
# Optional: Create visualization
|
|
if view:
|
|
echo_verbose("Creating visualization...", verbose)
|
|
quick_view_transform(result, output, title=f"{channel_model.replace('_', ' ').title()} - {Path(output).name}")
|
|
|
|
|
|
@transform.command(name="custom")
|
|
@click.argument("transform_name", required=False)
|
|
@click.argument("input", type=click.Path(exists=True), required=False)
|
|
@click.argument("output", type=click.Path(), required=False)
|
|
@click.option(
|
|
"--transform-dir",
|
|
type=click.Path(exists=True),
|
|
required=True,
|
|
help="Path to directory containing custom transform .py files",
|
|
)
|
|
@click.option("--list", "list_transforms", is_flag=True, help="List available custom transforms")
|
|
@click.option("--help-transform", is_flag=True, help="Show parameters for this transform")
|
|
@click.option("--params", multiple=True, help="Transform parameters as KEY=VALUE (can be repeated)")
|
|
@click.option("--view", is_flag=True, help="Save visualization PNG with constellation plot")
|
|
@click.option("--overwrite", is_flag=True, help="Overwrite output if it exists")
|
|
@click.option("--verbose", "-v", is_flag=True, help="Verbose output")
|
|
@click.option("--quiet", "-q", is_flag=True, help="Suppress output")
|
|
def custom(
|
|
transform_name,
|
|
input,
|
|
output,
|
|
transform_dir,
|
|
list_transforms,
|
|
help_transform,
|
|
params,
|
|
view,
|
|
overwrite,
|
|
verbose,
|
|
quiet,
|
|
):
|
|
"""Apply custom user-defined transforms to recordings.
|
|
|
|
Custom transforms are Python functions loaded from user-specified directory.
|
|
Each .py file in the directory is scanned for public functions that can be used.
|
|
|
|
Transform functions must have signature:
|
|
def my_transform(signal, **kwargs) -> signal_or_recording
|
|
where signal is a complex CxN array or Recording object.
|
|
|
|
Examples:
|
|
|
|
# List all custom transforms in directory
|
|
\b
|
|
utils transform custom --transform-dir ~/my_transforms --list
|
|
|
|
# Show parameters for a transform
|
|
\b
|
|
utils transform custom my_filter --transform-dir ~/my_transforms --help-transform
|
|
|
|
# Apply custom transform
|
|
\b
|
|
utils transform custom my_filter input.npy --transform-dir ~/my_transforms
|
|
|
|
# With parameters and visualization
|
|
\b
|
|
utils transform custom my_filter input.npy --transform-dir ~/my_transforms \\
|
|
--params cutoff_freq=5000 order=4 --view
|
|
"""
|
|
try:
|
|
available = load_custom_transforms(transform_dir)
|
|
except click.ClickException:
|
|
raise
|
|
|
|
if list_transforms:
|
|
click.echo(f"Available custom transforms in {transform_dir}:")
|
|
for name in sorted(available.keys()):
|
|
func = available[name]
|
|
source_file = getattr(func, "_transform_source_file", "unknown")
|
|
docstring = (func.__doc__ or "").split("\n")[0].strip()
|
|
click.echo(f" {name:30} {docstring:40} [{source_file}]")
|
|
return
|
|
|
|
if help_transform:
|
|
check_input_errors("transform_name", transform_name, available, input, help_transform)
|
|
show_transform_help(transform_name, available[transform_name])
|
|
return
|
|
|
|
check_input_errors("transform_name", transform_name, available, input, help_transform)
|
|
|
|
# Generate output filename if not provided
|
|
if output is None:
|
|
input_path = Path(input)
|
|
input_stem = input_path.stem
|
|
ext = input_path.suffix
|
|
suffix = generate_transform_suffix(transform_name, parse_transform_params(params))
|
|
output = str(input_path.parent / f"{input_stem}_{suffix}{ext}")
|
|
echo_verbose(f"Auto-generated output: {output}", verbose)
|
|
|
|
# Check if output exists
|
|
if not overwrite and Path(output).exists():
|
|
raise click.ClickException(f"Output file '{output}' already exists\n" f"Use --overwrite to replace")
|
|
|
|
echo_progress(f"Applying custom: {os.path.basename(input)} → {os.path.basename(output)}", quiet)
|
|
echo_verbose(f"Transform: {transform_name}", verbose)
|
|
|
|
# Load input
|
|
recording = load_input(input, verbose)
|
|
|
|
# Parse and apply transform
|
|
try:
|
|
transform_func = available[transform_name]
|
|
transform_params = parse_transform_params(params)
|
|
echo_verbose(f"Parameters: {transform_params}", verbose)
|
|
|
|
result = transform_func(recording, **transform_params)
|
|
except Exception as e:
|
|
raise click.ClickException(f"Transform failed: {e}")
|
|
|
|
# Track transform in metadata
|
|
updated_metadata = result.metadata.copy()
|
|
if "transforms_applied" not in updated_metadata:
|
|
updated_metadata["transforms_applied"] = []
|
|
|
|
updated_metadata["transforms_applied"].append(
|
|
{
|
|
"type": "custom",
|
|
"name": transform_name,
|
|
"source_file": getattr(available[transform_name], "_transform_source_file", "unknown"),
|
|
"params": parse_transform_params(params),
|
|
}
|
|
)
|
|
|
|
# Create new recording with updated metadata
|
|
result = Recording(data=result.data, metadata=updated_metadata, annotations=result.annotations)
|
|
|
|
# Save output
|
|
try:
|
|
save_recording(result, output, overwrite=overwrite, verbose=verbose)
|
|
echo_progress(f"Saved to: {output}", quiet)
|
|
except Exception as e:
|
|
raise click.ClickException(f"Failed to save output: {e}")
|
|
|
|
# Optional: Create visualization
|
|
if view:
|
|
echo_verbose("Creating visualization...", verbose)
|
|
quick_view_transform(result, output, title=f"{transform_name.replace('_', ' ').title()} - {Path(output).name}")
|