chore: update demo runtime, tests, and agent docs
This commit is contained in:
@@ -0,0 +1,394 @@
|
||||
#!/usr/bin/env python3
|
||||
"""
|
||||
Export all positive labeled batches from Scoliosis1K dataset as time windows.
|
||||
|
||||
Creates grid visualizations similar to visualizer._prepare_segmentation_input_view()
|
||||
for all positive class samples, arranged in sliding time windows.
|
||||
|
||||
Optimized UI with:
|
||||
- Subject ID and batch info footer
|
||||
- Dual frame counts (window-relative and sequence-relative)
|
||||
- Clean layout with proper spacing
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import json
|
||||
import pickle
|
||||
from pathlib import Path
|
||||
from typing import Final
|
||||
|
||||
import cv2
|
||||
import numpy as np
|
||||
from numpy.typing import NDArray
|
||||
|
||||
# Constants matching visualizer.py
|
||||
DISPLAY_HEIGHT: Final = 256
|
||||
DISPLAY_WIDTH: Final = 176
|
||||
SIL_HEIGHT: Final = 64
|
||||
SIL_WIDTH: Final = 44
|
||||
|
||||
# Footer settings
|
||||
FOOTER_HEIGHT: Final = 80 # Height for metadata footer
|
||||
FOOTER_BG_COLOR: Final = (40, 40, 40) # Dark gray background
|
||||
TEXT_COLOR: Final = (255, 255, 255) # White text
|
||||
ACCENT_COLOR: Final = (0, 255, 255) # Cyan for emphasis
|
||||
FONT: Final = cv2.FONT_HERSHEY_SIMPLEX
|
||||
FONT_SCALE: Final = 0.6
|
||||
FONT_THICKNESS: Final = 2
|
||||
|
||||
|
||||
def upscale_silhouette(
|
||||
silhouette: NDArray[np.float32] | NDArray[np.uint8],
|
||||
) -> NDArray[np.uint8]:
|
||||
"""Upscale silhouette to display size."""
|
||||
if silhouette.dtype == np.float32 or silhouette.dtype == np.float64:
|
||||
sil_u8 = (silhouette * 255).astype(np.uint8)
|
||||
else:
|
||||
sil_u8 = silhouette.astype(np.uint8)
|
||||
upscaled = cv2.resize(
|
||||
sil_u8, (DISPLAY_WIDTH, DISPLAY_HEIGHT), interpolation=cv2.INTER_NEAREST
|
||||
)
|
||||
return upscaled
|
||||
|
||||
|
||||
def create_optimized_visualization(
|
||||
silhouettes: NDArray[np.float32],
|
||||
subject_id: str,
|
||||
view_name: str,
|
||||
window_idx: int,
|
||||
start_frame: int,
|
||||
end_frame: int,
|
||||
n_frames_total: int,
|
||||
tile_height: int = DISPLAY_HEIGHT,
|
||||
tile_width: int = DISPLAY_WIDTH,
|
||||
) -> NDArray[np.uint8]:
|
||||
"""
|
||||
Create optimized visualization with grid and metadata footer.
|
||||
|
||||
Args:
|
||||
silhouettes: Array of shape (n_frames, 64, 44) float32
|
||||
subject_id: Subject identifier
|
||||
view_name: View identifier (e.g., "000_180")
|
||||
window_idx: Window index within sequence
|
||||
start_frame: Starting frame index in sequence
|
||||
end_frame: Ending frame index in sequence
|
||||
n_frames_total: Total frames in the sequence
|
||||
tile_height: Height of each tile in the grid
|
||||
tile_width: Width of each tile in the grid
|
||||
|
||||
Returns:
|
||||
Combined image with grid visualization and metadata footer
|
||||
"""
|
||||
n_frames = int(silhouettes.shape[0])
|
||||
tiles_per_row = int(np.ceil(np.sqrt(n_frames)))
|
||||
rows = int(np.ceil(n_frames / tiles_per_row))
|
||||
|
||||
# Create grid
|
||||
grid = np.zeros((rows * tile_height, tiles_per_row * tile_width), dtype=np.uint8)
|
||||
|
||||
# Place each silhouette in the grid
|
||||
for idx in range(n_frames):
|
||||
sil = silhouettes[idx]
|
||||
tile = upscale_silhouette(sil)
|
||||
r = idx // tiles_per_row
|
||||
c = idx % tiles_per_row
|
||||
y0, y1 = r * tile_height, (r + 1) * tile_height
|
||||
x0, x1 = c * tile_width, (c + 1) * tile_width
|
||||
grid[y0:y1, x0:x1] = tile
|
||||
|
||||
# Convert to BGR
|
||||
grid_bgr = cv2.cvtColor(grid, cv2.COLOR_GRAY2BGR)
|
||||
|
||||
# Add frame indices as text (both window-relative and sequence-relative)
|
||||
for idx in range(n_frames):
|
||||
r = idx // tiles_per_row
|
||||
c = idx % tiles_per_row
|
||||
y0 = r * tile_height
|
||||
x0 = c * tile_width
|
||||
|
||||
# Window frame count (top-left)
|
||||
cv2.putText(
|
||||
grid_bgr,
|
||||
f"{idx}", # Window-relative frame number
|
||||
(x0 + 8, y0 + 22),
|
||||
FONT,
|
||||
FONT_SCALE,
|
||||
ACCENT_COLOR,
|
||||
FONT_THICKNESS,
|
||||
cv2.LINE_AA,
|
||||
)
|
||||
|
||||
# Sequence frame count (bottom-left of tile)
|
||||
seq_frame = start_frame + idx
|
||||
cv2.putText(
|
||||
grid_bgr,
|
||||
f"#{seq_frame}", # Sequence-relative frame number
|
||||
(x0 + 8, y0 + tile_height - 10),
|
||||
FONT,
|
||||
0.45, # Slightly smaller font
|
||||
(180, 180, 180), # Light gray
|
||||
1,
|
||||
cv2.LINE_AA,
|
||||
)
|
||||
|
||||
# Create footer with metadata
|
||||
grid_width = grid_bgr.shape[1]
|
||||
footer = np.full((FOOTER_HEIGHT, grid_width, 3), FOOTER_BG_COLOR, dtype=np.uint8)
|
||||
|
||||
# Line 1: Subject ID and view
|
||||
line1 = f"Subject: {subject_id} | View: {view_name}"
|
||||
cv2.putText(
|
||||
footer,
|
||||
line1,
|
||||
(15, 25),
|
||||
FONT,
|
||||
0.7,
|
||||
TEXT_COLOR,
|
||||
FONT_THICKNESS,
|
||||
cv2.LINE_AA,
|
||||
)
|
||||
|
||||
# Line 2: Window batch frame range
|
||||
line2 = f"Window {window_idx}: frames [{start_frame:03d} - {end_frame - 1:03d}] ({n_frames} frames)"
|
||||
cv2.putText(
|
||||
footer,
|
||||
line2,
|
||||
(15, 50),
|
||||
FONT,
|
||||
0.7,
|
||||
ACCENT_COLOR,
|
||||
FONT_THICKNESS,
|
||||
cv2.LINE_AA,
|
||||
)
|
||||
|
||||
# Line 3: Progress within sequence
|
||||
progress_pct = (end_frame / n_frames_total) * 100
|
||||
line3 = f"Sequence: {n_frames_total} frames total | Progress: {progress_pct:.1f}%"
|
||||
cv2.putText(
|
||||
footer,
|
||||
line3,
|
||||
(15, 72),
|
||||
FONT,
|
||||
0.6,
|
||||
(200, 200, 200),
|
||||
1,
|
||||
cv2.LINE_AA,
|
||||
)
|
||||
|
||||
# Combine grid and footer
|
||||
combined = np.vstack([grid_bgr, footer])
|
||||
|
||||
return combined
|
||||
|
||||
|
||||
def load_pkl_sequence(pkl_path: Path) -> NDArray[np.float32]:
|
||||
"""Load a .pkl file containing silhouette sequence."""
|
||||
with open(pkl_path, "rb") as f:
|
||||
data = pickle.load(f)
|
||||
|
||||
# Handle different possible structures
|
||||
if isinstance(data, np.ndarray):
|
||||
return data.astype(np.float32)
|
||||
elif isinstance(data, list):
|
||||
# List of frames
|
||||
return np.stack([np.array(frame) for frame in data]).astype(np.float32)
|
||||
else:
|
||||
raise ValueError(f"Unexpected data type in {pkl_path}: {type(data)}")
|
||||
|
||||
|
||||
def create_windows(
|
||||
sequence: NDArray[np.float32],
|
||||
window_size: int = 30,
|
||||
stride: int = 30,
|
||||
) -> list[NDArray[np.float32]]:
|
||||
"""
|
||||
Split a sequence into sliding windows.
|
||||
|
||||
Args:
|
||||
sequence: Array of shape (N, 64, 44)
|
||||
window_size: Number of frames per window
|
||||
stride: Stride between consecutive windows
|
||||
|
||||
Returns:
|
||||
List of window arrays, each of shape (window_size, 64, 44)
|
||||
"""
|
||||
n_frames = sequence.shape[0]
|
||||
windows = []
|
||||
|
||||
for start_idx in range(0, n_frames - window_size + 1, stride):
|
||||
end_idx = start_idx + window_size
|
||||
window = sequence[start_idx:end_idx]
|
||||
windows.append(window)
|
||||
|
||||
return windows
|
||||
|
||||
|
||||
def export_positive_batches(
|
||||
dataset_root: Path,
|
||||
output_dir: Path,
|
||||
window_size: int = 30,
|
||||
stride: int = 30,
|
||||
max_sequences: int | None = None,
|
||||
) -> None:
|
||||
"""
|
||||
Export all positive labeled batches from Scoliosis1K dataset as time windows.
|
||||
|
||||
Args:
|
||||
dataset_root: Path to Scoliosis1K-sil-pkl directory
|
||||
output_dir: Output directory for visualizations
|
||||
window_size: Number of frames per window (default 30)
|
||||
stride: Stride between consecutive windows (default 30 = non-overlapping)
|
||||
max_sequences: Maximum number of sequences to process (None = all)
|
||||
"""
|
||||
output_dir.mkdir(parents=True, exist_ok=True)
|
||||
|
||||
# Find all positive samples
|
||||
positive_samples: list[
|
||||
tuple[Path, str, str, str]
|
||||
] = [] # (pkl_path, subject_id, view_name, pkl_name)
|
||||
|
||||
for subject_dir in sorted(dataset_root.iterdir()):
|
||||
if not subject_dir.is_dir():
|
||||
continue
|
||||
subject_id = subject_dir.name
|
||||
|
||||
# Check for positive class directory (lowercase)
|
||||
positive_dir = subject_dir / "positive"
|
||||
if not positive_dir.exists():
|
||||
continue
|
||||
|
||||
# Iterate through views
|
||||
for view_dir in sorted(positive_dir.iterdir()):
|
||||
if not view_dir.is_dir():
|
||||
continue
|
||||
view_name = view_dir.name
|
||||
|
||||
# Find .pkl files
|
||||
for pkl_file in sorted(view_dir.glob("*.pkl")):
|
||||
positive_samples.append(
|
||||
(pkl_file, subject_id, view_name, pkl_file.stem)
|
||||
)
|
||||
|
||||
print(f"Found {len(positive_samples)} positive labeled sequences")
|
||||
|
||||
if max_sequences:
|
||||
positive_samples = positive_samples[:max_sequences]
|
||||
print(f"Processing first {max_sequences} sequences")
|
||||
|
||||
total_windows = 0
|
||||
|
||||
# Export each sequence's windows
|
||||
for seq_idx, (pkl_path, subject_id, view_name, pkl_name) in enumerate(
|
||||
positive_samples, 1
|
||||
):
|
||||
print(
|
||||
f"[{seq_idx}/{len(positive_samples)}] Processing {subject_id}/{view_name}/{pkl_name}..."
|
||||
)
|
||||
|
||||
# Load sequence
|
||||
try:
|
||||
sequence = load_pkl_sequence(pkl_path)
|
||||
except Exception as e:
|
||||
print(f" Error loading {pkl_path}: {e}")
|
||||
continue
|
||||
|
||||
# Ensure correct shape (N, 64, 44)
|
||||
if len(sequence.shape) == 2:
|
||||
# Single frame
|
||||
sequence = sequence[np.newaxis, ...]
|
||||
elif len(sequence.shape) == 3:
|
||||
# (N, H, W) - expected
|
||||
pass
|
||||
else:
|
||||
print(f" Unexpected shape {sequence.shape}, skipping")
|
||||
continue
|
||||
|
||||
n_frames = sequence.shape[0]
|
||||
print(f" Sequence has {n_frames} frames")
|
||||
|
||||
# Skip if sequence is shorter than window size
|
||||
if n_frames < window_size:
|
||||
print(f" Skipping: sequence too short (< {window_size} frames)")
|
||||
continue
|
||||
|
||||
# Normalize if needed
|
||||
if sequence.max() > 1.0:
|
||||
sequence = sequence / 255.0
|
||||
|
||||
# Create windows
|
||||
windows = create_windows(sequence, window_size=window_size, stride=stride)
|
||||
print(f" Created {len(windows)} windows (size={window_size}, stride={stride})")
|
||||
|
||||
# Export each window
|
||||
for window_idx, window in enumerate(windows):
|
||||
start_frame = window_idx * stride
|
||||
end_frame = start_frame + window_size
|
||||
|
||||
# Create visualization for this window with full metadata
|
||||
vis_image = create_optimized_visualization(
|
||||
silhouettes=window,
|
||||
subject_id=subject_id,
|
||||
view_name=view_name,
|
||||
window_idx=window_idx,
|
||||
start_frame=start_frame,
|
||||
end_frame=end_frame,
|
||||
n_frames_total=n_frames,
|
||||
)
|
||||
|
||||
# Save with descriptive filename including window index
|
||||
output_filename = (
|
||||
f"{subject_id}_{view_name}_{pkl_name}_win{window_idx:03d}.png"
|
||||
)
|
||||
output_path = output_dir / output_filename
|
||||
cv2.imwrite(str(output_path), vis_image)
|
||||
|
||||
# Save metadata for this window
|
||||
meta = {
|
||||
"subject_id": subject_id,
|
||||
"view": view_name,
|
||||
"pkl_name": pkl_name,
|
||||
"window_index": window_idx,
|
||||
"window_size": window_size,
|
||||
"stride": stride,
|
||||
"start_frame": start_frame,
|
||||
"end_frame": end_frame,
|
||||
"sequence_shape": sequence.shape,
|
||||
"n_frames_total": n_frames,
|
||||
"source_path": str(pkl_path),
|
||||
}
|
||||
meta_filename = (
|
||||
f"{subject_id}_{view_name}_{pkl_name}_win{window_idx:03d}.json"
|
||||
)
|
||||
meta_path = output_dir / meta_filename
|
||||
with open(meta_path, "w") as f:
|
||||
json.dump(meta, f, indent=2)
|
||||
|
||||
total_windows += 1
|
||||
|
||||
print(f" Exported {len(windows)} windows")
|
||||
|
||||
print(f"\nExport complete! Saved {total_windows} windows to {output_dir}")
|
||||
|
||||
|
||||
def main() -> None:
|
||||
"""Main entry point."""
|
||||
# Paths
|
||||
dataset_root = Path("/mnt/public/data/Scoliosis1K/Scoliosis1K-sil-pkl")
|
||||
output_dir = Path("/home/crosstyan/Code/OpenGait/output/positive_batches")
|
||||
|
||||
if not dataset_root.exists():
|
||||
print(f"Error: Dataset not found at {dataset_root}")
|
||||
return
|
||||
|
||||
# Export all positive batches with windowing
|
||||
export_positive_batches(
|
||||
dataset_root,
|
||||
output_dir,
|
||||
window_size=30, # 30 frames per window
|
||||
stride=30, # Non-overlapping windows
|
||||
)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
main()
|
||||
Reference in New Issue
Block a user