From 10aaa38eceb8bde84f662626cfcddba75fbe3887 Mon Sep 17 00:00:00 2001 From: Dmitry Batenkov Date: Thu, 12 Feb 2026 16:09:49 -0500 Subject: [PATCH 01/21] refresh button in the data dashboard --- collab_env/dashboard/app.py | 45 ++++++++++++++++++++++++++++++++++++- 1 file changed, 44 insertions(+), 1 deletion(-) diff --git a/collab_env/dashboard/app.py b/collab_env/dashboard/app.py index 0161c893..5c75f716 100644 --- a/collab_env/dashboard/app.py +++ b/collab_env/dashboard/app.py @@ -224,6 +224,11 @@ def __init__( self.clear_cache_button = pn.widgets.Button(name="Clear Cache", width=100) self.cache_info_pane = pn.pane.HTML("", width=300) + # Data browser refresh button + self.refresh_browser_button = pn.widgets.Button( + name="🔄 Refresh", button_type="primary", width=100 + ) + # PLY viewer state self.current_ply_viewer = None @@ -235,6 +240,7 @@ def __init__( self.save_button.on_click(self._save_edit) self.cancel_edit_button.on_click(self._cancel_edit) self.clear_cache_button.on_click(self._clear_cache) + self.refresh_browser_button.on_click(self._refresh_data_browser) self.convert_video_button.on_click(self._convert_video) self.bbox_viewer_button.on_click(self._open_bbox_viewer) self.mesh_3d_viewer_button.on_click(self._open_mesh_3d_viewer) @@ -271,6 +277,43 @@ def _load_sessions(self): f"

Error loading sessions: {e}

" ) + def _refresh_data_browser(self, event=None): + """Refresh data browser by re-reading buckets and resetting app state.""" + try: + logger.info("Refreshing data browser...") + self.status_pane.object = "

🔄 Refreshing data browser...

" + + # Reset app state + self.selected_session = "" + self.selected_file = "" + self.current_bucket_type = "curated" + self.session_select.value = "" + + # Clear current session data + self.current_session_files = [] + self.current_file_content = None + self.current_file_info = {} + self.display_to_path_map = {} + + # Hide UI elements + self.bucket_type_toggle.visible = False + self.file_tree.visible = False + self.file_tree.options = [] + self.file_viewer.object = "

Select a file to view its contents

" + self.file_name_header.object = "

No file selected

" + self.file_management_controls.visible = False + + # Re-load sessions from both buckets + self._load_sessions() + + logger.info("Data browser refreshed successfully") + + except Exception as e: + logger.error(f"Error refreshing data browser: {e}") + self.status_pane.object = ( + f"

Error refreshing: {e}

" + ) + def _on_session_change(self, event): """Handle session selection change.""" session_name = event.new @@ -2110,7 +2153,7 @@ def create_layout(self): cache_controls = pn.Column(self.cache_info_pane, self.clear_cache_button) nav_panel = pn.Column( - "## Data Browser", + pn.Row("## Data Browser", pn.Spacer(width=100), self.refresh_browser_button), self.session_select, self.bucket_type_toggle, self.file_tree, From ece939392bd1725ba753821c678f8eb0926321ba Mon Sep 17 00:00:00 2001 From: Dmitry Batenkov Date: Fri, 13 Feb 2026 17:50:30 -0500 Subject: [PATCH 02/21] tracking studio - initial version. WIP --- Dockerfile.tracking-studio | 59 ++ cloudbuild.yaml | 41 + collab_env/tracking_studio/__init__.py | 12 + collab_env/tracking_studio/app.py | 802 ++++++++++++++++++ .../tracking_studio/bytetrack_params.json | 37 + collab_env/tracking_studio/gcs_browser.py | 168 ++++ collab_env/tracking_studio/model_manager.py | 365 ++++++++ collab_env/tracking_studio/video_converter.py | 96 +++ collab_env/tracking_studio/video_processor.py | 374 ++++++++ docs/tracking/tracking_web_gui.md | 345 ++++++++ pyproject.toml | 6 +- scripts/tracking/bytetrack_video_inference.py | 170 ++++ .../configs/botsort_thermal_rats.yaml | 21 + .../configs/bytetrack_thermal_rats.yaml | 16 + scripts/tracking/roboflow_video_inference.py | 23 + scripts/tracking/run_tracking_studio.py | 9 + scripts/tracking/yolo_native_tracking.py | 109 +++ 17 files changed, 2652 insertions(+), 1 deletion(-) create mode 100644 Dockerfile.tracking-studio create mode 100644 cloudbuild.yaml create mode 100644 collab_env/tracking_studio/__init__.py create mode 100644 collab_env/tracking_studio/app.py create mode 100644 collab_env/tracking_studio/bytetrack_params.json create mode 100644 collab_env/tracking_studio/gcs_browser.py create mode 100644 collab_env/tracking_studio/model_manager.py create mode 100644 collab_env/tracking_studio/video_converter.py create mode 100644 collab_env/tracking_studio/video_processor.py create mode 100644 docs/tracking/tracking_web_gui.md create mode 100644 scripts/tracking/bytetrack_video_inference.py create mode 100644 scripts/tracking/configs/botsort_thermal_rats.yaml create mode 100644 scripts/tracking/configs/bytetrack_thermal_rats.yaml create mode 100644 scripts/tracking/roboflow_video_inference.py create mode 100755 scripts/tracking/run_tracking_studio.py create mode 100644 scripts/tracking/yolo_native_tracking.py diff --git a/Dockerfile.tracking-studio b/Dockerfile.tracking-studio new file mode 100644 index 00000000..f3924c13 --- /dev/null +++ b/Dockerfile.tracking-studio @@ -0,0 +1,59 @@ +FROM python:3.10-slim + +# Install system dependencies +RUN apt-get update && apt-get install -y \ + ffmpeg \ + libgl1 \ + libglib2.0-0 \ + libsm6 \ + libxext6 \ + libxrender-dev \ + && rm -rf /var/lib/apt/lists/* + +WORKDIR /workspace + +# Copy and install Python dependencies +COPY pyproject.toml . +RUN pip install --no-cache-dir \ + nicegui \ + ultralytics \ + supervision \ + google-cloud-storage \ + gcsfs \ + inference \ + opencv-python-headless \ + pandas \ + numpy \ + loguru + +# Copy application code +COPY collab_env/ collab_env/ +COPY scripts/ scripts/ + +# Copy GCS credentials +COPY config-local/ config/ + +# Pre-download YOLO models (YOLO11 and YOLO26 variants) +RUN mkdir -p /workspace/models && \ + python -c "from ultralytics import YOLO; \ + print('Downloading YOLO11 models...'); \ + YOLO('yolo11n.pt'); \ + YOLO('yolo11s.pt'); \ + YOLO('yolo11m.pt');" && \ + mv ~/.cache/ultralytics/*.pt /workspace/models/ 2>/dev/null || true + +# Note: YOLO26 models might not be available yet, will download on first use +# RUN python -c "from ultralytics import YOLO; \ +# print('Downloading YOLO26 models...'); \ +# YOLO('yolo26n-fast.pt'); \ +# YOLO('yolo26s-fast.pt'); \ +# YOLO('yolo26m-fast.pt');" && \ +# mv ~/.cache/ultralytics/*.pt /workspace/models/ 2>/dev/null || true + +# Create tmp directories +RUN mkdir -p /tmp/videos /tmp/outputs /tmp/uploads + +ENV PORT=8080 +ENV PYTHONUNBUFFERED=1 + +CMD ["python", "scripts/run_tracking_studio.py"] diff --git a/cloudbuild.yaml b/cloudbuild.yaml new file mode 100644 index 00000000..f3865f58 --- /dev/null +++ b/cloudbuild.yaml @@ -0,0 +1,41 @@ +steps: + - name: 'gcr.io/cloud-builders/docker' + args: + - 'build' + - '-t' + - 'gcr.io/$PROJECT_ID/tracking-studio:$SHORT_SHA' + - '-t' + - 'gcr.io/$PROJECT_ID/tracking-studio:latest' + - '-f' + - 'Dockerfile.tracking-studio' + - '.' + + - name: 'gcr.io/cloud-builders/docker' + args: ['push', 'gcr.io/$PROJECT_ID/tracking-studio:$SHORT_SHA'] + + - name: 'gcr.io/cloud-builders/docker' + args: ['push', 'gcr.io/$PROJECT_ID/tracking-studio:latest'] + + - name: 'gcr.io/google.com/cloudsdktool/cloud-sdk' + entrypoint: gcloud + args: + - 'run' + - 'deploy' + - 'tracking-studio' + - '--image=gcr.io/$PROJECT_ID/tracking-studio:$SHORT_SHA' + - '--platform=managed' + - '--region=us-central1' + - '--memory=4Gi' + - '--cpu=2' + - '--timeout=900' + - '--concurrency=1' + - '--max-instances=5' + - '--set-env-vars=ROBOFLOW_API_KEY=${_ROBOFLOW_API_KEY}' + - '--allow-unauthenticated' + +images: + - 'gcr.io/$PROJECT_ID/tracking-studio:$SHORT_SHA' + - 'gcr.io/$PROJECT_ID/tracking-studio:latest' + +options: + machineType: 'N1_HIGHCPU_8' diff --git a/collab_env/tracking_studio/__init__.py b/collab_env/tracking_studio/__init__.py new file mode 100644 index 00000000..38508125 --- /dev/null +++ b/collab_env/tracking_studio/__init__.py @@ -0,0 +1,12 @@ +""" +NiceGUI-based Video Tracking Studio + +A web application for interactive video tracking with support for: +- GCS bucket video browsing and upload +- YOLO and Roboflow model selection +- Real-time tracking visualization +- ByteTrack with Re-ID support +- CSV output export +""" + +__version__ = "0.1.0" diff --git a/collab_env/tracking_studio/app.py b/collab_env/tracking_studio/app.py new file mode 100644 index 00000000..c1f59956 --- /dev/null +++ b/collab_env/tracking_studio/app.py @@ -0,0 +1,802 @@ +""" +Main NiceGUI Tracking Studio Application + +Single-page interactive app for video tracking with: +- GCS bucket browsing and video upload +- Model selection (YOLO and Roboflow) +- Real-time tracking visualization +- CSV output download +""" + +from nicegui import ui, app +import asyncio +from pathlib import Path +import uuid +import os +import io +from loguru import logger + +from .gcs_browser import GCSVideoBrowser +from .model_manager import ModelManager +from .video_processor import VideoTracker +from .video_converter import convert_to_h264, needs_conversion + + +# Load ByteTrack parameter definitions +def load_bytetrack_params(): + """Load ByteTrack parameter schema from JSON""" + import json + params_file = Path(__file__).parent / "bytetrack_params.json" + with open(params_file, 'r') as f: + return json.load(f) + +bytetrack_params_schema = load_bytetrack_params() + + +# Initialize services +def get_credentials_path(): + """Get GCS credentials path from environment or default""" + return os.getenv( + "GCS_CREDENTIALS", + "/workspace/config/collab-data-463313-c340ad86b28e.json", + ) + + +try: + gcs_browser = GCSVideoBrowser(credentials_path=get_credentials_path()) +except Exception as e: + logger.error(f"Failed to initialize GCS browser: {e}") + gcs_browser = None + +model_manager = ModelManager() + + +@ui.page("/") +async def index(): + """Main tracking studio page""" + session_id = str(uuid.uuid4())[:8] + + # State variables (stored in page context) + import threading + state = { + "selected_bucket": None, + "selected_video_path": None, + "selected_model": None, + "processing": False, + "results": None, + "uploaded_video": None, + "uploaded_model": None, # Uploaded model .pt file + "stop_event": None, # Hard stop + "pause_event": None, # Pause/resume + "skip_frames_event": None, # Skip forward signal + "video_path": None, # Path to current video being processed + "video_loaded": False, # Video is loaded and ready for playback + "model_loaded": False, # Model is loaded and ready for tracking + "loaded_model": None, # Reference to loaded model object + } + + # UI Layout + with ui.column().classes("w-full p-4 gap-3"): + # Header + with ui.row().classes("w-full items-center mb-2"): + ui.label("🎯 Video Tracking Studio").classes("text-2xl font-bold") + ui.space() + ui.label("Real-time object detection and tracking").classes("text-sm text-gray-600") + + # Row 1: Video source selection (Bucket | Folder | Video | Upload) + with ui.card().classes("w-full shadow-md p-2"): + with ui.row().classes("w-full gap-2 items-end"): + # GCS selection + if gcs_browser: + bucket_select = ui.select( + label="Bucket", options=[], + ).style("width: 250px") + try: + buckets = gcs_browser.list_buckets() + bucket_select.options = buckets + if buckets: + bucket_select.value = buckets[0] + except Exception as e: + logger.error(f"Failed to list buckets: {e}") + + folder_select = ui.select( + label="Folder", options=[""], value="", clearable=True + ).style("width: 350px") + + video_select = ui.select(label="Video", options=[]).classes("flex-grow") + + async def update_folders(e): + """Update folder list when bucket changes""" + try: + bucket = bucket_select.value + if bucket: + folders = gcs_browser.list_folders(bucket, "") + folder_select.options = [""] + folders + folder_select.value = "" + folder_select.update() + await update_video_list(None) + except Exception as error: + logger.error(f"Failed to list folders: {error}") + + async def update_video_list(e): + """Update video list when bucket or folder changes""" + try: + bucket = bucket_select.value + folder = folder_select.value or "" + if bucket: + videos = gcs_browser.list_videos(bucket, folder) + video_select.options = [v["rel_path"] for v in videos] + video_select.update() + except Exception as error: + logger.error(f"Failed to list videos: {error}") + + def enable_load_video_btn(e=None): + """Enable Load Video button when video is selected""" + if video_select.value or state.get("uploaded_video"): + load_video_btn.enable() + + bucket_select.on("update:model-value", update_folders) + folder_select.on("update:model-value", update_video_list) + video_select.on("update:model-value", enable_load_video_btn) + + if bucket_select.value: + ui.timer(0.1, lambda: update_folders(None), once=True) + + # Upload widget + async def handle_upload(e): + """Handle user video upload""" + try: + upload_path = Path(f"/tmp/uploads/{session_id}") + upload_path.mkdir(parents=True, exist_ok=True) + uploaded_file = upload_path / e.name + uploaded_file.write_bytes(e.content.read()) + state["uploaded_video"] = uploaded_file + ui.notify(f"Uploaded: {e.name}") + load_video_btn.enable() # Enable Load Video button + except Exception as error: + logger.error(f"Upload failed: {error}") + ui.notify(f"Upload failed: {error}", type="negative") + + upload = ui.upload( + on_upload=handle_upload, + auto_upload=True, + ).props("accept=video/mp4,video/quicktime,video/x-msvideo dense flat").props("label=Upload").style("width: 120px; height: 40px") + + # Row 2: Model/Params (left) | Controls + Preview (right) + with ui.row().classes("w-full gap-3"): + # LEFT: Model + Parameters (stacked) + with ui.column().classes("gap-3").style("flex: 0 0 280px; min-width: 280px"): + # Model card + with ui.card().classes("w-full shadow-md p-3"): + ui.label("🤖 Model").classes("text-sm font-semibold mb-2") + + # Radio with three options + with ui.row().classes("gap-2"): + model_source = ui.radio(["YOLO", "Roboflow", "Custom"], value="YOLO").classes("text-xs") + + # YOLO model selection (default visible) + yolo_container = ui.column().classes("w-full mt-2") + with yolo_container: + yolo_model_input = ui.input( + label="Model Name", + placeholder="e.g., yolo11n.pt", + value="yolo11n.pt" + ).classes("w-full").tooltip("Enter any YOLO model name (will auto-download if available)") + + # Roboflow model selection (hidden by default) + rf_container = ui.column().classes("w-full mt-2 gap-2") + rf_container.visible = False + with rf_container: + rf_project_input = ui.input( + label="Project ID", + placeholder="workspace/project", + value="dima-sdrkv/ratsmerged20260211" + ).classes("w-full") + + # Define the function BEFORE referencing it in the button + async def list_rf_models(): + """Query Roboflow for available model versions""" + project_id = rf_project_input.value + if not project_id: + ui.notify("Please enter project ID", type="warning") + return + + try: + rf_list_btn.disable() + # Call directly (synchronous HTTP request, no need for threading) + versions = model_manager.list_roboflow_project_models(project_id) + if versions: + rf_version_select.options = versions + rf_version_select.value = versions[0] + rf_version_select.enable() + ui.notify(f"Found {len(versions)} versions", type="positive") + else: + ui.notify("No versions found", type="warning") + except Exception as error: + logger.error(f"Failed to list models: {error}") + ui.notify(f"Error: {error}", type="negative") + finally: + rf_list_btn.enable() + + # Now create the button and version select (after function definition) + with ui.row().classes("w-full gap-2"): + rf_list_btn = ui.button("List Models", on_click=list_rf_models).props("size=sm color=primary") + + rf_version_select = ui.select( + label="Version", + options=[], + ).classes("w-full") + rf_version_select.disable() + + # Custom model upload (hidden by default) + custom_container = ui.column().classes("w-full mt-2") + custom_container.visible = False + with custom_container: + # Upload widget for model weights + async def handle_model_upload(e): + """Handle model .pt file upload""" + try: + model_upload_path = Path(f"/tmp/models/{session_id}") + model_upload_path.mkdir(parents=True, exist_ok=True) + uploaded_model_file = model_upload_path / e.name + uploaded_model_file.write_bytes(e.content.read()) + state["uploaded_model"] = uploaded_model_file + ui.notify(f"Model uploaded: {e.name}", type="positive") + logger.info(f"Model uploaded to: {uploaded_model_file}") + load_model_btn.enable() # Enable Load Model button + except Exception as error: + logger.error(f"Model upload failed: {error}") + ui.notify(f"Model upload failed: {error}", type="negative") + + ui.upload( + on_upload=handle_model_upload, + auto_upload=True, + ).props("accept=.pt dense flat").props("label=Upload Model (.pt)").classes("w-full") + + # Toggle visibility based on model source + def toggle_model_ui(e=None): + value = model_source.value + if value == "YOLO": + yolo_container.visible = True + rf_container.visible = False + custom_container.visible = False + # Enable load button if YOLO model name is entered + if yolo_model_input.value: + load_model_btn.enable() + elif value == "Roboflow": + yolo_container.visible = False + rf_container.visible = True + custom_container.visible = False + # Enable load button if Roboflow model is selected + if rf_version_select.value: + load_model_btn.enable() + else: # Custom + yolo_container.visible = False + rf_container.visible = False + custom_container.visible = True + # Enable load button if custom model uploaded + if state.get("uploaded_model"): + load_model_btn.enable() + + def enable_load_model_btn(e=None): + """Enable Load Model button when model is selected""" + if model_source.value == "YOLO" and yolo_model_input.value: + load_model_btn.enable() + elif model_source.value == "Roboflow" and rf_version_select.value: + load_model_btn.enable() + elif model_source.value == "Custom" and state.get("uploaded_model"): + load_model_btn.enable() + + model_source.on("update:model-value", toggle_model_ui) + yolo_model_input.on("update:model-value", enable_load_model_btn) + rf_version_select.on("update:model-value", enable_load_model_btn) + + # Parameters card + with ui.card().classes("w-full shadow-md p-3"): + ui.label("⚙️ Parameters").classes("text-sm font-semibold mb-2") + + # Detection confidence (not in ByteTrack params) + with ui.column().classes("w-full gap-1"): + conf_label = ui.label("Confidence: 0.50").classes("text-xs") + conf_slider = ui.slider(min=0.1, max=0.9, step=0.05, value=0.5).classes("w-full").tooltip("Detection confidence threshold") + conf_slider.on("update:model-value", lambda e: conf_label.set_text(f"Confidence: {e.args:.2f}")) + + # Dynamic ByteTrack parameters from JSON + param_widgets = {} # Store references to UI elements + with ui.column().classes("w-full gap-1 mt-2"): + for param_name, param_config in bytetrack_params_schema.items(): + if param_config["type"] == "float": + # Float slider + default_val = param_config["default"] + min_val, max_val = param_config["range"] + + # Create label with tooltip + param_label = ui.label(f"{param_name.replace('_', ' ').title()}: {default_val:.2f}").classes("text-xs") + param_label.tooltip(param_config["description"]) + + # Create slider + step = 0.05 if max_val <= 1.0 else 0.1 + param_slider = ui.slider( + min=min_val, + max=max_val, + step=step, + value=default_val + ).classes("w-full") + + # Update label on change + param_slider.on( + "update:model-value", + lambda e, lbl=param_label, name=param_name: lbl.set_text( + f"{name.replace('_', ' ').title()}: {e.args:.2f}" + ) + ) + param_widgets[param_name] = param_slider + + elif param_config["type"] == "int": + # Int slider + default_val = param_config["default"] + min_val = param_config["range"][0] + max_val = param_config["range"][1] if param_config["range"][1] else 300 + + param_label = ui.label(f"{param_name.replace('_', ' ').title()}: {default_val}").classes("text-xs") + param_label.tooltip(param_config["description"]) + + param_slider = ui.slider( + min=min_val, + max=max_val, + step=1, + value=default_val + ).classes("w-full") + + param_slider.on( + "update:model-value", + lambda e, lbl=param_label, name=param_name: lbl.set_text( + f"{name.replace('_', ' ').title()}: {int(e.args)}" + ) + ) + param_widgets[param_name] = param_slider + + elif param_config["type"] == "bool": + # Checkbox + param_checkbox = ui.checkbox( + param_name.replace('_', ' ').title(), + value=param_config["default"] + ).classes("text-xs") + param_checkbox.tooltip(param_config["description"]) + param_widgets[param_name] = param_checkbox + + elif param_config["type"] == "string": + # Dropdown for options + param_select = ui.select( + label=param_name.replace('_', ' ').title(), + options=param_config["options"], + value=param_config["default"] + ).classes("w-full text-xs") + param_select.tooltip(param_config["description"]) + param_widgets[param_name] = param_select + + # Skip frames (for fast-forward, not a ByteTrack param) + with ui.row().classes("w-full items-center gap-2 mt-2"): + skip_frames_label = ui.label("Skip Frames: 1 frame").classes("text-xs") + skip_frames_slider = ui.slider(min=1, max=30, step=1, value=1).style("width: 100px") + skip_frames_slider.tooltip("Process every Nth frame (1 = all frames)") + skip_frames_slider.on("update:model-value", lambda e: skip_frames_label.set_text(f"Skip Frames: {int(e.args)} {'frame' if int(e.args) == 1 else 'frames'}")) + + # GUI refresh rate (display updates, not a ByteTrack param) + with ui.row().classes("w-full items-center gap-2 mt-1"): + display_update_label = ui.label("Display Update: 10 frames").classes("text-xs") + display_update_slider = ui.slider(min=1, max=30, step=1, value=10).style("width: 100px") + display_update_slider.tooltip("Update display every Nth frame (lower = smoother, more network traffic)") + display_update_slider.on("update:model-value", lambda e: display_update_label.set_text(f"Display Update: {int(e.args)} {'frame' if int(e.args) == 1 else 'frames'}")) + + # RIGHT: Controls + Preview (stacked vertically) + with ui.column().classes("flex-grow gap-3"): + # Controls card + with ui.card().classes("w-full shadow-md p-3"): + # Row 1: Load buttons + with ui.row().classes("w-full items-center gap-2 mb-2"): + load_video_btn = ui.button("Load Video").props("color=primary icon=video_file") + load_video_btn.disable() # Enabled when video selected + + load_model_btn = ui.button("Load Model").props("color=primary icon=model_training") + load_model_btn.disable() # Enabled when model selected + + ui.separator().props("vertical") + + with ui.column().classes("flex-grow gap-1"): + status_label = ui.label("Select video and model").classes("text-xs") + + # Row 2: Playback controls + with ui.row().classes("w-full items-center gap-2"): + start_btn = ui.button("Start Tracking").props("color=positive icon=play_arrow") + start_btn.disable() # Enabled when both video and model loaded + + pause_btn = ui.button("Pause").props("color=warning icon=pause") + pause_btn.disable() + + stop_btn = ui.button("Stop").props("color=negative icon=stop") + stop_btn.disable() + + ui.separator().props("vertical") + + with ui.column().classes("flex-grow gap-1"): + progress_label = ui.label("Ready").classes("text-xs") + progress = ui.linear_progress(value=0).props("size=15px color=primary") + + # Time slider for seeking + with ui.column().classes("w-full gap-1 mt-2"): + time_label = ui.label("Frame: 0 / 0").classes("text-xs text-gray-600") + time_slider = ui.slider(min=0, max=100, value=0).props("lazy").classes("w-full") + time_slider.disable() + + async def preview_frame_on_drag(e): + """Show video frame preview during drag (no detection/tracking)""" + if not state.get("video_path"): + return + + target_frame = int(e.args) + try: + # Open video for preview (separate from processing thread) + import cv2 + import base64 + cap = cv2.VideoCapture(str(state["video_path"])) + cap.set(cv2.CAP_PROP_POS_FRAMES, target_frame) + ret, frame = cap.read() + cap.release() + + if ret: + # Show raw frame without annotations + _, buffer = cv2.imencode(".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, 85]) + img_base64 = base64.b64encode(buffer).decode('utf-8') + video_display.set_source(f"data:image/jpeg;base64,{img_base64}") + total_frames = state.get("total_frames", target_frame) + time_label.text = f"Frame: {target_frame} / {total_frames}" + except Exception as err: + logger.warning(f"Preview failed: {err}") + + def seek_to_frame(e): + """Seek to specific frame when slider is released""" + if state.get("processing") and state.get("skip_frames_event"): + target_frame = int(e.args) + current_frame = state.get("current_frame", 0) + if target_frame != current_frame: + # Clear any pending seeks + state["skip_frames_event"]["skip_amount"] = target_frame - current_frame + ui.notify(f"Seeking to frame {target_frame}...", type="info") + + # Live preview during drag + time_slider.on("update:model-value", preview_frame_on_drag) + # Actual seek on release + time_slider.on("change", seek_to_frame) + + # Preview card + with ui.card().classes("w-full shadow-md p-3"): + ui.label("Live Preview").classes("text-sm font-semibold mb-2") + video_display = ui.interactive_image().classes("w-full border-2 border-gray-200 rounded bg-gray-50").style("max-height: 500px; object-fit: contain;") + + # Results (initially hidden, separate row) + results_container = ui.card().classes("w-full shadow-md p-3 hidden") + with results_container: + with ui.row().classes("w-full items-center gap-3"): + ui.label("✅ Results").classes("text-sm font-semibold") + stats_label = ui.label().classes("text-sm flex-grow") + download_track_btn = ui.button("Download CSV").props("color=primary icon=download size=sm") + + # Event handlers + async def load_video(): + """Load and prepare video for viewing/tracking""" + from nicegui import context + + try: + status_label.text = "Loading video..." + load_video_btn.disable() + + # Capture client context before threading + client = context.client + + # Get video (either download from GCS or use uploaded) + if state.get("uploaded_video"): + # Use uploaded video + local_video = state["uploaded_video"] + status_label.text = "Using uploaded video..." + elif gcs_browser and bucket_select.value and video_select.value: + # Download from GCS + status_label.text = "Downloading video..." + bucket = bucket_select.value + folder = folder_select.value or "" + video_name = video_select.value + gcs_path = f"{bucket}/{video_name}" + + local_video_dir = Path(f"/tmp/videos/{session_id}") + local_video_dir.mkdir(parents=True, exist_ok=True) + local_video = local_video_dir / Path(video_name).name + + await asyncio.to_thread(gcs_browser.download_video, gcs_path, str(local_video)) + else: + raise ValueError("No video selected. Please select or upload a video.") + + # Convert to H.264 if needed + if await asyncio.to_thread(needs_conversion, local_video): + status_label.text = "Converting to H.264..." + converted_video = local_video.parent / f"{local_video.stem}_h264.mp4" + await asyncio.to_thread(convert_to_h264, local_video, converted_video) + local_video = converted_video + with client: + ui.notify("Video converted to H.264") + + # Restore context for UI updates after threading + with client: + # Store video path in state + state["video_path"] = local_video + state["video_loaded"] = True + + # Set up video display and time slider + import cv2 + import base64 + cap = cv2.VideoCapture(str(local_video)) + total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + state["total_frames"] = total_frames + + # Display first frame + cap.set(cv2.CAP_PROP_POS_FRAMES, 0) + ret, frame = cap.read() + if ret: + _, buffer = cv2.imencode(".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, 85]) + img_base64 = base64.b64encode(buffer).decode('utf-8') + video_display.set_source(f"data:image/jpeg;base64,{img_base64}") + + cap.release() + + # Enable time slider for playback + time_slider.enable() + time_slider.set_value(0) + time_slider.props(f"max={total_frames}") + time_label.text = f"Frame: 0 / {total_frames}" + + status_label.text = "Video loaded ✓" + ui.notify("Video loaded successfully", type="positive") + + # Enable Start button if model is also loaded + if state["model_loaded"]: + start_btn.enable() + + except Exception as e: + logger.error(f"Failed to load video: {e}", exc_info=True) + with client: + status_label.text = f"Error loading video" + ui.notify(f"Error: {str(e)}", type="negative") + finally: + with client: + load_video_btn.enable() + + async def load_model(): + """Load selected model""" + from nicegui import context + + try: + status_label.text = "Loading model..." + load_model_btn.disable() + + # Capture client context before threading + client = context.client + + if model_source.value == "YOLO": + model = await asyncio.to_thread( + model_manager.load_yolo_model, yolo_model_input.value + ) + elif model_source.value == "Roboflow": + # Roboflow mode: download from API + project_id = rf_project_input.value + version = rf_version_select.value + + if not project_id or not version: + raise ValueError("Please select a Roboflow model (project ID and version)") + + # Validate project ID format (should be workspace/project) + project_parts = project_id.split('/') + if len(project_parts) != 2: + raise ValueError( + f"Invalid project ID: '{project_id}'\n" + f"Expected format: workspace/project (e.g., 'dima-sdrkv/ratsmerged20260211')" + ) + + # Construct full model ID: workspace/project/version + model_id = f"{project_id}/{version}" + logger.info(f"Loading Roboflow model: {model_id}") + + model = await asyncio.to_thread( + model_manager.load_roboflow_model, model_id + ) + else: # Custom + # Custom mode: use uploaded model file + if state.get("uploaded_model"): + model_path = str(state["uploaded_model"]) + logger.info(f"Loading uploaded model: {model_path}") + model = await asyncio.to_thread( + model_manager.load_roboflow_model, model_path + ) + else: + raise ValueError("Please upload a model .pt file") + + # Restore context for UI updates after threading + with client: + # Store model in state + state["loaded_model"] = model + state["model_loaded"] = True + + # Detect tracker type for display + from ultralytics import YOLO + tracker_type = "YOLO Native" if isinstance(model, YOLO) else "Supervision" + state["tracker_type"] = tracker_type + + status_label.text = f"Model loaded ✓ ({tracker_type} tracking)" + ui.notify("Model loaded successfully", type="positive") + + # Enable Start button if video is also loaded + if state["video_loaded"]: + start_btn.enable() + + except Exception as e: + logger.error(f"Failed to load model: {e}", exc_info=True) + # Restore context for error UI updates + with client: + status_label.text = f"Error loading model" + ui.notify(f"Error: {str(e)}", type="negative") + finally: + with client: + load_model_btn.enable() + + def pause_tracking(): + """Pause/resume tracking""" + if state["pause_event"]: + if state["pause_event"].is_set(): + # Currently paused, resume + state["pause_event"].clear() + pause_btn.props("icon=pause") + pause_btn.text = "Pause" + progress_label.text = "Resuming..." + ui.notify("Resumed", type="info") + else: + # Currently running, pause + state["pause_event"].set() + pause_btn.props("icon=play_arrow") + pause_btn.text = "Resume" + progress_label.text = "Paused" + ui.notify("Paused", type="warning") + + + def stop_tracking(): + """Hard stop - terminates processing""" + if state["stop_event"]: + state["stop_event"].set() + progress_label.text = "Stopping..." + ui.notify("Stopping tracking...", type="negative") + async def start_tracking(): + """Start tracking on already-loaded video with already-loaded model""" + if not state.get("video_loaded") or not state.get("model_loaded"): + ui.notify("Please load video and model first", type="warning") + return + + state["processing"] = True + state["stop_event"] = threading.Event() # Hard stop + state["pause_event"] = threading.Event() # Pause (starts clear = not paused) + state["skip_frames_event"] = {"skip_amount": 0} # Skip forward + state["current_frame"] = 0 + start_btn.disable() + pause_btn.enable() + stop_btn.enable() + results_container.classes(add="hidden") + + try: + # Show tracker type in progress label + tracker_type = state.get("tracker_type", "Unknown") + progress_label.text = f"Starting tracking ({tracker_type})..." + progress.value = 0 + + # Use already-loaded video and model from state + local_video = state["video_path"] + model = state["loaded_model"] + + # Frame callback for real-time UI updates + # Capture display update interval from slider + display_interval = int(display_update_slider.value) + + async def frame_callback(annotated_frame, frame_idx, total_frames): + """Update UI with current frame (throttled for performance)""" + state["current_frame"] = frame_idx + + # Update every N frames based on slider setting + if frame_idx % display_interval != 0 and frame_idx != total_frames - 1: + return + + # Convert frame to bytes for display + import cv2 + import base64 + _, buffer = cv2.imencode(".jpg", annotated_frame, [cv2.IMWRITE_JPEG_QUALITY, 85]) + img_base64 = base64.b64encode(buffer).decode('utf-8') + video_display.set_source(f"data:image/jpeg;base64,{img_base64}") + + # Update progress + progress.value = 0.1 + 0.8 * (frame_idx / total_frames) + progress_label.text = f"Tracking: Frame {frame_idx + 1}/{total_frames}" + + # Update time slider + time_slider.set_value(frame_idx) + time_label.text = f"Frame: {frame_idx} / {total_frames}" + + # Build tracker config from dynamic parameter widgets + tracker_config = { + "skip_frames": int(skip_frames_slider.value), # Fast-forward + } + # Add ByteTrack parameters from param_widgets + for param_name, widget in param_widgets.items(): + if hasattr(widget, 'value'): + tracker_config[param_name] = widget.value + + # Initialize tracker with dynamic parameters + tracker = VideoTracker( + model=model, + tracker_config=tracker_config, + confidence=conf_slider.value, + frame_callback=frame_callback, + stop_event=state["stop_event"], + pause_event=state["pause_event"], + skip_frames_event=state["skip_frames_event"], + ) + + output_dir = f"/tmp/outputs/{session_id}" + results = await tracker.process_video_realtime(str(local_video), output_dir) + + # Show results + progress.value = 1.0 + progress_label.text = "Complete!" + + state["results"] = results + + stats_label.text = ( + f"Processed {results['stats']['total_frames']} frames | " + f"{results['stats']['total_detections']} detections | " + f"{results['stats']['unique_tracks']} unique tracks" + ) + + # Setup download button (only tracking CSV) + download_track_btn.on_click(lambda: ui.download(results["tracking_csv"])) + + results_container.classes(remove="hidden") + + except Exception as e: + logger.error(f"Tracking failed: {e}", exc_info=True) + try: + progress_label.text = f"Error: {str(e)}" + ui.notify(f"Error: {str(e)}", type="negative") + except Exception as notify_error: + logger.error(f"Failed to show error notification: {notify_error}") + try: + progress_label.text = f"Error: {str(e)}" + except: + pass + + finally: + state["processing"] = False + state["stop_event"] = None + state["pause_event"] = None + state["skip_frames_event"] = None + start_btn.enable() + pause_btn.disable() + stop_btn.disable() + + # Wire up buttons to event handlers (after functions are defined) + load_video_btn.on_click(load_video) + load_model_btn.on_click(load_model) + start_btn.on_click(start_tracking) + pause_btn.on_click(lambda: pause_tracking()) + stop_btn.on_click(lambda: stop_tracking()) + + +# Run the NiceGUI app directly (no function wrapper for reload compatibility) +ui.run( + host="0.0.0.0", + port=int(os.getenv("PORT", 8080)), + reload=True, # Enable auto-reload for development + title="Tracking Studio", +) diff --git a/collab_env/tracking_studio/bytetrack_params.json b/collab_env/tracking_studio/bytetrack_params.json new file mode 100644 index 00000000..5700217a --- /dev/null +++ b/collab_env/tracking_studio/bytetrack_params.json @@ -0,0 +1,37 @@ +{ + "track_high_thresh": { + "type": "float", + "default": 0.25, + "range": [0.0, 1.0], + "description": "Detection confidence threshold that separates high-confidence and low-confidence detections. Detections above this score enter Stage 1 (primary IoU matching). Detections between track_low_thresh and this value enter Stage 2 (secondary matching to rescue lost tracks). Raise to restrict primary matching to only the most confident detections; lower to feed more detections into Stage 1." + }, + "track_low_thresh": { + "type": "float", + "default": 0.1, + "range": [0.0, 1.0], + "description": "Absolute minimum detection confidence. Detections below this score are discarded entirely and never participate in any matching stage. Lower to recover very marginal detections at the cost of more noise; raise to filter out weak false positives." + }, + "new_track_thresh": { + "type": "float", + "default": 0.25, + "range": [0.0, 1.0], + "description": "Minimum detection confidence required to initialize a new track. Unmatched detections from Stage 1 must exceed this score to spawn a new track ID. Higher values prevent spurious tracks from false positives; lower values allow tracks to start from weaker detections." + }, + "track_buffer": { + "type": "int", + "default": 30, + "range": [1, null], + "description": "Number of frames a lost track is kept alive before permanent deletion. Internally scaled by frame rate: max_time_lost = int(fps / 30.0 * track_buffer). Higher values let tracks survive longer occlusions but increase ID switch risk when the object reappears far from its last position. Lower values remove lost tracks faster." + }, + "match_thresh": { + "type": "float", + "default": 0.8, + "range": [0.0, 1.0], + "description": "IoU-distance gating threshold for Stage 1 (high-confidence) association. Passed as cost_limit to the linear assignment solver. Since cost = 1 - IoU, a threshold of 0.8 accepts matches with IoU >= 0.2. Higher values are more lenient (easier to match); lower values require stronger spatial overlap. Tune with detector quality." + }, + "fuse_score": { + "type": "bool", + "default": true, + "description": "When enabled, IoU similarity is multiplied by the detection confidence score before matching: fused_cost = 1 - (iou_similarity * detection_score). This biases matching toward high-confidence detections. Disable if your detector's confidence scores are poorly calibrated or inconsistent." + } +} diff --git a/collab_env/tracking_studio/gcs_browser.py b/collab_env/tracking_studio/gcs_browser.py new file mode 100644 index 00000000..96288089 --- /dev/null +++ b/collab_env/tracking_studio/gcs_browser.py @@ -0,0 +1,168 @@ +""" +GCS Video Browser Component + +Provides interface for browsing and downloading videos from Google Cloud Storage. +""" + +from pathlib import Path +from typing import List, Dict +from loguru import logger + +from collab_env.data.gcs_utils import GCSClient + + +class GCSVideoBrowser: + """Browser for selecting and downloading videos from GCS buckets""" + + def __init__(self, credentials_path: str): + """ + Initialize GCS browser. + + Args: + credentials_path: Path to GCS service account credentials JSON + """ + self.gcs = GCSClient(credentials_path=credentials_path) + logger.info("GCS Video Browser initialized") + + def list_buckets(self) -> List[str]: + """ + List all available GCS buckets. + + Returns: + List of bucket names + """ + try: + buckets = self.gcs.list_buckets() + logger.info(f"Found {len(buckets)} buckets") + return buckets + except Exception as e: + logger.error(f"Failed to list buckets: {e}") + return [] + + def list_folders(self, bucket: str, prefix: str = "") -> List[str]: + """ + List immediate subfolders in a bucket path. + + Note: GCS doesn't have real folders - they're just prefixes in object names. + This function extracts unique first-level directory prefixes. + + Args: + bucket: GCS bucket name + prefix: Path prefix within bucket (should end with / if not empty) + + Returns: + List of folder names (relative to prefix) + """ + try: + # Ensure prefix ends with / if not empty + if prefix and not prefix.endswith("/"): + prefix = prefix + "/" + + # Get all objects recursively to find folder-like structures + pattern = f"{bucket}/{prefix}**" if prefix else f"{bucket}/**" + all_paths = self.gcs.glob(pattern) + + # Extract unique immediate subdirectories + unique_folders = set() + for path in all_paths: + # Remove bucket prefix + rel_path = path.replace(f"{bucket}/", "") + + # Remove the current prefix if any + if prefix: + if not rel_path.startswith(prefix): + continue + rel_path = rel_path[len(prefix):] + + # Get first directory component after prefix + if "/" in rel_path: + folder = rel_path.split("/")[0] + if folder: # Skip empty strings + unique_folders.add(folder) + + folder_list = sorted(list(unique_folders)) + logger.info(f"Found {len(folder_list)} folder prefixes in {bucket}/{prefix}") + return folder_list + + except Exception as e: + logger.error(f"Failed to list folders in {bucket}/{prefix}: {e}") + return [] + + def list_videos(self, bucket: str, prefix: str = "") -> List[Dict[str, str]]: + """ + List video files (.mp4, .mov, .avi) in a bucket path. + + Args: + bucket: GCS bucket name + prefix: Path prefix within bucket + + Returns: + List of dicts with video metadata: {name, path, rel_path} + """ + try: + # Build pattern for video files - ensure prefix ends with / if not empty + if prefix and not prefix.endswith("/"): + prefix = prefix + "/" + + # Search for multiple video formats + video_extensions = ["*.mp4", "*.mov", "*.avi", "*.MP4", "*.MOV", "*.AVI"] + all_files = [] + + for ext in video_extensions: + pattern = f"{bucket}/{prefix}**/{ext}" if prefix else f"{bucket}/**/{ext}" + files = self.gcs.glob(pattern) + all_files.extend(files) + + videos = [] + seen_paths = set() # Avoid duplicates from case-insensitive extensions + + for file_path in all_files: + if file_path in seen_paths: + continue + seen_paths.add(file_path) + + # Extract filename + filename = file_path.split("/")[-1] + + # Get relative path from bucket + rel_path = file_path.replace(f"{bucket}/", "") + + videos.append( + { + "name": filename, + "path": file_path, + "rel_path": rel_path, + } + ) + + logger.info(f"Found {len(videos)} videos in {bucket}/{prefix}") + return sorted(videos, key=lambda x: x["name"]) + + except Exception as e: + logger.error(f"Failed to list videos in {bucket}/{prefix}: {e}") + return [] + + def download_video(self, gcs_path: str, local_path: str) -> str: + """ + Download video from GCS to local path. + + Args: + gcs_path: Full GCS path (e.g., "bucket/path/video.mp4" or "gs://bucket/path/video.mp4") + local_path: Local destination path + + Returns: + Local path to downloaded video + """ + try: + # Remove gs:// prefix if present + if gcs_path.startswith("gs://"): + gcs_path = gcs_path[5:] + + logger.info(f"Downloading {gcs_path} to {local_path}") + self.gcs.download_file(gcs_path, local_path, overwrite=True) + logger.info(f"Successfully downloaded video to {local_path}") + return local_path + + except Exception as e: + logger.error(f"Failed to download video from {gcs_path}: {e}") + raise diff --git a/collab_env/tracking_studio/model_manager.py b/collab_env/tracking_studio/model_manager.py new file mode 100644 index 00000000..e2705f6a --- /dev/null +++ b/collab_env/tracking_studio/model_manager.py @@ -0,0 +1,365 @@ +""" +Model Manager Component + +Handles loading and managing detection models (YOLO and Roboflow). +""" + +import os +from pathlib import Path +from typing import List +from loguru import logger + +from ultralytics import YOLO + + +class ModelManager: + """Manager for detection models (YOLO and Roboflow)""" + + def __init__(self, roboflow_api_key: str = None): + """ + Initialize model manager. + + Args: + roboflow_api_key: Roboflow API key (or read from env) + """ + self.roboflow_api_key = roboflow_api_key or os.getenv("ROBOFLOW_API_KEY") + self.local_models_dir = Path("/workspace/models") + + # Check if running locally (models in ~/.cache/ultralytics) + if not self.local_models_dir.exists(): + # Use default Ultralytics cache directory + self.local_models_dir = Path.home() / ".cache" / "ultralytics" + + logger.info(f"Model directory: {self.local_models_dir}") + + def list_local_yolo_models(self) -> List[str]: + """ + Return available YOLO models (YOLO11 and YOLO26 variants). + + Returns: + List of model filenames + """ + # Auto-downloadable models (Ultralytics will download them) + auto_downloadable = [ + "yolo11n.pt", + "yolo11s.pt", + "yolo11m.pt", + ] + + # Models that must exist locally (not auto-downloadable) + local_only = [ + "yolo26n-fast.pt", + "yolo26s-fast.pt", + "yolo26m-fast.pt", + ] + + available = [] + + # Add auto-downloadable models (always available) + available.extend(auto_downloadable) + + # Add local-only models only if they exist + for model in local_only: + model_path = self.local_models_dir / model + if model_path.exists(): + available.append(model) + logger.debug(f"Found local YOLO26 model: {model}") + + logger.info(f"Available YOLO models: {available}") + return available + + def load_yolo_model(self, model_name: str) -> YOLO: + """ + Load YOLO model - will download automatically if available. + + Args: + model_name: Model filename (e.g., "yolo11n.pt", "yolo26n-fast.pt") + + Returns: + Loaded YOLO model + """ + try: + logger.info(f"Loading YOLO model: {model_name}") + # Pass directly to YOLO - it will handle local files or auto-download + model = YOLO(model_name) + logger.info(f"Successfully loaded YOLO model: {model_name}") + return model + + except Exception as e: + logger.error(f"Failed to load YOLO model {model_name}: {e}") + raise ValueError( + f"Failed to load model '{model_name}'.\n\n" + f"Possible solutions:\n" + f"- Check the model name is correct\n" + f"- Download manually and place in {self.local_models_dir}\n" + f"- Use the 'Custom' upload option to upload your .pt file" + ) from e + + def _validate_roboflow_model_id(self, model_id: str) -> str: + """ + Validate and format Roboflow model ID. + + Accepts: + - project/version (e.g., "ratsmerged20260211/1") + - workspace/project/version (e.g., "myworkspace/ratsmerged20260211/1") + + Returns properly formatted model ID. + """ + parts = model_id.split('/') + + if len(parts) == 2: + # project/version format + logger.info(f"Model ID format: project/version ({model_id})") + return model_id + elif len(parts) == 3: + # workspace/project/version format + logger.info(f"Model ID format: workspace/project/version ({model_id})") + return model_id + else: + raise ValueError( + f"Invalid model ID format: {model_id}\n" + f"Expected: 'project/version' or 'workspace/project/version'" + ) + + def load_roboflow_model(self, model_id: str): + """ + Load Roboflow model using Inference SDK or local file path. + + Args: + model_id: Model ID in format "project/version", "workspace/project/version", + or a local file path to a .pt file + + Returns: + Loaded Roboflow model or YOLO model from local file + """ + # Check if model_id is a local file path + if model_id.startswith('/') or model_id.startswith('~') or model_id.endswith('.pt'): + logger.info(f"Loading Roboflow model from local file: {model_id}") + model_path = Path(model_id).expanduser() + + if not model_path.exists(): + raise FileNotFoundError(f"Model file not found: {model_path}") + + logger.info(f"Loading YOLO model from: {model_path}") + model = YOLO(str(model_path)) + logger.info(f"Successfully loaded Roboflow model from local file: {model_id}") + return model + + if not self.roboflow_api_key: + raise ValueError( + "ROBOFLOW_API_KEY not set. Please provide API key in environment or constructor." + ) + + # Validate model ID format + model_id = self._validate_roboflow_model_id(model_id) + + # Try downloading .pt file first (for YOLO native tracking) + # This provides better performance and supports all ByteTrack parameters + try: + logger.info(f"Downloading Roboflow model weights for YOLO native tracking: {model_id}") + model = self._load_roboflow_with_pipeline(model_id) + logger.info(f"Successfully loaded Roboflow model with native tracking: {model_id}") + return model + + except Exception as download_error: + # Fallback to get_model() (inference API) if download fails + logger.warning(f"Download failed: {download_error}") + logger.info(f"Attempting fallback: loading with inference API (Supervision tracking)") + + try: + from inference import get_model + + # Extract project/version from workspace/project/version if needed + parts = model_id.split('/') + if len(parts) == 3: + # workspace/project/version -> project/version + project_version = f"{parts[1]}/{parts[2]}" + logger.info(f"Trying to load Roboflow model with get_model(): {project_version}") + model = get_model(model_id=project_version, api_key=self.roboflow_api_key) + else: + # Already project/version format + logger.info(f"Trying to load Roboflow model with get_model(): {model_id}") + model = get_model(model_id=model_id, api_key=self.roboflow_api_key) + + logger.info(f"Successfully loaded Roboflow model via inference API: {model_id}") + return model + + except ImportError: + logger.error("inference library not installed. Install with: pip install inference") + raise + except Exception as inference_error: + # Both methods failed + logger.error(f"Inference API also failed: {inference_error}") + error_msg = ( + f"Failed to load Roboflow model '{model_id}'.\n\n" + f"Tried:\n" + f"1. Downloading model weights (.pt file): {str(download_error)}\n" + f"2. Loading via inference API: {str(inference_error)}\n\n" + f"Possible solutions:\n" + f"- Verify model ID format: workspace/project/version (e.g., 'dima-sdrkv/ratsmerged20260211/1')\n" + f"- Check model exists at https://app.roboflow.com/\n" + f"- Ensure ROBOFLOW_API_KEY has access to this model\n" + f"- Try uploading the .pt file directly using 'Custom' option" + ) + raise ValueError(error_msg) from inference_error + + def _load_roboflow_with_pipeline(self, model_id: str): + """ + Fallback: Download Roboflow model weights via /ptFile endpoint. + + This downloads the model weights once via API, then runs inference locally. + Much faster than HTTP inference for every frame. + """ + import requests + + logger.info(f"Downloading Roboflow model weights for local inference: {model_id}") + + # Parse model ID to get workspace/project/version + parts = model_id.split('/') + if len(parts) == 2: + # project/version format - need workspace + raise ValueError( + f"Model ID '{model_id}' missing workspace.\n" + f"For model download, use format: workspace/project/version" + ) + elif len(parts) == 3: + # workspace/project/version format + workspace, project, version = parts + else: + raise ValueError(f"Invalid model ID format: {model_id}") + + # Create cache directory for downloaded models + cache_dir = self.local_models_dir / "roboflow_cache" + cache_dir.mkdir(parents=True, exist_ok=True) + + # Check if model already downloaded + model_cache_name = f"{workspace}_{project}_v{version}.pt" + cached_model_path = cache_dir / model_cache_name + + if cached_model_path.exists(): + logger.info(f"Using cached Roboflow model: {cached_model_path}") + return YOLO(str(cached_model_path)) + + # Download model weights from Roboflow using /ptFile endpoint + logger.info("Fetching model weights URL from Roboflow API...") + + try: + # Call /ptFile endpoint to get signed download URL + ptfile_url = f"https://api.roboflow.com/{workspace}/{project}/{version}/ptFile" + logger.info(f"Requesting weights URL from: {ptfile_url}") + + response = requests.get( + ptfile_url, + params={"api_key": self.roboflow_api_key}, + timeout=10 + ) + response.raise_for_status() + + # Parse response to get weightsUrl + data = response.json() + if 'weightsUrl' not in data: + raise ValueError(f"No weightsUrl in response: {data}") + + weights_url = data['weightsUrl'] + logger.info(f"Got weights URL, downloading...") + + # Download the .pt file from signed URL + response = requests.get(weights_url, stream=True, timeout=120) + response.raise_for_status() + + # Save to cache + with open(cached_model_path, 'wb') as f: + for chunk in response.iter_content(chunk_size=8192): + f.write(chunk) + + logger.info(f"Downloaded model weights: {cached_model_path} ({cached_model_path.stat().st_size} bytes)") + + # Load with Ultralytics YOLO + model = YOLO(str(cached_model_path)) + logger.info(f"Successfully loaded Roboflow model for local inference: {model_id}") + return model + + except requests.exceptions.HTTPError as e: + error_msg = ( + f"Failed to download Roboflow model weights for '{model_id}'.\n\n" + f"HTTP Error {e.response.status_code}: {e.response.text[:200]}\n\n" + f"Possible solutions:\n" + f"1. Verify model ID format is workspace/project/version\n" + f"2. Check ROBOFLOW_API_KEY has access to this model\n" + f"3. Ensure model exists at https://app.roboflow.com/\n" + f"4. Upload model weights manually via 'Custom' option" + ) + logger.error(error_msg) + raise ValueError(error_msg) from e + except Exception as e: + error_msg = ( + f"Failed to download Roboflow model weights for '{model_id}'.\n\n" + f"Error: {str(e)}\n\n" + f"Try uploading model weights manually via 'Custom' option." + ) + logger.error(error_msg) + raise ValueError(error_msg) from e + + def list_roboflow_project_models(self, project_id: str) -> List[str]: + """ + Query Roboflow API for available model versions in a project. + + Args: + project_id: Project ID in format "workspace/project" (e.g., "dima-sdrkv/ratsmerged20260211") + + Returns: + List of version numbers (e.g., ["1", "2", "3"]) + """ + import requests + + if not self.roboflow_api_key: + raise ValueError("ROBOFLOW_API_KEY not set") + + parts = project_id.split('/') + if len(parts) != 2: + raise ValueError("Project ID must be in format: workspace/project") + + workspace, project = parts + + try: + url = f"https://api.roboflow.com/{workspace}/{project}" + logger.info(f"Querying Roboflow project models: {url}") + + response = requests.get( + url, + params={"api_key": self.roboflow_api_key}, + timeout=10 + ) + response.raise_for_status() + + data = response.json() + + # Extract version numbers from response + versions = [] + if 'versions' in data: + for version_data in data['versions']: + # Try different fields that might contain the version number + version_num = version_data.get('id') + + # If id is a full path (workspace/project/version), extract just the version + if version_num and isinstance(version_num, str) and '/' in version_num: + version_num = version_num.split('/')[-1] # Get last part + + # Also check for a 'version' field + if not version_num: + version_num = version_data.get('version') + + if version_num: + versions.append(str(version_num)) + + logger.info(f"Found {len(versions)} versions: {versions}") + return sorted(versions, key=lambda x: int(x) if x.isdigit() else 0, reverse=True) + + except requests.exceptions.HTTPError as e: + error_msg = f"Failed to query Roboflow project: HTTP {e.response.status_code}" + logger.error(error_msg) + raise ValueError(error_msg) from e + except Exception as e: + error_msg = f"Failed to query Roboflow project: {str(e)}" + logger.error(error_msg) + raise ValueError(error_msg) from e diff --git a/collab_env/tracking_studio/video_converter.py b/collab_env/tracking_studio/video_converter.py new file mode 100644 index 00000000..e7ed06b0 --- /dev/null +++ b/collab_env/tracking_studio/video_converter.py @@ -0,0 +1,96 @@ +""" +Video Format Converter Component + +Ensures videos are in H.264 format for browser compatibility. +""" + +import subprocess +from pathlib import Path +from loguru import logger + + +def needs_conversion(video_path: Path) -> bool: + """ + Check if video needs H.264 conversion. + + Args: + video_path: Path to video file + + Returns: + True if conversion needed, False otherwise + """ + try: + cmd = [ + "ffprobe", + "-v", + "error", + "-select_streams", + "v:0", + "-show_entries", + "stream=codec_name", + "-of", + "default=noprint_wrappers=1:nokey=1", + str(video_path), + ] + + result = subprocess.run(cmd, capture_output=True, text=True, check=True) + codec = result.stdout.strip() + + logger.info(f"Video codec: {codec}") + return codec != "h264" + + except subprocess.CalledProcessError as e: + logger.error(f"Failed to check video codec: {e}") + # Assume conversion needed if check fails + return True + except FileNotFoundError: + logger.error("ffprobe not found. Please install ffmpeg.") + raise + + +def convert_to_h264(input_path: Path, output_path: Path) -> Path: + """ + Convert video to H.264 format using ffmpeg. + + Args: + input_path: Original video file + output_path: Output path for converted video + + Returns: + Path to converted video + """ + try: + logger.info(f"Converting {input_path} to H.264 format") + + cmd = [ + "ffmpeg", + "-i", + str(input_path), + "-c:v", + "libx264", + "-preset", + "fast", + "-crf", + "23", + "-c:a", + "aac", + "-b:a", + "128k", + "-movflags", + "+faststart", # Web optimization + "-y", # Overwrite output + str(output_path), + ] + + subprocess.run(cmd, check=True, capture_output=True) + + logger.info(f"Successfully converted video to {output_path}") + return output_path + + except subprocess.CalledProcessError as e: + logger.error(f"Failed to convert video: {e}") + logger.error(f"stderr: {e.stderr.decode() if e.stderr else 'N/A'}") + raise + except FileNotFoundError: + logger.error("ffmpeg not found. Please install ffmpeg.") + raise diff --git a/collab_env/tracking_studio/video_processor.py b/collab_env/tracking_studio/video_processor.py new file mode 100644 index 00000000..e04746da --- /dev/null +++ b/collab_env/tracking_studio/video_processor.py @@ -0,0 +1,374 @@ +""" +Video Processor Component + +Core tracking pipeline with ByteTrack. +""" + +import asyncio +import threading +import cv2 +import supervision as sv +from ultralytics import YOLO +import pandas as pd +from pathlib import Path +from typing import Callable, Dict, List, Union, Any +import numpy as np +import tempfile +import yaml +from loguru import logger + + +class VideoTracker: + """Video tracking processor with detection and tracking""" + + def __init__( + self, + model: Union[YOLO, Any], # YOLO or Roboflow model + tracker_config: Dict, # ByteTrack parameters + confidence: float = 0.5, + frame_callback: Callable[[np.ndarray, int, int], None] = None, + stop_event: threading.Event = None, + pause_event: threading.Event = None, + skip_frames_event: Dict = None, + ): + """ + Initialize video tracker. + + Args: + model: Detection model (YOLO or Roboflow) + tracker_config: Tracker configuration dict + confidence: Detection confidence threshold + frame_callback: Async callback for frame updates (frame, frame_idx, total_frames) + stop_event: Threading event to signal hard stop + pause_event: Threading event to signal pause/resume + skip_frames_event: Dict with skip_amount for forward seeking + """ + self.model = model + self.confidence = confidence + self.frame_callback = frame_callback + self.stop_event = stop_event or threading.Event() + self.pause_event = pause_event or threading.Event() + self.skip_frames_event = skip_frames_event or {"skip_amount": 0} + + # Store tracker config for use with model.track() + self.tracker_config = tracker_config + + # Check if model supports native tracking + self.use_native_tracking = isinstance(model, YOLO) + + # For Roboflow inference models (fallback), initialize supervision tracker + if not self.use_native_tracking: + logger.info("Using supervision ByteTrack (Roboflow inference model fallback)") + self.tracker = sv.ByteTrack( + track_activation_threshold=tracker_config.get("track_high_thresh", 0.25), + lost_track_buffer=tracker_config.get("track_buffer", 30), + minimum_matching_threshold=tracker_config.get("match_thresh", 0.8), + minimum_consecutive_frames=1, + frame_rate=30, + ) + self.tracker_yaml_path = None + else: + logger.info("Using Ultralytics native ByteTrack (supports all parameters)") + self.tracker = None + # Create temporary ByteTrack YAML config from parameters + self.tracker_yaml_path = self._create_bytetrack_config(tracker_config) + + # Fast-forward: Skip frames for faster preview + self.skip_frames = tracker_config.get("skip_frames", 1) # 1 = process every frame + + # Annotators for visualization + self.box_annotator = sv.BoxAnnotator() + self.label_annotator = sv.LabelAnnotator() + + logger.info( + f"VideoTracker initialized (confidence: {self.confidence}, native_tracking: {self.use_native_tracking})" + ) + + def _create_bytetrack_config(self, config: Dict) -> str: + """ + Create a temporary ByteTrack YAML config file from parameters. + + Args: + config: Tracker configuration dict + + Returns: + Path to temporary YAML config file + """ + # Map our parameter names to Ultralytics ByteTrack YAML format + bytetrack_yaml = { + "tracker_type": "bytetrack", + "track_high_thresh": config.get("track_high_thresh", 0.25), + "track_low_thresh": config.get("track_low_thresh", 0.1), + "new_track_thresh": config.get("new_track_thresh", 0.25), + "track_buffer": config.get("track_buffer", 30), + "match_thresh": config.get("match_thresh", 0.8), + "fuse_score": config.get("fuse_score", True), + } + + # Create temporary YAML file + temp_file = tempfile.NamedTemporaryFile( + mode='w', + suffix='.yaml', + delete=False, + prefix='bytetrack_' + ) + + with temp_file as f: + yaml.dump(bytetrack_yaml, f, default_flow_style=False) + + logger.info(f"Created ByteTrack config: {temp_file.name}") + logger.debug(f"Config values: {bytetrack_yaml}") + + return temp_file.name + + def _process_video_sync( + self, video_path: str, output_dir: str, event_loop + ) -> Dict[str, Any]: + """ + Synchronous video processing function (runs in background thread). + + Args: + video_path: Path to input video + output_dir: Directory for output CSV + event_loop: Main asyncio event loop for scheduling UI updates + + Returns: + Dict with tracking_csv path and stats + """ + logger.info(f"Processing video in background thread: {video_path}") + + cap = cv2.VideoCapture(video_path) + if not cap.isOpened(): + raise ValueError(f"Failed to open video: {video_path}") + + fps = cap.get(cv2.CAP_PROP_FPS) + total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + + logger.info( + f"Video info: {total_frames} frames, {fps} fps, {width}x{height}" + ) + + detections_list = [] + tracking_list = [] + + frame_idx = 0 + while frame_idx < total_frames: + # Check if stop was requested (hard stop) + if self.stop_event.is_set(): + logger.info(f"Stop requested at frame {frame_idx}, stopping processing") + break + + # Check if pause was requested + while self.pause_event.is_set(): + import time + time.sleep(0.1) # Wait while paused + if self.stop_event.is_set(): + break + + # Check if seek was requested (forward or backward) + if self.skip_frames_event["skip_amount"] != 0: + skip_to = frame_idx + self.skip_frames_event["skip_amount"] + # Clamp to valid range + skip_to = max(0, min(skip_to, total_frames - 1)) + logger.info(f"Seeking from frame {frame_idx} to {skip_to}") + cap.set(cv2.CAP_PROP_POS_FRAMES, skip_to) + self.skip_frames_event["skip_amount"] = 0 # Reset + frame_idx = skip_to + continue + + ret, frame = cap.read() + if not ret: + logger.warning(f"Failed to read frame {frame_idx}, stopping") + break + + # Fast-forward: Skip frames if requested + if self.skip_frames > 1 and frame_idx % self.skip_frames != 0: + frame_idx += 1 + continue + + # 1. Run detection and tracking + try: + if self.use_native_tracking: + # Use Ultralytics native tracking (supports all ByteTrack parameters) + results = self.model.track( + source=frame, + conf=self.confidence, + persist=True, # Maintain track IDs across frames + tracker=self.tracker_yaml_path, # Custom ByteTrack config + verbose=False, + )[0] + + # Convert to supervision Detections (with track IDs) + tracked_detections = sv.Detections.from_ultralytics(results) + + # Also get detections without tracking for stats + detections = tracked_detections + else: + # Roboflow inference model (fallback to supervision ByteTrack) + logger.debug(f"Running Roboflow inference on frame {frame_idx}...") + results = self.model.infer(frame, confidence=self.confidence)[0] + detections = sv.Detections.from_inference(results) + logger.debug(f"Frame {frame_idx}: {len(detections)} detections") + + # Update tracker (adds track IDs via supervision ByteTrack) + tracked_detections = self.tracker.update_with_detections(detections) + + except Exception as e: + logger.error(f"Detection/tracking failed on frame {frame_idx}: {e}", exc_info=True) + detections = sv.Detections.empty() + tracked_detections = sv.Detections.empty() + + # 2. Save raw detections (for stats only, not exported) + for i, (bbox, conf, class_id) in enumerate( + zip(detections.xyxy, detections.confidence, detections.class_id) + ): + detections_list.append( + { + "frame": frame_idx, + "x1": bbox[0], + "y1": bbox[1], + "x2": bbox[2], + "y2": bbox[3], + "confidence": conf, + "class": class_id, + } + ) + + # 3. Tracked detections now have track IDs (from native tracking or supervision) + + # 4. Save tracking with IDs (matches output_tracked_bboxes_csv format) + # Only save if we have track IDs (handles cases where no detections exist) + if tracked_detections.tracker_id is not None and len(tracked_detections) > 0: + for bbox, track_id, conf, class_id in zip( + tracked_detections.xyxy, + tracked_detections.tracker_id, + tracked_detections.confidence, + tracked_detections.class_id, + ): + tracking_list.append( + { + "track_id": int(track_id), + "frame": frame_idx, + "x1": int(bbox[0]), + "y1": int(bbox[1]), + "x2": int(bbox[2]), + "y2": int(bbox[3]), + "confidence": float(conf), + "class": int(class_id), + } + ) + + # 5. Annotate frame for display + annotated_frame = frame.copy() + annotated_frame = self.box_annotator.annotate( + annotated_frame, tracked_detections + ) + + # Create labels with track IDs + labels = [ + f"#{track_id} {conf:.2f}" + for track_id, conf in zip( + tracked_detections.tracker_id, tracked_detections.confidence + ) + ] + annotated_frame = self.label_annotator.annotate( + annotated_frame, tracked_detections, labels=labels + ) + + # 6. Send frame to UI (schedule callback in main event loop) + if self.frame_callback and event_loop: + # Schedule callback in main event loop from background thread + future = asyncio.run_coroutine_threadsafe( + self.frame_callback(annotated_frame, frame_idx, total_frames), + event_loop + ) + # Wait for UI update to complete (with timeout to prevent blocking) + try: + future.result(timeout=2.0) + except Exception as e: + logger.warning(f"Frame callback failed: {e}") + + # Increment frame counter for next iteration + frame_idx += 1 + + cap.release() + + # Cleanup temporary tracker config file if created + if self.tracker_yaml_path: + try: + import os + os.unlink(self.tracker_yaml_path) + logger.debug(f"Cleaned up temporary tracker config: {self.tracker_yaml_path}") + except Exception as e: + logger.warning(f"Failed to cleanup tracker config: {e}") + + logger.info( + f"Processing complete: {total_frames} frames, " + f"{len(detections_list)} detections, " + f"{len(set(t['track_id'] for t in tracking_list))} unique tracks" + ) + + # 7. Save tracking CSV (matches output_tracked_bboxes_csv format) + tracking_df = pd.DataFrame(tracking_list) + + output_path = Path(output_dir) + output_path.mkdir(parents=True, exist_ok=True) + + # Only save tracking CSV (not detections - user doesn't need them) + tracking_csv = output_path / "tracking.csv" + + if len(tracking_list) > 0: + # Ensure column order matches: track_id,frame,x1,y1,x2,y2,confidence,class + tracking_df = tracking_df[ + ["track_id", "frame", "x1", "y1", "x2", "y2", "confidence", "class"] + ] + tracking_df.to_csv(tracking_csv, index=False) + logger.info(f"Saved tracking CSV to {tracking_csv}") + else: + # Create empty CSV with correct headers + pd.DataFrame( + columns=["track_id", "frame", "x1", "y1", "x2", "y2", "confidence", "class"] + ).to_csv(tracking_csv, index=False) + logger.warning("No tracks found, saved empty CSV") + + return { + "tracking_csv": str(tracking_csv), + "stats": { + "total_frames": total_frames, + "total_detections": len(detections_list), + "unique_tracks": ( + tracking_df["track_id"].nunique() if len(tracking_list) > 0 else 0 + ), + "fps": fps, + }, + } + + async def process_video_realtime( + self, video_path: str, output_dir: str + ) -> Dict[str, Any]: + """ + Process video frame-by-frame with real-time UI updates. + + This runs the heavy processing in a background thread to prevent + blocking the asyncio event loop and WebSocket connections. + + Args: + video_path: Path to input video + output_dir: Directory for output CSV + + Returns: + Dict with tracking_csv path and stats + """ + # Get current event loop for scheduling UI updates from background thread + loop = asyncio.get_running_loop() + + # Run processing in background thread + logger.info("Starting video processing in background thread...") + result = await asyncio.to_thread( + self._process_video_sync, video_path, output_dir, loop + ) + + logger.info("Video processing complete") + return result diff --git a/docs/tracking/tracking_web_gui.md b/docs/tracking/tracking_web_gui.md new file mode 100644 index 00000000..bb885994 --- /dev/null +++ b/docs/tracking/tracking_web_gui.md @@ -0,0 +1,345 @@ +# Tracking Studio Web GUI + +Interactive web-based application for real-time video object detection and tracking using YOLO and ByteTrack. + +## Overview + +The Tracking Studio provides a user-friendly interface for: +- Loading videos from Google Cloud Storage or local uploads +- Selecting detection models (YOLO, Roboflow, or custom .pt files) +- Tuning ByteTrack parameters in real-time +- Visualizing tracking results with live preview +- Exporting tracking data to CSV + +## Quick Start + +### Running the Application + +```bash +# From the repository root +python scripts/run_tracking_studio.py + +# Or directly +python -m collab_env.tracking_studio.app +``` + +The application will start on `http://localhost:8080` + +### Environment Setup + +**Optional: For Google Cloud Storage integration** +```bash +export GCS_CREDENTIALS=/path/to/credentials.json +``` + +**Optional: For Roboflow models** +```bash +export ROBOFLOW_API_KEY=your_api_key_here +``` + +## Workflow + +### 1. Load Video + +**From Google Cloud Storage:** +1. Select **Bucket** (e.g., `collab-data-463313`) +2. Select **Folder** (optional subfolder) +3. Select **Video** from the list +4. Click **Load Video** + +**From Local Upload:** +1. Click **Upload** button +2. Select video file (.mp4, .mov, .avi) +3. Click **Load Video** + +The first frame will display in the preview area. + +### 2. Load Model + +Choose one of three model sources: + +#### YOLO Models +- Enter any YOLO model name (e.g., `yolo11n.pt`, `yolo26n.pt`) +- Models will auto-download if available from Ultralytics +- Click **Load Model** + +#### Roboflow Models +- Enter **Project ID** in format: `workspace/project` +- Click **List Models** to fetch available versions +- Select a **Version** from dropdown +- Click **Load Model** + +**Supported types:** +- Object detection models (standard) +- Instance segmentation models (extracts bounding boxes only) + +#### Custom Models +- Upload your own `.pt` file +- Click **Load Model** + +### 3. Configure Parameters + +**Detection Confidence** +- Threshold for detection scores (0.1 - 0.9) +- Higher = fewer false positives, may miss detections +- Lower = more detections, may include noise + +**ByteTrack Parameters** (see [ByteTrack Parameters](#bytetrack-parameters) below) + +**Skip Frames** +- Process every Nth frame (1 = all frames, 30 = every 30th frame) +- Use higher values for faster preview on long videos +- Final tracking still captures data for skipped frames + +**Display Update** +- Update preview every Nth frame (1-30, default: 10) +- Lower values = smoother preview (more network traffic) +- Higher values = less frequent updates (lower bandwidth) +- At 30fps video: 10 frames = ~3 updates/second, 5 frames = ~6 updates/second + +### 4. Start Tracking + +1. Click **Start Tracking** +2. Watch live preview with bounding boxes and track IDs +3. Use **Pause** to temporarily halt processing +4. Use **Stop** to terminate early +5. Drag the time slider to jump to specific frames + +### 5. Export Results + +When complete, click **Download CSV** to save tracking data. + +## ByteTrack Parameters + +ByteTrack uses a two-stage association algorithm to track objects across frames: + +### `track_high_thresh` (default: 0.25) +Detection confidence threshold separating high-confidence and low-confidence detections. +- Detections **above** this → Stage 1 (primary IoU matching) +- Detections **between** `track_low_thresh` and this → Stage 2 (secondary matching) +- **Raise** to restrict primary matching to most confident detections +- **Lower** to feed more detections into Stage 1 + +### `track_low_thresh` (default: 0.1) +Absolute minimum detection confidence. +- Detections **below** this are discarded entirely +- **Lower** to recover marginal detections (more noise) +- **Raise** to filter weak false positives + +### `new_track_thresh` (default: 0.25) +Minimum confidence required to initialize a new track. +- Unmatched detections from Stage 1 must exceed this to spawn new track IDs +- **Higher** prevents spurious tracks from false positives +- **Lower** allows tracks to start from weaker detections + +### `track_buffer` (default: 30) +Number of frames a lost track is kept alive before deletion. +- Internally scaled by frame rate: `max_time_lost = int(fps / 30.0 * track_buffer)` +- **Higher** values let tracks survive longer occlusions (more ID switches risk) +- **Lower** values remove lost tracks faster + +### `match_thresh` (default: 0.8) +IoU-distance gating threshold for Stage 1 association. +- Cost = 1 - IoU, so threshold of 0.8 accepts matches with IoU ≥ 0.2 +- **Higher** values are more lenient (easier to match) +- **Lower** values require stronger spatial overlap +- Tune based on detector quality + +### `fuse_score` (default: true) +Multiply IoU similarity by detection confidence before matching. +- Formula: `fused_cost = 1 - (iou_similarity * detection_score)` +- **Enable** to bias matching toward high-confidence detections +- **Disable** if detector's confidence scores are poorly calibrated + +## Output Format + +### Tracking CSV + +Format: `tracking.csv` + +| Column | Type | Description | +|--------|------|-------------| +| `track_id` | int | Unique object track identifier | +| `frame` | int | Frame number (0-indexed) | +| `x1` | int | Bounding box top-left X | +| `y1` | int | Bounding box top-left Y | +| `x2` | int | Bounding box bottom-right X | +| `y2` | int | Bounding box bottom-right Y | +| `confidence` | float | Detection confidence score | +| `class` | int | Object class ID | + +Example: +```csv +track_id,frame,x1,y1,x2,y2,confidence,class +1,0,245,150,345,280,0.87,0 +1,1,247,152,346,281,0.85,0 +2,1,450,200,550,320,0.92,0 +``` + +## Model Support + +### YOLO Models (Ultralytics) +- ✅ YOLO11 series (`yolo11n.pt`, `yolo11s.pt`, `yolo11m.pt`, etc.) +- ✅ YOLO26 series (`yolo26n.pt`, `yolo26s.pt`, etc.) +- ✅ Custom trained YOLO models (.pt files) +- Uses **native Ultralytics tracking** (supports all 6 ByteTrack parameters) + +### Roboflow Models +- ✅ Object detection models +- ✅ Instance segmentation models (bounding boxes only, masks ignored) +- Uses **model download + local inference** (YOLO-compatible weights) +- Fallback to **supervision ByteTrack** if native tracking unavailable + +### Custom Models +- ✅ Upload any YOLO-compatible `.pt` file +- Must be trainable with Ultralytics YOLO framework + +## Architecture + +### Components + +**Frontend: NiceGUI** +- Reactive web interface +- Real-time frame updates via WebSocket +- Slider-based parameter tuning + +**Backend: FastAPI (via NiceGUI)** +- Async video processing with `asyncio.to_thread()` +- Background thread handles heavy CV operations +- Event loop scheduling for UI updates + +**Video Processing: [video_processor.py](../../collab_env/tracking_studio/video_processor.py)** +- OpenCV video capture +- Frame-by-frame detection + tracking +- Supervision library for annotation +- Temporary YAML config for ByteTrack parameters + +**Model Management: [model_manager.py](../../collab_env/tracking_studio/model_manager.py)** +- YOLO model loading (Ultralytics) +- Roboflow model loading (inference SDK + fallback download) +- Model caching for faster reloads + +## Video Format Support + +### Supported Formats +- MP4 (H.264 codec) - **recommended** +- MOV (QuickTime) +- AVI (uncompressed) + +### Automatic Conversion +Videos not in H.264 format are automatically converted on load: +- Source: mjpeg, raw, etc. +- Target: H.264 MP4 (1080p max, 30fps) +- Uses FFmpeg via `video_converter.py` + +## Playback Controls + +### Real-time Controls +- **Start Tracking**: Begin processing video +- **Pause**: Temporarily halt (can resume) +- **Stop**: Hard stop (terminates processing) + +### Seeking +- Drag **time slider** during processing to jump to specific frame +- Shows raw frame preview during drag (no tracking) +- Releases to seek tracker forward/backward + +### Preview Updates +- Updates every 10 frames for performance (~3x per second at 30fps) +- Shows annotated frames with bounding boxes and track IDs + +## Performance Tips + +### For Long Videos +1. Use **Skip Frames** to process every Nth frame for faster preview +2. Increase **track_buffer** to maintain tracks across skipped frames +3. Lower **detection confidence** if missing objects + +### For Crowded Scenes +1. Raise **track_high_thresh** to focus on confident detections +2. Lower **match_thresh** to require tighter spatial overlap +3. Increase **new_track_thresh** to reduce spurious tracks + +### For Fast Motion +1. Lower **match_thresh** to accept looser spatial matches +2. Increase **track_buffer** to keep tracks alive longer +3. Process all frames (Skip = 1) for smoother tracking + +## Troubleshooting + +### "No detections found" +- Lower **detection confidence** slider +- Check video quality and lighting +- Try different model (e.g., `yolo11m.pt` instead of `yolo11n.pt`) + +### "Too many false positives" +- Raise **detection confidence** +- Increase **new_track_thresh** +- Raise **track_high_thresh** + +### "Track IDs jumping/switching" +- Lower **track_high_thresh** to feed more detections into Stage 1 +- Raise **match_thresh** for more lenient matching +- Increase **track_buffer** to keep lost tracks alive longer +- Enable **fuse_score** if disabled + +### "Video conversion failed" +- Check FFmpeg installation: `ffmpeg -version` +- Ensure video file is not corrupted +- Try converting manually: `ffmpeg -i input.mp4 -c:v libx264 output.mp4` + +### "Model loading failed" +- YOLO: Check model name spelling and internet connection +- Roboflow: Verify `ROBOFLOW_API_KEY` is set and has access +- Custom: Ensure `.pt` file is YOLO-compatible format + +## Advanced Usage + +### Running with Docker + +```bash +# Build image +docker build -t tracking-studio . + +# Run with GCS credentials +docker run -p 8080:8080 \ + -v /path/to/credentials.json:/workspace/config/credentials.json \ + -e GCS_CREDENTIALS=/workspace/config/credentials.json \ + tracking-studio +``` + +### Batch Processing + +For offline batch processing without the GUI, use the [full tracking pipeline notebook](full_pipeline.ipynb) or direct API: + +```python +from collab_env.tracking_studio.video_processor import VideoTracker +from collab_env.tracking_studio.model_manager import ModelManager + +# Load model +manager = ModelManager() +model = manager.load_yolo_model("yolo11n.pt") + +# Configure tracker +tracker_config = { + "track_high_thresh": 0.25, + "track_low_thresh": 0.1, + "new_track_thresh": 0.25, + "track_buffer": 30, + "match_thresh": 0.8, + "fuse_score": True, +} + +tracker = VideoTracker(model=model, tracker_config=tracker_config, confidence=0.5) + +# Process video +results = await tracker.process_video_realtime("input.mp4", "/tmp/output") +print(f"Saved to: {results['tracking_csv']}") +``` + +## References + +- [ByteTrack Paper](https://arxiv.org/abs/2110.06864) +- [Ultralytics YOLO](https://docs.ultralytics.com/) +- [Supervision Library](https://supervision.roboflow.com/) +- [Roboflow Inference](https://inference.roboflow.com/) diff --git a/pyproject.toml b/pyproject.toml index 4d05e844..0751b6db 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -54,7 +54,10 @@ dependencies = [ "rich", "pyarrow", "starbars", - "importlib-metadata" + "importlib-metadata", + "nicegui>=1.4.0", + "supervision>=0.18.0", + "inference>=0.28.0" # Upgraded for yolo26n-seg support ] [tool.setuptools] # NEW @@ -70,6 +73,7 @@ packages = [ "collab_env.sim.util", "collab_env.tracking", "collab_env.tracking.model", + "collab_env.tracking_studio", "collab_env.utils" ] diff --git a/scripts/tracking/bytetrack_video_inference.py b/scripts/tracking/bytetrack_video_inference.py new file mode 100644 index 00000000..f0bf1e3f --- /dev/null +++ b/scripts/tracking/bytetrack_video_inference.py @@ -0,0 +1,170 @@ +"""Real-time ByteTracker inference on video with live visualization""" + +import cv2 +import numpy as np +import supervision as sv +from ultralytics import YOLO +from pathlib import Path +import argparse + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--model_path", type=str, required=True, help="Path to YOLO model") + parser.add_argument("path_to_video", type=str, help="Path to video file") + parser.add_argument("--confidence", type=float, default=0.5, help="Confidence threshold") + parser.add_argument("--track_activation", type=float, default=0.2, help="Track activation threshold") + parser.add_argument("--lost_buffer", type=int, default=90, help="Lost track buffer frames") + parser.add_argument("--match_threshold", type=float, default=0.8, help="Minimum matching threshold") + parser.add_argument("--min_frames", type=int, default=5, help="Minimum consecutive frames") + + args = parser.parse_args() + + # Load model + model = YOLO(args.model_path) + + # Initialize tracker + tracker = sv.ByteTrack( + track_activation_threshold=args.track_activation, + lost_track_buffer=args.lost_buffer, + minimum_matching_threshold=args.match_threshold, + minimum_consecutive_frames=args.min_frames + ) + + # Initialize annotators + box_annotator = sv.BoxAnnotator(thickness=1) + mask_annotator = sv.MaskAnnotator(opacity=0.4) + label_annotator = sv.LabelAnnotator( + text_scale=0.3, + text_thickness=1, + text_padding=3, + text_position=sv.Position.TOP_LEFT, + color=sv.Color.BLACK, + text_color=sv.Color.WHITE, + border_radius=2, + smart_position=True + ) + + # Open video + cap = cv2.VideoCapture(args.path_to_video) + + if not cap.isOpened(): + print(f"Error: Could not open video {args.path_to_video}") + return + + # Get video properties + fps = cap.get(cv2.CAP_PROP_FPS) + total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + + print(f"Video: {width}x{height} @ {fps:.2f} fps, {total_frames} frames") + print(f"Confidence: {args.confidence}") + print(f"ByteTracker params: activation={args.track_activation}, buffer={args.lost_buffer}, match={args.match_threshold}, min_frames={args.min_frames}") + print("\nPress 'q' to quit, 'p' to pause/unpause, SPACE to step frame when paused") + + frame_idx = 0 + paused = False + + # Calculate padding for YOLO + target_height = ((height + 31) // 32) * 32 + target_width = ((width + 31) // 32) * 32 + + while True: + if not paused: + ret, frame = cap.read() + if not ret: + print("\nEnd of video") + break + + # Run YOLO detection + results = model.predict( + source=frame, + conf=args.confidence, + verbose=False, + imgsz=(target_width, target_height) + ) + + # Process detections + if results and results[0].boxes: + boxes = results[0].boxes + masks = results[0].masks + + # Resize masks to original frame size + if masks is not None: + mask_array = masks.data.cpu().numpy() + resized_masks = np.zeros((mask_array.shape[0], height, width)) + for i in range(mask_array.shape[0]): + resized_masks[i] = cv2.resize( + mask_array[i], + (width, height), + interpolation=cv2.INTER_LINEAR + ) + resized_masks = resized_masks > 0.5 + + # Create detections with masks + detections = sv.Detections( + xyxy=boxes.xyxy.cpu().numpy(), + mask=resized_masks, + confidence=boxes.conf.cpu().numpy(), + class_id=boxes.cls.cpu().numpy().astype(np.int32), + ) + else: + # No masks, just boxes + detections = sv.Detections( + xyxy=boxes.xyxy.cpu().numpy(), + confidence=boxes.conf.cpu().numpy(), + class_id=boxes.cls.cpu().numpy().astype(np.int32), + ) + + # Update tracker + detections = tracker.update_with_detections(detections) + + # Create labels + labels = [ + f"#{int(tid)} ({conf:.2f})" + for tid, conf in zip(detections.tracker_id, detections.confidence) + ] + + # Annotate frame + annotated_frame = frame.copy() + if detections.mask is not None: + annotated_frame = mask_annotator.annotate(annotated_frame, detections=detections) + else: + annotated_frame = box_annotator.annotate(annotated_frame, detections=detections) + annotated_frame = label_annotator.annotate(annotated_frame, detections=detections, labels=labels) + + # Add info overlay + info_text = f"Frame: {frame_idx}/{total_frames} | Detections: {len(detections)}" + cv2.putText(annotated_frame, info_text, (10, 20), + cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1) + else: + annotated_frame = frame.copy() + info_text = f"Frame: {frame_idx}/{total_frames} | Detections: 0" + cv2.putText(annotated_frame, info_text, (10, 20), + cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 1) + + frame_idx += 1 + + # Display frame + cv2.imshow('ByteTracker Inference', annotated_frame) + + # Handle key presses + key = cv2.waitKey(1 if not paused else 0) & 0xFF + + if key == ord('q'): + print("\nQuitting...") + break + elif key == ord('p'): + paused = not paused + print(f"\n{'Paused' if paused else 'Resumed'}") + elif key == ord(' ') and paused: + # Step one frame forward + ret, frame = cap.read() + if ret: + frame_idx += 1 + + cap.release() + cv2.destroyAllWindows() + +if __name__ == "__main__": + main() diff --git a/scripts/tracking/configs/botsort_thermal_rats.yaml b/scripts/tracking/configs/botsort_thermal_rats.yaml new file mode 100644 index 00000000..ec33ec3f --- /dev/null +++ b/scripts/tracking/configs/botsort_thermal_rats.yaml @@ -0,0 +1,21 @@ +# Custom BoT-SORT config optimized for thermal rat tracking +# BoT-SORT handles camera motion better than ByteTrack + +tracker_type: botsort + +# Detection thresholds +track_high_thresh: 0.3 # High confidence threshold +track_low_thresh: 0.1 # Low confidence threshold +new_track_thresh: 0.2 # Threshold for creating new tracks + +# Track management +track_buffer: 120 # Frames to keep lost tracks alive +match_thresh: 0.7 # IoU threshold for matching + +# BoT-SORT specific +cmc_method: sparseOptFlow # Camera motion compensation method +with_reid: False # ReID features + +# Kalman filter settings (for motion prediction) +std_weight_position: 0.05 +std_weight_velocity: 0.00625 diff --git a/scripts/tracking/configs/bytetrack_thermal_rats.yaml b/scripts/tracking/configs/bytetrack_thermal_rats.yaml new file mode 100644 index 00000000..7c144a15 --- /dev/null +++ b/scripts/tracking/configs/bytetrack_thermal_rats.yaml @@ -0,0 +1,16 @@ +# Custom ByteTrack config optimized for thermal rat tracking +# Lower thresholds to reduce dropped detections + +tracker_type: bytetrack + +# Detection thresholds +track_high_thresh: 0.3 # High confidence threshold (lowered from default 0.5) +track_low_thresh: 0.1 # Low confidence threshold for re-identification (lowered from 0.1) +new_track_thresh: 0.2 # Threshold for creating new tracks (lowered from 0.4) + +# Track management +track_buffer: 120 # Frames to keep lost tracks alive (increased from 30) +match_thresh: 0.7 # IoU threshold for matching (lowered from 0.8 for more lenient matching) + +# Optional features +with_reid: False # ReID features (not needed for rats) diff --git a/scripts/tracking/roboflow_video_inference.py b/scripts/tracking/roboflow_video_inference.py new file mode 100644 index 00000000..3f56c643 --- /dev/null +++ b/scripts/tracking/roboflow_video_inference.py @@ -0,0 +1,23 @@ +# Import the InferencePipeline object +from inference import InferencePipeline +# Import the built in render_boxes sink for visualizing results +from inference.core.interfaces.stream.sinks import render_boxes + +if __name__ == "__main__": + + import argparse + parser = argparse.ArgumentParser() + parser.add_argument("--model_id", type=str, required=True) + parser.add_argument("path_to_video", type=str) + + args = parser.parse_args() + + + # initialize a pipeline object + pipeline = InferencePipeline.init( + model_id=args.model_id, # Roboflow model to use + video_reference=args.path_to_video, # Path to video, device id (int, usually 0 for built in webcams), or RTSP stream url + on_prediction=render_boxes, # Function to run after each prediction + ) + pipeline.start() + pipeline.join() diff --git a/scripts/tracking/run_tracking_studio.py b/scripts/tracking/run_tracking_studio.py new file mode 100755 index 00000000..a32ebf0c --- /dev/null +++ b/scripts/tracking/run_tracking_studio.py @@ -0,0 +1,9 @@ +#!/usr/bin/env python3 +""" +Entry point for the Tracking Studio NiceGUI application. + +Run with: python scripts/run_tracking_studio.py +""" + +# Simply import the app module - ui.run() is called at module level +import collab_env.tracking_studio.app diff --git a/scripts/tracking/yolo_native_tracking.py b/scripts/tracking/yolo_native_tracking.py new file mode 100644 index 00000000..155c1bcd --- /dev/null +++ b/scripts/tracking/yolo_native_tracking.py @@ -0,0 +1,109 @@ +"""Real-time tracking using YOLO's native track() method""" + +import cv2 +import numpy as np +from ultralytics import YOLO +import argparse + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--model_path", type=str, required=True, help="Path to YOLO model") + parser.add_argument("path_to_video", type=str, help="Path to video file") + parser.add_argument("--confidence", type=float, default=0.2, help="Confidence threshold") + parser.add_argument("--tracker", type=str, default="bytetrack.yaml", + help="Tracker config: bytetrack.yaml, botsort.yaml, or path to custom .yaml") + parser.add_argument("--iou", type=float, default=0.5, help="IOU threshold for NMS") + parser.add_argument("--no-persist", action="store_true", help="Don't persist tracks between frames (default: persist=True)") + + args = parser.parse_args() + + # Load model + model = YOLO(args.model_path) + + # Open video + cap = cv2.VideoCapture(args.path_to_video) + + if not cap.isOpened(): + print(f"Error: Could not open video {args.path_to_video}") + return + + # Get video properties + fps = cap.get(cv2.CAP_PROP_FPS) + total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + + print(f"Video: {width}x{height} @ {fps:.2f} fps, {total_frames} frames") + print(f"Tracker: {args.tracker}") + print(f"Confidence: {args.confidence}, IOU: {args.iou}") + print("\nPress 'q' to quit, 'p' to pause/unpause, SPACE to step frame when paused") + + frame_idx = 0 + paused = False + current_frame = None + + # Calculate padding for YOLO + target_height = ((height + 31) // 32) * 32 + target_width = ((width + 31) // 32) * 32 + + while True: + if not paused: + ret, frame = cap.read() + if not ret: + print("\nEnd of video") + break + + # Run YOLO tracking (track() method does detection + tracking in one step!) + results = model.track( + source=frame, + conf=args.confidence, + iou=args.iou, + tracker=args.tracker, + persist=not args.no_persist, # Persist tracks between frames (True by default) + verbose=False, + imgsz=(target_width, target_height), + device='mps' if hasattr(model, 'device') else 'cpu' # Use MPS on Mac if available + ) + + # Get annotated frame with tracking visualization + # YOLO's plot() method draws boxes, masks, and track IDs automatically + annotated_frame = results[0].plot() + + # Add custom info overlay + if results[0].boxes is not None and results[0].boxes.id is not None: + n_detections = len(results[0].boxes.id) + track_ids = results[0].boxes.id.cpu().numpy().astype(int) + unique_tracks = len(np.unique(track_ids)) + info_text = f"Frame: {frame_idx}/{total_frames} | Detections: {n_detections} | Unique IDs: {unique_tracks}" + else: + info_text = f"Frame: {frame_idx}/{total_frames} | Detections: 0" + + cv2.putText(annotated_frame, info_text, (10, 20), + cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0, 255, 0), 2) + + current_frame = annotated_frame + frame_idx += 1 + + # Display frame + if current_frame is not None: + cv2.imshow('YOLO Native Tracking', current_frame) + + # Handle key presses + key = cv2.waitKey(1 if not paused else 0) & 0xFF + + if key == ord('q'): + print("\nQuitting...") + break + elif key == ord('p'): + paused = not paused + print(f"\n{'Paused' if paused else 'Resumed'}") + elif key == ord(' ') and paused: + # Step one frame forward when paused + paused = False + continue + + cap.release() + cv2.destroyAllWindows() + +if __name__ == "__main__": + main() From f41fa4be8f1d7d88c506e489b68c12d8d02822f6 Mon Sep 17 00:00:00 2001 From: Dmitry Batenkov Date: Fri, 13 Feb 2026 18:27:21 -0500 Subject: [PATCH 03/21] deployment plan updated --- docs/tracking/tracking_studio_deploy_plan.md | 92 ++++++++++++++++++++ 1 file changed, 92 insertions(+) create mode 100644 docs/tracking/tracking_studio_deploy_plan.md diff --git a/docs/tracking/tracking_studio_deploy_plan.md b/docs/tracking/tracking_studio_deploy_plan.md new file mode 100644 index 00000000..7de4d243 --- /dev/null +++ b/docs/tracking/tracking_studio_deploy_plan.md @@ -0,0 +1,92 @@ +# Deploy Tracking Studio to Cloud Run + +## Context +The tracking studio NiceGUI app is implemented in `collab_env/tracking_studio/`. The Cloud Run infra (`Dockerfile.tracking-studio`, `cloudbuild.yaml`) exists but has several blockers that prevent a successful deployment. + +## Blockers Found + +### 1. Wrong script path in Dockerfile CMD (CRITICAL) +- **File**: `Dockerfile.tracking-studio:59` +- CMD is `python scripts/run_tracking_studio.py` but the file lives at `scripts/tracking/run_tracking_studio.py` +- Container will crash on startup + +### 2. `config-local/` not available in Cloud Build (CRITICAL) +- **File**: `Dockerfile.tracking-studio:34` +- `COPY config-local/ config/` will fail because `config-local/` is in `.gitignore` +- Cloud Build clones the repo, so gitignored files aren't available + +### 3. GCSClient doesn't support Application Default Credentials (CRITICAL) +- **File**: `collab_env/data/gcs_utils.py:33` +- `GCSClient.__init__()` asserts the credentials file exists and uses `service_account.Credentials.from_service_account_file()` +- On Cloud Run, the recommended auth is ADC via the service account - no credential file needed +- Need to add ADC fallback to `GCSClient` + +### 4. `reload=True` in production +- **File**: `collab_env/tracking_studio/app.py:800` +- NiceGUI's reload mode uses file watchers and a different startup method, which can cause issues in containers +- Should be `False` in production (controlled by env var) + +### 5. Roboflow API key needs Secret Manager +- **File**: `cloudbuild.yaml:33` +- Currently uses `--set-env-vars=ROBOFLOW_API_KEY=${_ROBOFLOW_API_KEY}` (build substitution) +- Key should be stored in GCP Secret Manager for security + +## Plan + +### Step 1: Fix Dockerfile +In `Dockerfile.tracking-studio`: +- Fix CMD path: `scripts/run_tracking_studio.py` -> `scripts/tracking/run_tracking_studio.py` +- Remove `COPY config-local/ config/` (not available in Cloud Build, ADC replaces it) +- Add `ENV NICEGUI_RELOAD=false` + +### Step 2: Add ADC support to GCSClient +In `collab_env/data/gcs_utils.py`, modify `GCSClient.__init__()`: +- If credentials_path is provided and file exists -> use service account file (current behavior) +- If credentials_path is `None` or file doesn't exist -> fall back to ADC: + - `storage.Client(project=project_id)` (ADC auto-detected) + - `gcsfs.GCSFileSystem(project=project_id, token='google_default')` +- Replace `assert os.path.exists()` with conditional logic +- Log which auth method is being used + +### Step 3: Update tracking studio app for production +In `collab_env/tracking_studio/app.py`: +- `get_credentials_path()`: return `None` when env var not set and default path doesn't exist (triggers ADC in GCSClient) +- `ui.run(reload=...)`: use `os.getenv("NICEGUI_RELOAD", "true").lower() == "true"` so Dockerfile can disable it + +### Step 4: Set up Roboflow API key in Secret Manager +In `cloudbuild.yaml`, change the deploy step to use `--set-secrets` instead of `--set-env-vars` for the API key: +```yaml +- '--set-secrets=ROBOFLOW_API_KEY=roboflow-api-key:latest' +``` +This references a secret named `roboflow-api-key` in Secret Manager. + +**Manual prerequisite** (run once before first deploy): +```bash +# Create the secret +echo -n "YOUR_ROBOFLOW_KEY" | gcloud secrets create roboflow-api-key --data-file=- + +# Grant Cloud Run service account access +gcloud secrets add-iam-policy-binding roboflow-api-key \ + --member="serviceAccount:PROJECT_NUMBER-compute@developer.gserviceaccount.com" \ + --role="roles/secretmanager.secretAccessor" +``` + +### Step 5: Deploy +```bash +gcloud builds submit --config=cloudbuild.yaml +``` +No substitution needed since the Roboflow key comes from Secret Manager. + +## Files to Modify +1. `Dockerfile.tracking-studio` - fix CMD path, remove config-local COPY, add NICEGUI_RELOAD=false +2. `collab_env/data/gcs_utils.py` - add ADC fallback in GCSClient.__init__() +3. `collab_env/tracking_studio/app.py` - update get_credentials_path(), make reload configurable +4. `cloudbuild.yaml` - switch ROBOFLOW_API_KEY from substitution to Secret Manager + +## Verification +1. Build Docker image locally: `docker build -f Dockerfile.tracking-studio -t tracking-studio .` +2. Run locally: `docker run -p 8080:8080 tracking-studio` (GCS will be disabled without creds, that's fine) +3. Verify the app loads at http://localhost:8080 with upload + YOLO working +4. Create secret in Secret Manager (manual one-time step) +5. Deploy via `gcloud builds submit --config=cloudbuild.yaml` +6. Verify the Cloud Run URL serves the app and GCS browsing works From 8b006d5b18cbef027e2d263a7eb34d24a97aa4ef Mon Sep 17 00:00:00 2001 From: Dmitry Batenkov Date: Tue, 17 Feb 2026 15:43:02 -0500 Subject: [PATCH 04/21] deployment to cloudrun --- .dockerignore | 59 ++----------- .gcloudignore | 13 +++ Dockerfile.tracking-studio | 39 +++------ cloudbuild.yaml | 10 +-- collab_env/data/gcs_utils.py | 44 ++++++---- collab_env/tracking_studio/app.py | 15 ++-- docs/tracking/tracking_studio_deploy_plan.md | 92 -------------------- scripts/tracking/run_tracking_studio.py | 2 +- 8 files changed, 71 insertions(+), 203 deletions(-) create mode 100644 .gcloudignore delete mode 100644 docs/tracking/tracking_studio_deploy_plan.md diff --git a/.dockerignore b/.dockerignore index 00ed5497..635cecca 100644 --- a/.dockerignore +++ b/.dockerignore @@ -1,54 +1,7 @@ -# Git -.git -.gitignore +# Exclude everything by default +* -# Python -__pycache__/ -*.py[cod] -*$py.class -*.so -.Python -build/ -develop-eggs/ -dist/ -downloads/ -eggs/ -.eggs/ -lib/ -lib64/ -parts/ -sdist/ -var/ -wheels/ -*.egg-info/ -.installed.cfg -*.egg - -# Virtual Environment -venv/ -env/ -ENV/ - -# IDE -.idea/ -.vscode/ -*.swp -*.swo - -# Docker -Dockerfile -.dockerignore - -# Documentation -docs/ -*.md -*.rst - -# Tests -tests/ -.pytest_cache/ -.coverage -htmlcov/ - -# Misc -.DS_Store \ No newline at end of file +# Allow only what the Dockerfile needs +!collab_env/ +!scripts/ +!pyproject.toml diff --git a/.gcloudignore b/.gcloudignore new file mode 100644 index 00000000..06b38a7c --- /dev/null +++ b/.gcloudignore @@ -0,0 +1,13 @@ +# Upload only what the Docker build needs +# .gcloudignore controls what gcloud builds submit uploads + +# Start by ignoring everything +* + +# Allow only what Dockerfile.tracking-studio needs +!collab_env/ +!scripts/ +!pyproject.toml +!Dockerfile.tracking-studio +!cloudbuild.yaml +!.dockerignore diff --git a/Dockerfile.tracking-studio b/Dockerfile.tracking-studio index f3924c13..7916ee7f 100644 --- a/Dockerfile.tracking-studio +++ b/Dockerfile.tracking-studio @@ -10,17 +10,23 @@ RUN apt-get update && apt-get install -y \ libxrender-dev \ && rm -rf /var/lib/apt/lists/* +# Install uv for fast dependency resolution +RUN pip install uv + WORKDIR /workspace -# Copy and install Python dependencies +# Install CPU-only PyTorch first (prevents ultralytics pulling 900MB CUDA version) +RUN uv pip install --system --no-cache \ + torch torchvision --index-url https://download.pytorch.org/whl/cpu + +# Install Python dependencies (Roboflow models load via .pt download, no inference SDK needed) COPY pyproject.toml . -RUN pip install --no-cache-dir \ +RUN uv pip install --system --no-cache \ nicegui \ ultralytics \ supervision \ google-cloud-storage \ gcsfs \ - inference \ opencv-python-headless \ pandas \ numpy \ @@ -30,30 +36,13 @@ RUN pip install --no-cache-dir \ COPY collab_env/ collab_env/ COPY scripts/ scripts/ -# Copy GCS credentials -COPY config-local/ config/ - -# Pre-download YOLO models (YOLO11 and YOLO26 variants) -RUN mkdir -p /workspace/models && \ - python -c "from ultralytics import YOLO; \ - print('Downloading YOLO11 models...'); \ - YOLO('yolo11n.pt'); \ - YOLO('yolo11s.pt'); \ - YOLO('yolo11m.pt');" && \ - mv ~/.cache/ultralytics/*.pt /workspace/models/ 2>/dev/null || true - -# Note: YOLO26 models might not be available yet, will download on first use -# RUN python -c "from ultralytics import YOLO; \ -# print('Downloading YOLO26 models...'); \ -# YOLO('yolo26n-fast.pt'); \ -# YOLO('yolo26s-fast.pt'); \ -# YOLO('yolo26m-fast.pt');" && \ -# mv ~/.cache/ultralytics/*.pt /workspace/models/ 2>/dev/null || true - -# Create tmp directories +# Create model cache and tmp directories +RUN mkdir -p /workspace/models RUN mkdir -p /tmp/videos /tmp/outputs /tmp/uploads ENV PORT=8080 ENV PYTHONUNBUFFERED=1 +ENV PYTHONPATH=/workspace +ENV NICEGUI_RELOAD=false -CMD ["python", "scripts/run_tracking_studio.py"] +CMD ["python", "scripts/tracking/run_tracking_studio.py"] diff --git a/cloudbuild.yaml b/cloudbuild.yaml index f3865f58..ac526f6d 100644 --- a/cloudbuild.yaml +++ b/cloudbuild.yaml @@ -3,16 +3,11 @@ steps: args: - 'build' - '-t' - - 'gcr.io/$PROJECT_ID/tracking-studio:$SHORT_SHA' - - '-t' - 'gcr.io/$PROJECT_ID/tracking-studio:latest' - '-f' - 'Dockerfile.tracking-studio' - '.' - - name: 'gcr.io/cloud-builders/docker' - args: ['push', 'gcr.io/$PROJECT_ID/tracking-studio:$SHORT_SHA'] - - name: 'gcr.io/cloud-builders/docker' args: ['push', 'gcr.io/$PROJECT_ID/tracking-studio:latest'] @@ -22,7 +17,7 @@ steps: - 'run' - 'deploy' - 'tracking-studio' - - '--image=gcr.io/$PROJECT_ID/tracking-studio:$SHORT_SHA' + - '--image=gcr.io/$PROJECT_ID/tracking-studio:latest' - '--platform=managed' - '--region=us-central1' - '--memory=4Gi' @@ -30,11 +25,10 @@ steps: - '--timeout=900' - '--concurrency=1' - '--max-instances=5' - - '--set-env-vars=ROBOFLOW_API_KEY=${_ROBOFLOW_API_KEY}' + - '--set-secrets=ROBOFLOW_API_KEY=roboflow-api-key:latest' - '--allow-unauthenticated' images: - - 'gcr.io/$PROJECT_ID/tracking-studio:$SHORT_SHA' - 'gcr.io/$PROJECT_ID/tracking-studio:latest' options: diff --git a/collab_env/data/gcs_utils.py b/collab_env/data/gcs_utils.py index 0c8bcf17..aa204783 100644 --- a/collab_env/data/gcs_utils.py +++ b/collab_env/data/gcs_utils.py @@ -22,32 +22,40 @@ def __init__( ): """ Args: - credentials_path: Path to GCS credentials file. If not provided, will use the default path. + credentials_path: Path to GCS credentials file. If not provided, will try + the default path, then fall back to Application Default Credentials. """ if credentials_path is None: - credentials_path = expand_path( + default_path = expand_path( DEFAULT_GCS_CREDENTIALS_PATH, get_project_root() ) - - self.credentials_path = credentials_path - assert os.path.exists(self.credentials_path), ( - f"Credentials file {self.credentials_path} does not exist" - ) - logger.info(f"Using credentials from {self.credentials_path}") - self.credentials = service_account.Credentials.from_service_account_file( - self.credentials_path - ) + if os.path.exists(default_path): + credentials_path = default_path self.project_id = project_id - logger.info(f"Using project {self.project_id}") - self._gcs = gcsfs.GCSFileSystem( - self.project_id, token=str(self.credentials_path) - ) - self._storage_client = storage.Client( - self.project_id, credentials=self.credentials - ) + if credentials_path and os.path.exists(str(credentials_path)): + logger.info(f"Using service account credentials from {credentials_path}") + self.credentials_path = credentials_path + self.credentials = service_account.Credentials.from_service_account_file( + str(credentials_path) + ) + self._gcs = gcsfs.GCSFileSystem( + self.project_id, token=str(credentials_path) + ) + self._storage_client = storage.Client( + self.project_id, credentials=self.credentials + ) + else: + logger.info("Using Application Default Credentials (no credentials file found)") + self.credentials_path = None + self.credentials = None + self._gcs = gcsfs.GCSFileSystem( + project=self.project_id, token="google_default" + ) + self._storage_client = storage.Client(project=self.project_id) + logger.info(f"Using project {self.project_id}") self.is_initialized = True @property diff --git a/collab_env/tracking_studio/app.py b/collab_env/tracking_studio/app.py index c1f59956..e669c37b 100644 --- a/collab_env/tracking_studio/app.py +++ b/collab_env/tracking_studio/app.py @@ -35,11 +35,14 @@ def load_bytetrack_params(): # Initialize services def get_credentials_path(): - """Get GCS credentials path from environment or default""" - return os.getenv( - "GCS_CREDENTIALS", - "/workspace/config/collab-data-463313-c340ad86b28e.json", - ) + """Get GCS credentials path from environment or default. Returns None for ADC.""" + env_path = os.getenv("GCS_CREDENTIALS") + if env_path: + return env_path + default = "/workspace/config/collab-data-463313-c340ad86b28e.json" + if os.path.exists(default): + return default + return None # GCSClient will use Application Default Credentials try: @@ -797,6 +800,6 @@ async def frame_callback(annotated_frame, frame_idx, total_frames): ui.run( host="0.0.0.0", port=int(os.getenv("PORT", 8080)), - reload=True, # Enable auto-reload for development + reload=os.getenv("NICEGUI_RELOAD", "true").lower() == "true", title="Tracking Studio", ) diff --git a/docs/tracking/tracking_studio_deploy_plan.md b/docs/tracking/tracking_studio_deploy_plan.md deleted file mode 100644 index 7de4d243..00000000 --- a/docs/tracking/tracking_studio_deploy_plan.md +++ /dev/null @@ -1,92 +0,0 @@ -# Deploy Tracking Studio to Cloud Run - -## Context -The tracking studio NiceGUI app is implemented in `collab_env/tracking_studio/`. The Cloud Run infra (`Dockerfile.tracking-studio`, `cloudbuild.yaml`) exists but has several blockers that prevent a successful deployment. - -## Blockers Found - -### 1. Wrong script path in Dockerfile CMD (CRITICAL) -- **File**: `Dockerfile.tracking-studio:59` -- CMD is `python scripts/run_tracking_studio.py` but the file lives at `scripts/tracking/run_tracking_studio.py` -- Container will crash on startup - -### 2. `config-local/` not available in Cloud Build (CRITICAL) -- **File**: `Dockerfile.tracking-studio:34` -- `COPY config-local/ config/` will fail because `config-local/` is in `.gitignore` -- Cloud Build clones the repo, so gitignored files aren't available - -### 3. GCSClient doesn't support Application Default Credentials (CRITICAL) -- **File**: `collab_env/data/gcs_utils.py:33` -- `GCSClient.__init__()` asserts the credentials file exists and uses `service_account.Credentials.from_service_account_file()` -- On Cloud Run, the recommended auth is ADC via the service account - no credential file needed -- Need to add ADC fallback to `GCSClient` - -### 4. `reload=True` in production -- **File**: `collab_env/tracking_studio/app.py:800` -- NiceGUI's reload mode uses file watchers and a different startup method, which can cause issues in containers -- Should be `False` in production (controlled by env var) - -### 5. Roboflow API key needs Secret Manager -- **File**: `cloudbuild.yaml:33` -- Currently uses `--set-env-vars=ROBOFLOW_API_KEY=${_ROBOFLOW_API_KEY}` (build substitution) -- Key should be stored in GCP Secret Manager for security - -## Plan - -### Step 1: Fix Dockerfile -In `Dockerfile.tracking-studio`: -- Fix CMD path: `scripts/run_tracking_studio.py` -> `scripts/tracking/run_tracking_studio.py` -- Remove `COPY config-local/ config/` (not available in Cloud Build, ADC replaces it) -- Add `ENV NICEGUI_RELOAD=false` - -### Step 2: Add ADC support to GCSClient -In `collab_env/data/gcs_utils.py`, modify `GCSClient.__init__()`: -- If credentials_path is provided and file exists -> use service account file (current behavior) -- If credentials_path is `None` or file doesn't exist -> fall back to ADC: - - `storage.Client(project=project_id)` (ADC auto-detected) - - `gcsfs.GCSFileSystem(project=project_id, token='google_default')` -- Replace `assert os.path.exists()` with conditional logic -- Log which auth method is being used - -### Step 3: Update tracking studio app for production -In `collab_env/tracking_studio/app.py`: -- `get_credentials_path()`: return `None` when env var not set and default path doesn't exist (triggers ADC in GCSClient) -- `ui.run(reload=...)`: use `os.getenv("NICEGUI_RELOAD", "true").lower() == "true"` so Dockerfile can disable it - -### Step 4: Set up Roboflow API key in Secret Manager -In `cloudbuild.yaml`, change the deploy step to use `--set-secrets` instead of `--set-env-vars` for the API key: -```yaml -- '--set-secrets=ROBOFLOW_API_KEY=roboflow-api-key:latest' -``` -This references a secret named `roboflow-api-key` in Secret Manager. - -**Manual prerequisite** (run once before first deploy): -```bash -# Create the secret -echo -n "YOUR_ROBOFLOW_KEY" | gcloud secrets create roboflow-api-key --data-file=- - -# Grant Cloud Run service account access -gcloud secrets add-iam-policy-binding roboflow-api-key \ - --member="serviceAccount:PROJECT_NUMBER-compute@developer.gserviceaccount.com" \ - --role="roles/secretmanager.secretAccessor" -``` - -### Step 5: Deploy -```bash -gcloud builds submit --config=cloudbuild.yaml -``` -No substitution needed since the Roboflow key comes from Secret Manager. - -## Files to Modify -1. `Dockerfile.tracking-studio` - fix CMD path, remove config-local COPY, add NICEGUI_RELOAD=false -2. `collab_env/data/gcs_utils.py` - add ADC fallback in GCSClient.__init__() -3. `collab_env/tracking_studio/app.py` - update get_credentials_path(), make reload configurable -4. `cloudbuild.yaml` - switch ROBOFLOW_API_KEY from substitution to Secret Manager - -## Verification -1. Build Docker image locally: `docker build -f Dockerfile.tracking-studio -t tracking-studio .` -2. Run locally: `docker run -p 8080:8080 tracking-studio` (GCS will be disabled without creds, that's fine) -3. Verify the app loads at http://localhost:8080 with upload + YOLO working -4. Create secret in Secret Manager (manual one-time step) -5. Deploy via `gcloud builds submit --config=cloudbuild.yaml` -6. Verify the Cloud Run URL serves the app and GCS browsing works diff --git a/scripts/tracking/run_tracking_studio.py b/scripts/tracking/run_tracking_studio.py index a32ebf0c..e26a88d3 100755 --- a/scripts/tracking/run_tracking_studio.py +++ b/scripts/tracking/run_tracking_studio.py @@ -2,7 +2,7 @@ """ Entry point for the Tracking Studio NiceGUI application. -Run with: python scripts/run_tracking_studio.py +Run with: python scripts/tracking/run_tracking_studio.py """ # Simply import the app module - ui.run() is called at module level From bc348c94ff01e24117046e06f425321be01197b8 Mon Sep 17 00:00:00 2001 From: Dmitry Batenkov Date: Tue, 17 Feb 2026 16:07:39 -0500 Subject: [PATCH 05/21] added "detection only" option --- collab_env/tracking_studio/app.py | 33 +++-- collab_env/tracking_studio/video_processor.py | 113 +++++++++++------- 2 files changed, 97 insertions(+), 49 deletions(-) diff --git a/collab_env/tracking_studio/app.py b/collab_env/tracking_studio/app.py index e669c37b..0e58ed64 100644 --- a/collab_env/tracking_studio/app.py +++ b/collab_env/tracking_studio/app.py @@ -378,6 +378,13 @@ def enable_load_model_btn(e=None): param_select.tooltip(param_config["description"]) param_widgets[param_name] = param_select + # Detection-only mode toggle + detection_only_checkbox = ui.checkbox( + "Detection only (no tracking)", + value=False, + ).classes("text-xs mt-2") + detection_only_checkbox.tooltip("Run detection without ByteTrack — shows raw detections per frame") + # Skip frames (for fast-forward, not a ByteTrack param) with ui.row().classes("w-full items-center gap-2 mt-2"): skip_frames_label = ui.label("Skip Frames: 1 frame").classes("text-xs") @@ -691,9 +698,12 @@ async def start_tracking(): results_container.classes(add="hidden") try: - # Show tracker type in progress label - tracker_type = state.get("tracker_type", "Unknown") - progress_label.text = f"Starting tracking ({tracker_type})..." + # Show mode in progress label + if detection_only_checkbox.value: + progress_label.text = "Starting detection..." + else: + tracker_type = state.get("tracker_type", "Unknown") + progress_label.text = f"Starting tracking ({tracker_type})..." progress.value = 0 # Use already-loaded video and model from state @@ -741,6 +751,7 @@ async def frame_callback(annotated_frame, frame_idx, total_frames): model=model, tracker_config=tracker_config, confidence=conf_slider.value, + detection_only=detection_only_checkbox.value, frame_callback=frame_callback, stop_event=state["stop_event"], pause_event=state["pause_event"], @@ -756,11 +767,17 @@ async def frame_callback(annotated_frame, frame_idx, total_frames): state["results"] = results - stats_label.text = ( - f"Processed {results['stats']['total_frames']} frames | " - f"{results['stats']['total_detections']} detections | " - f"{results['stats']['unique_tracks']} unique tracks" - ) + if detection_only_checkbox.value: + stats_label.text = ( + f"Processed {results['stats']['total_frames']} frames | " + f"{results['stats']['total_detections']} detections" + ) + else: + stats_label.text = ( + f"Processed {results['stats']['total_frames']} frames | " + f"{results['stats']['total_detections']} detections | " + f"{results['stats']['unique_tracks']} unique tracks" + ) # Setup download button (only tracking CSV) download_track_btn.on_click(lambda: ui.download(results["tracking_csv"])) diff --git a/collab_env/tracking_studio/video_processor.py b/collab_env/tracking_studio/video_processor.py index e04746da..6b1d8ea3 100644 --- a/collab_env/tracking_studio/video_processor.py +++ b/collab_env/tracking_studio/video_processor.py @@ -26,6 +26,7 @@ def __init__( model: Union[YOLO, Any], # YOLO or Roboflow model tracker_config: Dict, # ByteTrack parameters confidence: float = 0.5, + detection_only: bool = False, frame_callback: Callable[[np.ndarray, int, int], None] = None, stop_event: threading.Event = None, pause_event: threading.Event = None, @@ -38,6 +39,7 @@ def __init__( model: Detection model (YOLO or Roboflow) tracker_config: Tracker configuration dict confidence: Detection confidence threshold + detection_only: If True, run detection without tracking (no track IDs) frame_callback: Async callback for frame updates (frame, frame_idx, total_frames) stop_event: Threading event to signal hard stop pause_event: Threading event to signal pause/resume @@ -45,6 +47,7 @@ def __init__( """ self.model = model self.confidence = confidence + self.detection_only = detection_only self.frame_callback = frame_callback self.stop_event = stop_event or threading.Event() self.pause_event = pause_event or threading.Event() @@ -56,8 +59,12 @@ def __init__( # Check if model supports native tracking self.use_native_tracking = isinstance(model, YOLO) - # For Roboflow inference models (fallback), initialize supervision tracker - if not self.use_native_tracking: + if detection_only: + logger.info("Detection-only mode (no tracking)") + self.tracker = None + self.tracker_yaml_path = None + elif not self.use_native_tracking: + # For Roboflow inference models (fallback), initialize supervision tracker logger.info("Using supervision ByteTrack (Roboflow inference model fallback)") self.tracker = sv.ByteTrack( track_activation_threshold=tracker_config.get("track_high_thresh", 0.25), @@ -188,9 +195,23 @@ def _process_video_sync( frame_idx += 1 continue - # 1. Run detection and tracking + # 1. Run detection (and optionally tracking) try: - if self.use_native_tracking: + if self.detection_only: + # Detection only - no tracking + if self.use_native_tracking: + results = self.model( + source=frame, + conf=self.confidence, + verbose=False, + )[0] + detections = sv.Detections.from_ultralytics(results) + else: + results = self.model.infer(frame, confidence=self.confidence)[0] + detections = sv.Detections.from_inference(results) + tracked_detections = detections + + elif self.use_native_tracking: # Use Ultralytics native tracking (supports all ByteTrack parameters) results = self.model.track( source=frame, @@ -220,7 +241,7 @@ def _process_video_sync( detections = sv.Detections.empty() tracked_detections = sv.Detections.empty() - # 2. Save raw detections (for stats only, not exported) + # 2. Save detections for i, (bbox, conf, class_id) in enumerate( zip(detections.xyxy, detections.confidence, detections.class_id) ): @@ -236,11 +257,8 @@ def _process_video_sync( } ) - # 3. Tracked detections now have track IDs (from native tracking or supervision) - - # 4. Save tracking with IDs (matches output_tracked_bboxes_csv format) - # Only save if we have track IDs (handles cases where no detections exist) - if tracked_detections.tracker_id is not None and len(tracked_detections) > 0: + # 3. Save tracking data (with track IDs if tracking is enabled) + if not self.detection_only and tracked_detections.tracker_id is not None and len(tracked_detections) > 0: for bbox, track_id, conf, class_id in zip( tracked_detections.xyxy, tracked_detections.tracker_id, @@ -260,19 +278,26 @@ def _process_video_sync( } ) - # 5. Annotate frame for display + # 4. Annotate frame for display annotated_frame = frame.copy() annotated_frame = self.box_annotator.annotate( annotated_frame, tracked_detections ) - # Create labels with track IDs - labels = [ - f"#{track_id} {conf:.2f}" - for track_id, conf in zip( - tracked_detections.tracker_id, tracked_detections.confidence - ) - ] + if self.detection_only: + # Labels with confidence only + labels = [ + f"{conf:.2f}" + for conf in tracked_detections.confidence + ] + else: + # Labels with track IDs + labels = [ + f"#{track_id} {conf:.2f}" + for track_id, conf in zip( + tracked_detections.tracker_id, tracked_detections.confidence + ) + ] annotated_frame = self.label_annotator.annotate( annotated_frame, tracked_detections, labels=labels ) @@ -304,43 +329,49 @@ def _process_video_sync( except Exception as e: logger.warning(f"Failed to cleanup tracker config: {e}") + unique_tracks = len(set(t['track_id'] for t in tracking_list)) if tracking_list else 0 logger.info( f"Processing complete: {total_frames} frames, " f"{len(detections_list)} detections, " - f"{len(set(t['track_id'] for t in tracking_list))} unique tracks" + f"{unique_tracks} unique tracks" ) - # 7. Save tracking CSV (matches output_tracked_bboxes_csv format) - tracking_df = pd.DataFrame(tracking_list) - output_path = Path(output_dir) output_path.mkdir(parents=True, exist_ok=True) - # Only save tracking CSV (not detections - user doesn't need them) - tracking_csv = output_path / "tracking.csv" - - if len(tracking_list) > 0: - # Ensure column order matches: track_id,frame,x1,y1,x2,y2,confidence,class - tracking_df = tracking_df[ - ["track_id", "frame", "x1", "y1", "x2", "y2", "confidence", "class"] - ] - tracking_df.to_csv(tracking_csv, index=False) - logger.info(f"Saved tracking CSV to {tracking_csv}") + if self.detection_only: + # Save detections CSV (no track IDs) + output_csv = output_path / "detections.csv" + if len(detections_list) > 0: + det_df = pd.DataFrame(detections_list) + det_df = det_df[["frame", "x1", "y1", "x2", "y2", "confidence", "class"]] + det_df.to_csv(output_csv, index=False) + else: + pd.DataFrame( + columns=["frame", "x1", "y1", "x2", "y2", "confidence", "class"] + ).to_csv(output_csv, index=False) + logger.info(f"Saved detections CSV to {output_csv}") else: - # Create empty CSV with correct headers - pd.DataFrame( - columns=["track_id", "frame", "x1", "y1", "x2", "y2", "confidence", "class"] - ).to_csv(tracking_csv, index=False) - logger.warning("No tracks found, saved empty CSV") + # Save tracking CSV (with track IDs) + output_csv = output_path / "tracking.csv" + if len(tracking_list) > 0: + tracking_df = pd.DataFrame(tracking_list) + tracking_df = tracking_df[ + ["track_id", "frame", "x1", "y1", "x2", "y2", "confidence", "class"] + ] + tracking_df.to_csv(output_csv, index=False) + else: + pd.DataFrame( + columns=["track_id", "frame", "x1", "y1", "x2", "y2", "confidence", "class"] + ).to_csv(output_csv, index=False) + logger.info(f"Saved tracking CSV to {output_csv}") return { - "tracking_csv": str(tracking_csv), + "tracking_csv": str(output_csv), "stats": { "total_frames": total_frames, "total_detections": len(detections_list), - "unique_tracks": ( - tracking_df["track_id"].nunique() if len(tracking_list) > 0 else 0 - ), + "unique_tracks": unique_tracks, "fps": fps, }, } From da16c50f3a3e9d29f04c157ac5e46ebf88c54a58 Mon Sep 17 00:00:00 2001 From: Dmitry Batenkov Date: Tue, 17 Feb 2026 16:40:11 -0500 Subject: [PATCH 06/21] small gui fixes --- .gitignore | 3 ++ collab_env/tracking_studio/app.py | 22 +++++++++++++-- docs/tracking/tracking_web_gui.md | 47 +++++++++++++++++++++++-------- 3 files changed, 59 insertions(+), 13 deletions(-) diff --git a/.gitignore b/.gitignore index ac39baee..e09a4a89 100644 --- a/.gitignore +++ b/.gitignore @@ -87,6 +87,9 @@ venv/ # VSCode /.vscode +# Claude Code +.claude/ + # Automatically generated files docs/source/lightning_logs docs/preconvert diff --git a/collab_env/tracking_studio/app.py b/collab_env/tracking_studio/app.py index 0e58ed64..44f6782c 100644 --- a/collab_env/tracking_studio/app.py +++ b/collab_env/tracking_studio/app.py @@ -295,7 +295,8 @@ def enable_load_model_btn(e=None): rf_version_select.on("update:model-value", enable_load_model_btn) # Parameters card - with ui.card().classes("w-full shadow-md p-3"): + params_card = ui.card().classes("w-full shadow-md p-3") + with params_card: ui.label("⚙️ Parameters").classes("text-sm font-semibold mb-2") # Detection confidence (not in ByteTrack params) @@ -482,7 +483,11 @@ def seek_to_frame(e): # Preview card with ui.card().classes("w-full shadow-md p-3"): ui.label("Live Preview").classes("text-sm font-semibold mb-2") - video_display = ui.interactive_image().classes("w-full border-2 border-gray-200 rounded bg-gray-50").style("max-height: 500px; object-fit: contain;") + video_container = ui.element("div").classes("border-2 border-gray-200 rounded bg-gray-50").style( + "max-width: 100%; resize: horizontal; overflow: hidden;" + ) + with video_container: + video_display = ui.interactive_image().style("width: 100%; height: 100%; object-fit: contain;") # Results (initially hidden, separate row) results_container = ui.card().classes("w-full shadow-md p-3 hidden") @@ -551,6 +556,13 @@ async def load_video(): cap.set(cv2.CAP_PROP_POS_FRAMES, 0) ret, frame = cap.read() if ret: + # Size container to video's native dimensions, lock aspect ratio for resize + h, w = frame.shape[:2] + video_container.style( + f"width: {w}px; max-width: 100%; aspect-ratio: {w}/{h};" + f" resize: horizontal; overflow: hidden;" + ) + _, buffer = cv2.imencode(".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, 85]) img_base64 = base64.b64encode(buffer).decode('utf-8') video_display.set_source(f"data:image/jpeg;base64,{img_base64}") @@ -693,8 +705,11 @@ async def start_tracking(): state["skip_frames_event"] = {"skip_amount": 0} # Skip forward state["current_frame"] = 0 start_btn.disable() + pause_btn.text = "Pause" + pause_btn.props("icon=pause") pause_btn.enable() stop_btn.enable() + params_card.style("opacity: 0.5; pointer-events: none;") results_container.classes(add="hidden") try: @@ -802,8 +817,11 @@ async def frame_callback(annotated_frame, frame_idx, total_frames): state["pause_event"] = None state["skip_frames_event"] = None start_btn.enable() + pause_btn.text = "Pause" + pause_btn.props("icon=pause") pause_btn.disable() stop_btn.disable() + params_card.style(remove="opacity: 0.5; pointer-events: none;") # Wire up buttons to event handlers (after functions are defined) load_video_btn.on_click(load_video) diff --git a/docs/tracking/tracking_web_gui.md b/docs/tracking/tracking_web_gui.md index bb885994..2ae7d669 100644 --- a/docs/tracking/tracking_web_gui.md +++ b/docs/tracking/tracking_web_gui.md @@ -13,29 +13,54 @@ The Tracking Studio provides a user-friendly interface for: ## Quick Start +### Prerequisites + +1. **Python 3.10** with the project installed: + + ```bash + pip install -e . + ``` + +2. **FFmpeg** (for video format conversion): + + ```bash + # macOS + brew install ffmpeg + # Ubuntu/Debian + sudo apt install ffmpeg + ``` + +3. **GCS credentials** (for browsing videos in Google Cloud Storage): + - Place your service account JSON at `config-local/collab-data-463313-c340ad86b28e.json` + - Or set the env var: `export GCS_CREDENTIALS=/path/to/credentials.json` + - If neither is set, GCS browsing is disabled (video upload still works) + +4. **Roboflow API key** (only needed for Roboflow models): + + ```bash + export ROBOFLOW_API_KEY=your_api_key_here + ``` + + Get your key from [Roboflow settings](https://app.roboflow.com/settings/api) + ### Running the Application ```bash # From the repository root -python scripts/run_tracking_studio.py - -# Or directly -python -m collab_env.tracking_studio.app +python scripts/tracking/run_tracking_studio.py ``` The application will start on `http://localhost:8080` -### Environment Setup +### Cloud Run Deployment -**Optional: For Google Cloud Storage integration** +The tracking studio is deployed to Cloud Run via `cloudbuild.yaml`: ```bash -export GCS_CREDENTIALS=/path/to/credentials.json +gcloud builds submit --config=cloudbuild.yaml ``` -**Optional: For Roboflow models** -```bash -export ROBOFLOW_API_KEY=your_api_key_here -``` +The Roboflow API key is stored in GCP Secret Manager (`roboflow-api-key`). +GCS access uses the Cloud Run service account (Application Default Credentials). ## Workflow From b402ef21644d1e8713a501ec262ae03b2c6712b0 Mon Sep 17 00:00:00 2001 From: Dmitry Batenkov Date: Tue, 17 Feb 2026 16:46:54 -0500 Subject: [PATCH 07/21] update docs --- README.rst | 1 + docs/tracking/tracking_web_gui.md | 25 ------------------------- 2 files changed, 1 insertion(+), 25 deletions(-) diff --git a/README.rst b/README.rst index 5afcbc9a..c80e9612 100644 --- a/README.rst +++ b/README.rst @@ -141,5 +141,6 @@ Detailed documentation for specific modules: * `GNN Training `_ - Graph Neural Network training and rollouts * `Simulation `_ - Boids simulation and output format * `Tracking `_ - Animal tracking and thermal video processing +* `Tracking Studio `_ - Interactive web GUI for video object detection and tracking For contributing guidelines, see `CONTRIBUTING.md `_. diff --git a/docs/tracking/tracking_web_gui.md b/docs/tracking/tracking_web_gui.md index 2ae7d669..c4386a02 100644 --- a/docs/tracking/tracking_web_gui.md +++ b/docs/tracking/tracking_web_gui.md @@ -52,16 +52,6 @@ python scripts/tracking/run_tracking_studio.py The application will start on `http://localhost:8080` -### Cloud Run Deployment - -The tracking studio is deployed to Cloud Run via `cloudbuild.yaml`: -```bash -gcloud builds submit --config=cloudbuild.yaml -``` - -The Roboflow API key is stored in GCP Secret Manager (`roboflow-api-key`). -GCS access uses the Cloud Run service account (Application Default Credentials). - ## Workflow ### 1. Load Video @@ -318,21 +308,6 @@ Videos not in H.264 format are automatically converted on load: - Roboflow: Verify `ROBOFLOW_API_KEY` is set and has access - Custom: Ensure `.pt` file is YOLO-compatible format -## Advanced Usage - -### Running with Docker - -```bash -# Build image -docker build -t tracking-studio . - -# Run with GCS credentials -docker run -p 8080:8080 \ - -v /path/to/credentials.json:/workspace/config/credentials.json \ - -e GCS_CREDENTIALS=/workspace/config/credentials.json \ - tracking-studio -``` - ### Batch Processing For offline batch processing without the GUI, use the [full tracking pipeline notebook](full_pipeline.ipynb) or direct API: From 0a764c40abf22de320d4da53c164f5527763f9d7 Mon Sep 17 00:00:00 2001 From: Dmitry Batenkov Date: Tue, 17 Feb 2026 18:47:47 -0500 Subject: [PATCH 08/21] perf. updates --- collab_env/tracking_studio/app.py | 237 ++++++++++++------ collab_env/tracking_studio/video_converter.py | 57 +++-- collab_env/tracking_studio/video_processor.py | 55 ++-- 3 files changed, 206 insertions(+), 143 deletions(-) diff --git a/collab_env/tracking_studio/app.py b/collab_env/tracking_studio/app.py index 44f6782c..635aea24 100644 --- a/collab_env/tracking_studio/app.py +++ b/collab_env/tracking_studio/app.py @@ -14,6 +14,7 @@ import uuid import os import io +import base64 from loguru import logger from .gcs_browser import GCSVideoBrowser @@ -395,9 +396,9 @@ def enable_load_model_btn(e=None): # GUI refresh rate (display updates, not a ByteTrack param) with ui.row().classes("w-full items-center gap-2 mt-1"): - display_update_label = ui.label("Display Update: 10 frames").classes("text-xs") - display_update_slider = ui.slider(min=1, max=30, step=1, value=10).style("width: 100px") - display_update_slider.tooltip("Update display every Nth frame (lower = smoother, more network traffic)") + display_update_label = ui.label("Display Update: every frame").classes("text-xs") + display_update_slider = ui.slider(min=1, max=30, step=1, value=1).style("width: 100px") + display_update_slider.tooltip("Update display every Nth frame (1 = smoothest, higher = skip display frames)") display_update_slider.on("update:model-value", lambda e: display_update_label.set_text(f"Display Update: {int(e.args)} {'frame' if int(e.args) == 1 else 'frames'}")) # RIGHT: Controls + Preview (stacked vertically) @@ -430,50 +431,71 @@ def enable_load_model_btn(e=None): ui.separator().props("vertical") - with ui.column().classes("flex-grow gap-1"): - progress_label = ui.label("Ready").classes("text-xs") - progress = ui.linear_progress(value=0).props("size=15px color=primary") + status_indicator = ui.label("Ready").classes("text-xs flex-grow") # Time slider for seeking with ui.column().classes("w-full gap-1 mt-2"): time_label = ui.label("Frame: 0 / 0").classes("text-xs text-gray-600") - time_slider = ui.slider(min=0, max=100, value=0).props("lazy").classes("w-full") + time_slider = ui.slider(min=0, max=100, value=0).classes("w-full") time_slider.disable() + _preview_pending = [False] + async def preview_frame_on_drag(e): - """Show video frame preview during drag (no detection/tracking)""" - if not state.get("video_path"): + """Preview frame during slider drag — reads frame via cv2""" + if _preview_pending[0] or not state.get("video_path"): return - - target_frame = int(e.args) + _preview_pending[0] = True try: - # Open video for preview (separate from processing thread) import cv2 - import base64 - cap = cv2.VideoCapture(str(state["video_path"])) - cap.set(cv2.CAP_PROP_POS_FRAMES, target_frame) - ret, frame = cap.read() - cap.release() - - if ret: - # Show raw frame without annotations - _, buffer = cv2.imencode(".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, 85]) - img_base64 = base64.b64encode(buffer).decode('utf-8') - video_display.set_source(f"data:image/jpeg;base64,{img_base64}") - total_frames = state.get("total_frames", target_frame) - time_label.text = f"Frame: {target_frame} / {total_frames}" - except Exception as err: - logger.warning(f"Preview failed: {err}") - - def seek_to_frame(e): + target_frame = int(e.args) + video_display.content = '' # Clear SVG overlay + if not state.get("processing"): + reset_tracker_state() + + def read_frame(path, idx): + cap = cv2.VideoCapture(str(path)) + cap.set(cv2.CAP_PROP_POS_FRAMES, idx) + ret, f = cap.read() + cap.release() + return f if ret else None + + frame = await asyncio.to_thread( + read_frame, state["video_path"], target_frame + ) + if frame is not None: + _, buf = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, 50]) + b64 = base64.b64encode(buf).decode() + video_display.set_source(f'data:image/jpeg;base64,{b64}') + + total_frames = state.get("total_frames", target_frame) + time_label.text = f"Frame: {target_frame} / {total_frames}" + finally: + _preview_pending[0] = False + + async def seek_to_frame(e): """Seek to specific frame when slider is released""" + target_frame = int(e.args) if state.get("processing") and state.get("skip_frames_event"): - target_frame = int(e.args) current_frame = state.get("current_frame", 0) if target_frame != current_frame: - # Clear any pending seeks state["skip_frames_event"]["skip_amount"] = target_frame - current_frame ui.notify(f"Seeking to frame {target_frame}...", type="info") + elif state.get("video_path"): + # Not processing: show the frame at release position + def read_frame(path, idx): + import cv2 + cap = cv2.VideoCapture(str(path)) + cap.set(cv2.CAP_PROP_POS_FRAMES, idx) + ret, f = cap.read() + cap.release() + return f if ret else None + frame = await asyncio.to_thread(read_frame, state["video_path"], target_frame) + if frame is not None: + import cv2 + _, buf = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, 70]) + b64 = base64.b64encode(buf).decode() + video_display.set_source(f'data:image/jpeg;base64,{b64}') # Live preview during drag time_slider.on("update:model-value", preview_frame_on_drag) @@ -487,7 +509,9 @@ def seek_to_frame(e): "max-width: 100%; resize: horizontal; overflow: hidden;" ) with video_container: - video_display = ui.interactive_image().style("width: 100%; height: 100%; object-fit: contain;") + video_display = ui.interactive_image('').style( + "width: 100%;" + ) # Results (initially hidden, separate row) results_container = ui.card().classes("w-full shadow-md p-3 hidden") @@ -498,6 +522,13 @@ def seek_to_frame(e): download_track_btn = ui.button("Download CSV").props("color=primary icon=download size=sm") # Event handlers + def reset_tracker_state(): + """Reset YOLO model's internal tracker so track IDs start fresh.""" + model = state.get("loaded_model") + if model and hasattr(model, 'predictor') and model.predictor is not None: + # Full predictor reset — Ultralytics will create a fresh one on next call + model.predictor = None + async def load_video(): """Load and prepare video for viewing/tracking""" from nicegui import context @@ -530,11 +561,28 @@ async def load_video(): else: raise ValueError("No video selected. Please select or upload a video.") - # Convert to H.264 if needed + # Ensure browser-compatible H.264 MP4 if await asyncio.to_thread(needs_conversion, local_video): - status_label.text = "Converting to H.264..." converted_video = local_video.parent / f"{local_video.stem}_h264.mp4" - await asyncio.to_thread(convert_to_h264, local_video, converted_video) + # Check if codec is already h264 (just needs container remux) + import subprocess + try: + probe = subprocess.run( + ["ffprobe", "-v", "error", "-select_streams", "v:0", + "-show_entries", "stream=codec_name", + "-of", "default=noprint_wrappers=1:nokey=1", str(local_video)], + capture_output=True, text=True, check=True + ) + is_h264 = probe.stdout.strip() == "h264" + except Exception: + is_h264 = False + + if is_h264: + status_label.text = "Remuxing to MP4..." + await asyncio.to_thread(convert_to_h264, local_video, converted_video, remux_only=True) + else: + status_label.text = "Converting to H.264..." + await asyncio.to_thread(convert_to_h264, local_video, converted_video) local_video = converted_video with client: ui.notify("Video converted to H.264") @@ -545,29 +593,35 @@ async def load_video(): state["video_path"] = local_video state["video_loaded"] = True - # Set up video display and time slider + # Read video metadata and first frame import cv2 - import base64 cap = cv2.VideoCapture(str(local_video)) total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) + fps = cap.get(cv2.CAP_PROP_FPS) or 30 + w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) + h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) + ret, first_frame = cap.read() + cap.release() + state["total_frames"] = total_frames + state["video_fps"] = fps + state["video_width"] = w + state["video_height"] = h + + # Size container to video dimensions + video_container.style( + f"width: {w}px; max-width: 100%; aspect-ratio: {w}/{h};" + f" resize: horizontal; overflow: hidden;" + ) - # Display first frame - cap.set(cv2.CAP_PROP_POS_FRAMES, 0) - ret, frame = cap.read() + # Show first frame + video_display.content = '' + reset_tracker_state() if ret: - # Size container to video's native dimensions, lock aspect ratio for resize - h, w = frame.shape[:2] - video_container.style( - f"width: {w}px; max-width: 100%; aspect-ratio: {w}/{h};" - f" resize: horizontal; overflow: hidden;" - ) - - _, buffer = cv2.imencode(".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, 85]) - img_base64 = base64.b64encode(buffer).decode('utf-8') - video_display.set_source(f"data:image/jpeg;base64,{img_base64}") - - cap.release() + _, buf = cv2.imencode('.jpg', first_frame, [cv2.IMWRITE_JPEG_QUALITY, 80]) + b64 = base64.b64encode(buf).decode() + video_display.set_source(f'data:image/jpeg;base64,{b64}') + logger.info(f"Displayed first frame ({w}x{h})") # Enable time slider for playback time_slider.enable() @@ -676,14 +730,14 @@ def pause_tracking(): state["pause_event"].clear() pause_btn.props("icon=pause") pause_btn.text = "Pause" - progress_label.text = "Resuming..." + status_indicator.text = "Resuming..." ui.notify("Resumed", type="info") else: # Currently running, pause state["pause_event"].set() pause_btn.props("icon=play_arrow") pause_btn.text = "Resume" - progress_label.text = "Paused" + status_indicator.text = "Paused" ui.notify("Paused", type="warning") @@ -691,7 +745,7 @@ def stop_tracking(): """Hard stop - terminates processing""" if state["stop_event"]: state["stop_event"].set() - progress_label.text = "Stopping..." + status_indicator.text = "Stopping..." ui.notify("Stopping tracking...", type="negative") async def start_tracking(): """Start tracking on already-loaded video with already-loaded model""" @@ -700,6 +754,7 @@ async def start_tracking(): return state["processing"] = True + reset_tracker_state() state["stop_event"] = threading.Event() # Hard stop state["pause_event"] = threading.Event() # Pause (starts clear = not paused) state["skip_frames_event"] = {"skip_amount": 0} # Skip forward @@ -715,40 +770,60 @@ async def start_tracking(): try: # Show mode in progress label if detection_only_checkbox.value: - progress_label.text = "Starting detection..." + status_indicator.text = "Starting detection..." else: tracker_type = state.get("tracker_type", "Unknown") - progress_label.text = f"Starting tracking ({tracker_type})..." - progress.value = 0 + status_indicator.text = f"Starting tracking ({tracker_type})..." # Use already-loaded video and model from state local_video = state["video_path"] model = state["loaded_model"] # Frame callback for real-time UI updates - # Capture display update interval from slider display_interval = int(display_update_slider.value) - - async def frame_callback(annotated_frame, frame_idx, total_frames): - """Update UI with current frame (throttled for performance)""" - state["current_frame"] = frame_idx - - # Update every N frames based on slider setting - if frame_idx % display_interval != 0 and frame_idx != total_frames - 1: - return - - # Convert frame to bytes for display + _track_colors = [ + '#00FF00', '#FF0000', '#0080FF', '#FFFF00', + '#FF00FF', '#00FFFF', '#FF8000', '#8000FF', + '#00FF80', '#FF0080', '#80FF00', '#0040FF', + ] + + async def frame_callback(frame, detections, frame_idx, total_frames): + """Update UI: JPEG frame + SVG bbox overlay on every callback.""" import cv2 - import base64 - _, buffer = cv2.imencode(".jpg", annotated_frame, [cv2.IMWRITE_JPEG_QUALITY, 85]) - img_base64 = base64.b64encode(buffer).decode('utf-8') - video_display.set_source(f"data:image/jpeg;base64,{img_base64}") + state["current_frame"] = frame_idx - # Update progress - progress.value = 0.1 + 0.8 * (frame_idx / total_frames) - progress_label.text = f"Tracking: Frame {frame_idx + 1}/{total_frames}" + # Update base JPEG image (rate controlled by Display Update slider) + _, buf = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, 50]) + b64 = base64.b64encode(buf).decode() + video_display.set_source(f'data:image/jpeg;base64,{b64}') + + # Update SVG overlay with detection bboxes + svg_rects = [] + det_only = detection_only_checkbox.value + if len(detections) > 0: + for i, bbox in enumerate(detections.xyxy): + x1, y1, x2, y2 = bbox + bw, bh = x2 - x1, y2 - y1 + conf = detections.confidence[i] + + if det_only: + color = _track_colors[0] + label = f"{conf:.2f}" + else: + tid = int(detections.tracker_id[i]) if detections.tracker_id is not None else 0 + color = _track_colors[tid % len(_track_colors)] + label = f"#{tid} {conf:.2f}" + + svg_rects.append( + f'' + f'{label}' + ) + + video_display.content = '\n'.join(svg_rects) - # Update time slider time_slider.set_value(frame_idx) time_label.text = f"Frame: {frame_idx} / {total_frames}" @@ -767,6 +842,7 @@ async def frame_callback(annotated_frame, frame_idx, total_frames): tracker_config=tracker_config, confidence=conf_slider.value, detection_only=detection_only_checkbox.value, + display_interval=display_interval, frame_callback=frame_callback, stop_event=state["stop_event"], pause_event=state["pause_event"], @@ -777,8 +853,7 @@ async def frame_callback(annotated_frame, frame_idx, total_frames): results = await tracker.process_video_realtime(str(local_video), output_dir) # Show results - progress.value = 1.0 - progress_label.text = "Complete!" + status_indicator.text = "Complete!" state["results"] = results @@ -802,12 +877,12 @@ async def frame_callback(annotated_frame, frame_idx, total_frames): except Exception as e: logger.error(f"Tracking failed: {e}", exc_info=True) try: - progress_label.text = f"Error: {str(e)}" + status_indicator.text = f"Error: {str(e)}" ui.notify(f"Error: {str(e)}", type="negative") except Exception as notify_error: logger.error(f"Failed to show error notification: {notify_error}") try: - progress_label.text = f"Error: {str(e)}" + status_indicator.text = f"Error: {str(e)}" except: pass diff --git a/collab_env/tracking_studio/video_converter.py b/collab_env/tracking_studio/video_converter.py index e7ed06b0..aaaba55e 100644 --- a/collab_env/tracking_studio/video_converter.py +++ b/collab_env/tracking_studio/video_converter.py @@ -37,7 +37,14 @@ def needs_conversion(video_path: Path) -> bool: codec = result.stdout.strip() logger.info(f"Video codec: {codec}") - return codec != "h264" + if codec != "h264": + return True + # H.264 in non-MP4 container (e.g. .mov) may not play in all browsers + ext = Path(video_path).suffix.lower() + if ext not in (".mp4", ".m4v"): + logger.info(f"H.264 in {ext} container — will remux to .mp4 for browser compatibility") + return True + return False except subprocess.CalledProcessError as e: logger.error(f"Failed to check video codec: {e}") @@ -48,39 +55,43 @@ def needs_conversion(video_path: Path) -> bool: raise -def convert_to_h264(input_path: Path, output_path: Path) -> Path: +def convert_to_h264(input_path: Path, output_path: Path, remux_only: bool = False) -> Path: """ Convert video to H.264 format using ffmpeg. Args: input_path: Original video file output_path: Output path for converted video + remux_only: If True, copy streams without re-encoding (fast container change) Returns: Path to converted video """ try: - logger.info(f"Converting {input_path} to H.264 format") - - cmd = [ - "ffmpeg", - "-i", - str(input_path), - "-c:v", - "libx264", - "-preset", - "fast", - "-crf", - "23", - "-c:a", - "aac", - "-b:a", - "128k", - "-movflags", - "+faststart", # Web optimization - "-y", # Overwrite output - str(output_path), - ] + if remux_only: + logger.info(f"Remuxing {input_path} to MP4 container (no re-encoding)") + cmd = [ + "ffmpeg", + "-i", str(input_path), + "-c", "copy", # Copy all streams without re-encoding + "-movflags", "+faststart", + "-y", + str(output_path), + ] + else: + logger.info(f"Converting {input_path} to H.264 format") + cmd = [ + "ffmpeg", + "-i", str(input_path), + "-c:v", "libx264", + "-preset", "fast", + "-crf", "23", + "-c:a", "aac", + "-b:a", "128k", + "-movflags", "+faststart", + "-y", + str(output_path), + ] subprocess.run(cmd, check=True, capture_output=True) diff --git a/collab_env/tracking_studio/video_processor.py b/collab_env/tracking_studio/video_processor.py index 6b1d8ea3..f8c629b7 100644 --- a/collab_env/tracking_studio/video_processor.py +++ b/collab_env/tracking_studio/video_processor.py @@ -27,6 +27,7 @@ def __init__( tracker_config: Dict, # ByteTrack parameters confidence: float = 0.5, detection_only: bool = False, + display_interval: int = 10, frame_callback: Callable[[np.ndarray, int, int], None] = None, stop_event: threading.Event = None, pause_event: threading.Event = None, @@ -40,6 +41,7 @@ def __init__( tracker_config: Tracker configuration dict confidence: Detection confidence threshold detection_only: If True, run detection without tracking (no track IDs) + display_interval: Update display every Nth frame (1 = every frame) frame_callback: Async callback for frame updates (frame, frame_idx, total_frames) stop_event: Threading event to signal hard stop pause_event: Threading event to signal pause/resume @@ -48,10 +50,12 @@ def __init__( self.model = model self.confidence = confidence self.detection_only = detection_only + self.display_interval = max(1, display_interval) self.frame_callback = frame_callback self.stop_event = stop_event or threading.Event() self.pause_event = pause_event or threading.Event() self.skip_frames_event = skip_frames_event or {"skip_amount": 0} + self._pending_update = None # Track in-flight UI update # Store tracker config for use with model.track() self.tracker_config = tracker_config @@ -83,10 +87,6 @@ def __init__( # Fast-forward: Skip frames for faster preview self.skip_frames = tracker_config.get("skip_frames", 1) # 1 = process every frame - # Annotators for visualization - self.box_annotator = sv.BoxAnnotator() - self.label_annotator = sv.LabelAnnotator() - logger.info( f"VideoTracker initialized (confidence: {self.confidence}, native_tracking: {self.use_native_tracking})" ) @@ -278,42 +278,19 @@ def _process_video_sync( } ) - # 4. Annotate frame for display - annotated_frame = frame.copy() - annotated_frame = self.box_annotator.annotate( - annotated_frame, tracked_detections - ) - - if self.detection_only: - # Labels with confidence only - labels = [ - f"{conf:.2f}" - for conf in tracked_detections.confidence - ] - else: - # Labels with track IDs - labels = [ - f"#{track_id} {conf:.2f}" - for track_id, conf in zip( - tracked_detections.tracker_id, tracked_detections.confidence + # 4. Send frame + detections to UI for display + is_last = (frame_idx >= total_frames - 1) + should_display = (frame_idx % self.display_interval == 0) or is_last + + if should_display and self.frame_callback and event_loop: + # Skip if previous UI update is still in-flight (prevents queue buildup) + if self._pending_update is None or self._pending_update.done(): + self._pending_update = asyncio.run_coroutine_threadsafe( + self.frame_callback( + frame, tracked_detections, frame_idx, total_frames + ), + event_loop ) - ] - annotated_frame = self.label_annotator.annotate( - annotated_frame, tracked_detections, labels=labels - ) - - # 6. Send frame to UI (schedule callback in main event loop) - if self.frame_callback and event_loop: - # Schedule callback in main event loop from background thread - future = asyncio.run_coroutine_threadsafe( - self.frame_callback(annotated_frame, frame_idx, total_frames), - event_loop - ) - # Wait for UI update to complete (with timeout to prevent blocking) - try: - future.result(timeout=2.0) - except Exception as e: - logger.warning(f"Frame callback failed: {e}") # Increment frame counter for next iteration frame_idx += 1 From cb1a2233160c6da87a1eb9cf8fcb3509816709c4 Mon Sep 17 00:00:00 2001 From: Dmitry Batenkov Date: Tue, 17 Feb 2026 19:15:44 -0500 Subject: [PATCH 09/21] usability improvements --- cloudbuild.yaml | 2 +- collab_env/tracking_studio/app.py | 131 ++++++++++++------ collab_env/tracking_studio/model_manager.py | 47 ++++--- collab_env/tracking_studio/video_processor.py | 15 +- docs/tracking/tracking_web_gui.md | 14 ++ 5 files changed, 138 insertions(+), 71 deletions(-) diff --git a/cloudbuild.yaml b/cloudbuild.yaml index ac526f6d..48cb536b 100644 --- a/cloudbuild.yaml +++ b/cloudbuild.yaml @@ -26,7 +26,7 @@ steps: - '--concurrency=1' - '--max-instances=5' - '--set-secrets=ROBOFLOW_API_KEY=roboflow-api-key:latest' - - '--allow-unauthenticated' + - '--no-allow-unauthenticated' images: - 'gcr.io/$PROJECT_ID/tracking-studio:latest' diff --git a/collab_env/tracking_studio/app.py b/collab_env/tracking_studio/app.py index 635aea24..3d118545 100644 --- a/collab_env/tracking_studio/app.py +++ b/collab_env/tracking_studio/app.py @@ -172,24 +172,26 @@ async def handle_upload(e): with ui.column().classes("gap-3").style("flex: 0 0 280px; min-width: 280px"): # Model card with ui.card().classes("w-full shadow-md p-3"): - ui.label("🤖 Model").classes("text-sm font-semibold mb-2") + ui.label("Model").classes("text-sm font-semibold mb-2") - # Radio with three options - with ui.row().classes("gap-2"): - model_source = ui.radio(["YOLO", "Roboflow", "Custom"], value="YOLO").classes("text-xs") + model_source = ui.select( + label="Source", + options=["YOLO", "Roboflow", "Custom"], + value="Roboflow", + ).classes("w-full") - # YOLO model selection (default visible) + # YOLO model selection yolo_container = ui.column().classes("w-full mt-2") + yolo_container.visible = False with yolo_container: yolo_model_input = ui.input( label="Model Name", placeholder="e.g., yolo11n.pt", value="yolo11n.pt" - ).classes("w-full").tooltip("Enter any YOLO model name (will auto-download if available)") + ).classes("w-full").tooltip("Enter any YOLO model name (will auto-download)") - # Roboflow model selection (hidden by default) + # Roboflow model selection (default visible) rf_container = ui.column().classes("w-full mt-2 gap-2") - rf_container.visible = False with rf_container: rf_project_input = ui.input( label="Project ID", @@ -197,22 +199,34 @@ async def handle_upload(e): value="dima-sdrkv/ratsmerged20260211" ).classes("w-full") - # Define the function BEFORE referencing it in the button + # Store raw version data for detail dialog + _rf_versions_raw = {} + async def list_rf_models(): """Query Roboflow for available model versions""" project_id = rf_project_input.value if not project_id: ui.notify("Please enter project ID", type="warning") return - try: rf_list_btn.disable() - # Call directly (synchronous HTTP request, no need for threading) versions = model_manager.list_roboflow_project_models(project_id) if versions: - rf_version_select.options = versions - rf_version_select.value = versions[0] + options = {} + _rf_versions_raw.clear() + for v in versions: + parts = [f"v{v['version']}"] + if v['name']: + parts.append(v['name']) + parts.append(f"{v['images']} imgs") + if v['map']: + parts.append(f"mAP {v['map']}") + options[v['version']] = " | ".join(parts) + _rf_versions_raw[v['version']] = v.get('raw', {}) + rf_version_select.options = options + rf_version_select.value = versions[0]['version'] rf_version_select.enable() + rf_detail_btn.visible = True ui.notify(f"Found {len(versions)} versions", type="positive") else: ui.notify("No versions found", type="warning") @@ -222,9 +236,26 @@ async def list_rf_models(): finally: rf_list_btn.enable() - # Now create the button and version select (after function definition) - with ui.row().classes("w-full gap-2"): + def show_version_detail(): + """Show full JSON for the selected version in a dialog""" + import json + ver = rf_version_select.value + raw = _rf_versions_raw.get(ver, {}) + if not raw: + ui.notify("No version data available", type="warning") + return + with ui.dialog() as dlg, ui.card().style("min-width: 500px; max-height: 80vh;"): + ui.label(f"Version {ver} Details").classes("text-sm font-semibold") + ui.code(json.dumps(raw, indent=2, default=str)).classes( + "w-full text-xs" + ).style("max-height: 60vh; overflow: auto;") + ui.button("Close", on_click=dlg.close).props("size=sm flat") + dlg.open() + + with ui.row().classes("w-full gap-2 items-center"): rf_list_btn = ui.button("List Models", on_click=list_rf_models).props("size=sm color=primary") + rf_detail_btn = ui.button("Details", on_click=show_version_detail).props("size=sm flat") + rf_detail_btn.visible = False rf_version_select = ui.select( label="Version", @@ -232,11 +263,10 @@ async def list_rf_models(): ).classes("w-full") rf_version_select.disable() - # Custom model upload (hidden by default) + # Custom model upload custom_container = ui.column().classes("w-full mt-2") custom_container.visible = False with custom_container: - # Upload widget for model weights async def handle_model_upload(e): """Handle model .pt file upload""" try: @@ -247,7 +277,7 @@ async def handle_model_upload(e): state["uploaded_model"] = uploaded_model_file ui.notify(f"Model uploaded: {e.name}", type="positive") logger.info(f"Model uploaded to: {uploaded_model_file}") - load_model_btn.enable() # Enable Load Model button + load_model_btn.enable() except Exception as error: logger.error(f"Model upload failed: {error}") ui.notify(f"Model upload failed: {error}", type="negative") @@ -260,27 +290,10 @@ async def handle_model_upload(e): # Toggle visibility based on model source def toggle_model_ui(e=None): value = model_source.value - if value == "YOLO": - yolo_container.visible = True - rf_container.visible = False - custom_container.visible = False - # Enable load button if YOLO model name is entered - if yolo_model_input.value: - load_model_btn.enable() - elif value == "Roboflow": - yolo_container.visible = False - rf_container.visible = True - custom_container.visible = False - # Enable load button if Roboflow model is selected - if rf_version_select.value: - load_model_btn.enable() - else: # Custom - yolo_container.visible = False - rf_container.visible = False - custom_container.visible = True - # Enable load button if custom model uploaded - if state.get("uploaded_model"): - load_model_btn.enable() + yolo_container.visible = (value == "YOLO") + rf_container.visible = (value == "Roboflow") + custom_container.visible = (value == "Custom") + enable_load_model_btn() def enable_load_model_btn(e=None): """Enable Load Model button when model is selected""" @@ -389,17 +402,21 @@ def enable_load_model_btn(e=None): # Skip frames (for fast-forward, not a ByteTrack param) with ui.row().classes("w-full items-center gap-2 mt-2"): - skip_frames_label = ui.label("Skip Frames: 1 frame").classes("text-xs") + skip_frames_label = ui.label("Skip: every frame").classes("text-xs") skip_frames_slider = ui.slider(min=1, max=30, step=1, value=1).style("width: 100px") skip_frames_slider.tooltip("Process every Nth frame (1 = all frames)") - skip_frames_slider.on("update:model-value", lambda e: skip_frames_label.set_text(f"Skip Frames: {int(e.args)} {'frame' if int(e.args) == 1 else 'frames'}")) + skip_frames_slider.on("update:model-value", lambda e: skip_frames_label.set_text( + "Skip: every frame" if int(e.args) == 1 else f"Skip: every {int(e.args)} frames" + )) # GUI refresh rate (display updates, not a ByteTrack param) with ui.row().classes("w-full items-center gap-2 mt-1"): - display_update_label = ui.label("Display Update: every frame").classes("text-xs") + display_update_label = ui.label("Display: every frame").classes("text-xs") display_update_slider = ui.slider(min=1, max=30, step=1, value=1).style("width: 100px") - display_update_slider.tooltip("Update display every Nth frame (1 = smoothest, higher = skip display frames)") - display_update_slider.on("update:model-value", lambda e: display_update_label.set_text(f"Display Update: {int(e.args)} {'frame' if int(e.args) == 1 else 'frames'}")) + display_update_slider.tooltip("Update display every Nth frame (1 = every frame, higher = skip display frames)") + display_update_slider.on("update:model-value", lambda e: display_update_label.set_text( + "Display: every frame" if int(e.args) == 1 else f"Display: every {int(e.args)} frames" + )) # RIGHT: Controls + Preview (stacked vertically) with ui.column().classes("flex-grow gap-3"): @@ -513,6 +530,12 @@ def read_frame(path, idx): "width: 100%;" ) + # Debug: actual parameters passed to detector/tracker + debug_params_card = ui.card().classes("w-full shadow-md p-3 hidden") + with debug_params_card: + ui.label("Active Parameters").classes("text-xs font-semibold mb-1") + debug_params_label = ui.label("").classes("text-xs font-mono text-gray-600").style("white-space: pre-wrap;") + # Results (initially hidden, separate row) results_container = ui.card().classes("w-full shadow-md p-3 hidden") with results_container: @@ -758,7 +781,9 @@ async def start_tracking(): state["stop_event"] = threading.Event() # Hard stop state["pause_event"] = threading.Event() # Pause (starts clear = not paused) state["skip_frames_event"] = {"skip_amount": 0} # Skip forward - state["current_frame"] = 0 + # Resume from current slider position (preserved after stop) + start_frame = int(time_slider.value) if time_slider.value else 0 + state["current_frame"] = start_frame start_btn.disable() pause_btn.text = "Pause" pause_btn.props("icon=pause") @@ -836,6 +861,18 @@ async def frame_callback(frame, detections, frame_idx, total_frames): if hasattr(widget, 'value'): tracker_config[param_name] = widget.value + # Log and display actual parameters + active_params = { + "start_frame": start_frame, + "confidence": conf_slider.value, + "detection_only": detection_only_checkbox.value, + "display_interval": display_interval, + **tracker_config, + } + logger.info(f"Tracking params: {active_params}") + debug_params_label.text = " ".join(f"{k}={v}" for k, v in active_params.items()) + debug_params_card.classes(remove="hidden") + # Initialize tracker with dynamic parameters tracker = VideoTracker( model=model, @@ -850,7 +887,9 @@ async def frame_callback(frame, detections, frame_idx, total_frames): ) output_dir = f"/tmp/outputs/{session_id}" - results = await tracker.process_video_realtime(str(local_video), output_dir) + results = await tracker.process_video_realtime( + str(local_video), output_dir, start_frame=start_frame + ) # Show results status_indicator.text = "Complete!" diff --git a/collab_env/tracking_studio/model_manager.py b/collab_env/tracking_studio/model_manager.py index e2705f6a..0cca7ab4 100644 --- a/collab_env/tracking_studio/model_manager.py +++ b/collab_env/tracking_studio/model_manager.py @@ -300,7 +300,7 @@ def _load_roboflow_with_pipeline(self, model_id: str): logger.error(error_msg) raise ValueError(error_msg) from e - def list_roboflow_project_models(self, project_id: str) -> List[str]: + def list_roboflow_project_models(self, project_id: str) -> List[dict]: """ Query Roboflow API for available model versions in a project. @@ -308,9 +308,10 @@ def list_roboflow_project_models(self, project_id: str) -> List[str]: project_id: Project ID in format "workspace/project" (e.g., "dima-sdrkv/ratsmerged20260211") Returns: - List of version numbers (e.g., ["1", "2", "3"]) + List of dicts with keys: version, name, images, map """ import requests + from datetime import datetime if not self.roboflow_api_key: raise ValueError("ROBOFLOW_API_KEY not set") @@ -334,26 +335,34 @@ def list_roboflow_project_models(self, project_id: str) -> List[str]: data = response.json() - # Extract version numbers from response versions = [] if 'versions' in data: - for version_data in data['versions']: - # Try different fields that might contain the version number - version_num = version_data.get('id') - - # If id is a full path (workspace/project/version), extract just the version - if version_num and isinstance(version_num, str) and '/' in version_num: - version_num = version_num.split('/')[-1] # Get last part - - # Also check for a 'version' field + for vd in data['versions']: + version_num = vd.get('id', '') + if isinstance(version_num, str) and '/' in version_num: + version_num = version_num.split('/')[-1] if not version_num: - version_num = version_data.get('version') - - if version_num: - versions.append(str(version_num)) - - logger.info(f"Found {len(versions)} versions: {versions}") - return sorted(versions, key=lambda x: int(x) if x.isdigit() else 0, reverse=True) + version_num = vd.get('version') + if not version_num: + continue + + map_val = vd.get('model', {}).get('map', '') + if map_val and str(map_val) != 'NaN': + map_str = f"{float(map_val):.1f}%" + else: + map_str = "" + + versions.append({ + "version": str(version_num), + "name": vd.get('name', ''), + "images": vd.get('images', 0), + "map": map_str, + "raw": vd, + }) + + versions.sort(key=lambda x: int(x['version']) if x['version'].isdigit() else 0, reverse=True) + logger.info(f"Found {len(versions)} versions: {[v['version'] for v in versions]}") + return versions except requests.exceptions.HTTPError as e: error_msg = f"Failed to query Roboflow project: HTTP {e.response.status_code}" diff --git a/collab_env/tracking_studio/video_processor.py b/collab_env/tracking_studio/video_processor.py index f8c629b7..e1aa4a78 100644 --- a/collab_env/tracking_studio/video_processor.py +++ b/collab_env/tracking_studio/video_processor.py @@ -129,7 +129,7 @@ def _create_bytetrack_config(self, config: Dict) -> str: return temp_file.name def _process_video_sync( - self, video_path: str, output_dir: str, event_loop + self, video_path: str, output_dir: str, event_loop, start_frame: int = 0 ) -> Dict[str, Any]: """ Synchronous video processing function (runs in background thread). @@ -138,6 +138,7 @@ def _process_video_sync( video_path: Path to input video output_dir: Directory for output CSV event_loop: Main asyncio event loop for scheduling UI updates + start_frame: Frame index to start processing from (0-based) Returns: Dict with tracking_csv path and stats @@ -160,7 +161,10 @@ def _process_video_sync( detections_list = [] tracking_list = [] - frame_idx = 0 + frame_idx = start_frame + if start_frame > 0: + cap.set(cv2.CAP_PROP_POS_FRAMES, start_frame) + logger.info(f"Starting from frame {start_frame}") while frame_idx < total_frames: # Check if stop was requested (hard stop) if self.stop_event.is_set(): @@ -354,7 +358,7 @@ def _process_video_sync( } async def process_video_realtime( - self, video_path: str, output_dir: str + self, video_path: str, output_dir: str, start_frame: int = 0 ) -> Dict[str, Any]: """ Process video frame-by-frame with real-time UI updates. @@ -365,6 +369,7 @@ async def process_video_realtime( Args: video_path: Path to input video output_dir: Directory for output CSV + start_frame: Frame index to start processing from (0-based) Returns: Dict with tracking_csv path and stats @@ -373,9 +378,9 @@ async def process_video_realtime( loop = asyncio.get_running_loop() # Run processing in background thread - logger.info("Starting video processing in background thread...") + logger.info(f"Starting video processing in background thread (frame {start_frame})...") result = await asyncio.to_thread( - self._process_video_sync, video_path, output_dir, loop + self._process_video_sync, video_path, output_dir, loop, start_frame ) logger.info("Video processing complete") diff --git a/docs/tracking/tracking_web_gui.md b/docs/tracking/tracking_web_gui.md index c4386a02..bede6029 100644 --- a/docs/tracking/tracking_web_gui.md +++ b/docs/tracking/tracking_web_gui.md @@ -337,6 +337,20 @@ results = await tracker.process_video_realtime("input.mp4", "/tmp/output") print(f"Saved to: {results['tracking_csv']}") ``` +## Cloud Deployment + +The Tracking Studio can be deployed to Google Cloud Run using the provided `Dockerfile.tracking-studio` and `cloudbuild.yaml`: + +```bash +gcloud builds submit --config=cloudbuild.yaml +``` + +This builds a CPU-only Docker image and deploys to Cloud Run with 4GB RAM and 2 vCPUs. + +**Not recommended for interactive use.** The real-time frame preview relies on WebSocket streaming between the browser and server. Cloud Run's request-based scaling, cold starts, and network latency make the interactive experience significantly worse than running locally. For best results, run the studio on a local machine or a persistent VM with a GPU. + +Cloud Run deployment is better suited for batch processing or short demo sessions where latency is acceptable. + ## References - [ByteTrack Paper](https://arxiv.org/abs/2110.06864) From d501ba476c4fc597c9e0c137952dbd6cec876e07 Mon Sep 17 00:00:00 2001 From: Dmitry Batenkov Date: Tue, 17 Feb 2026 19:29:47 -0500 Subject: [PATCH 10/21] save/restore preferences --- collab_env/tracking_studio/app.py | 85 ++++++++++++++++++++++++++++--- 1 file changed, 78 insertions(+), 7 deletions(-) diff --git a/collab_env/tracking_studio/app.py b/collab_env/tracking_studio/app.py index 3d118545..b8dfdaa4 100644 --- a/collab_env/tracking_studio/app.py +++ b/collab_env/tracking_studio/app.py @@ -17,11 +17,36 @@ import base64 from loguru import logger +import json as _json + from .gcs_browser import GCSVideoBrowser from .model_manager import ModelManager from .video_processor import VideoTracker from .video_converter import convert_to_h264, needs_conversion +# Persistent preferences file (last used model/video settings) +_PREFS_FILE = Path.home() / ".tracking_studio.json" + + +def load_preferences() -> dict: + """Load saved preferences from dot file.""" + try: + if _PREFS_FILE.exists(): + return _json.loads(_PREFS_FILE.read_text()) + except Exception as e: + logger.warning(f"Failed to load preferences: {e}") + return {} + + +def save_preferences(prefs: dict): + """Save preferences to dot file (merges with existing).""" + try: + existing = load_preferences() + existing.update(prefs) + _PREFS_FILE.write_text(_json.dumps(existing, indent=2)) + except Exception as e: + logger.warning(f"Failed to save preferences: {e}") + # Load ByteTrack parameter definitions def load_bytetrack_params(): @@ -59,6 +84,7 @@ def get_credentials_path(): async def index(): """Main tracking studio page""" session_id = str(uuid.uuid4())[:8] + prefs = load_preferences() # State variables (stored in page context) import threading @@ -143,8 +169,23 @@ def enable_load_video_btn(e=None): folder_select.on("update:model-value", update_video_list) video_select.on("update:model-value", enable_load_video_btn) + # Restore last used bucket from preferences + saved_bucket = prefs.get("video_bucket") + if saved_bucket and saved_bucket in buckets: + bucket_select.value = saved_bucket + if bucket_select.value: - ui.timer(0.1, lambda: update_folders(None), once=True) + async def _restore_gcs_selection(): + await update_folders(None) + saved_folder = prefs.get("video_folder", "") + if saved_folder and saved_folder in folder_select.options: + folder_select.value = saved_folder + await update_video_list(None) + saved_video = prefs.get("video_name") + if saved_video and saved_video in video_select.options: + video_select.value = saved_video + load_video_btn.enable() + ui.timer(0.1, _restore_gcs_selection, once=True) # Upload widget async def handle_upload(e): @@ -177,26 +218,27 @@ async def handle_upload(e): model_source = ui.select( label="Source", options=["YOLO", "Roboflow", "Custom"], - value="Roboflow", + value=prefs.get("model_source", "Roboflow"), ).classes("w-full") # YOLO model selection yolo_container = ui.column().classes("w-full mt-2") - yolo_container.visible = False + yolo_container.visible = (prefs.get("model_source", "Roboflow") == "YOLO") with yolo_container: yolo_model_input = ui.input( label="Model Name", placeholder="e.g., yolo11n.pt", - value="yolo11n.pt" + value=prefs.get("yolo_model_name", "yolo11n.pt") ).classes("w-full").tooltip("Enter any YOLO model name (will auto-download)") - # Roboflow model selection (default visible) + # Roboflow model selection rf_container = ui.column().classes("w-full mt-2 gap-2") + rf_container.visible = (prefs.get("model_source", "Roboflow") == "Roboflow") with rf_container: rf_project_input = ui.input( label="Project ID", placeholder="workspace/project", - value="dima-sdrkv/ratsmerged20260211" + value=prefs.get("rf_project_id", "") ).classes("w-full") # Store raw version data for detail dialog @@ -265,7 +307,7 @@ def show_version_detail(): # Custom model upload custom_container = ui.column().classes("w-full mt-2") - custom_container.visible = False + custom_container.visible = (prefs.get("model_source", "Roboflow") == "Custom") with custom_container: async def handle_model_upload(e): """Handle model .pt file upload""" @@ -308,6 +350,18 @@ def enable_load_model_btn(e=None): yolo_model_input.on("update:model-value", enable_load_model_btn) rf_version_select.on("update:model-value", enable_load_model_btn) + # Auto-fetch Roboflow versions if saved project exists + if prefs.get("model_source") == "Roboflow" and prefs.get("rf_project_id"): + async def _restore_rf_version(): + await list_rf_models() + saved_ver = prefs.get("rf_version") + if saved_ver and saved_ver in (rf_version_select.options or {}): + rf_version_select.value = saved_ver + enable_load_model_btn() + ui.timer(0.1, _restore_rf_version, once=True) + elif prefs.get("model_source") == "YOLO": + enable_load_model_btn() + # Parameters card params_card = ui.card().classes("w-full shadow-md p-3") with params_card: @@ -655,6 +709,14 @@ async def load_video(): status_label.text = "Video loaded ✓" ui.notify("Video loaded successfully", type="positive") + # Save video selection to preferences + video_prefs = {} + if gcs_browser and bucket_select.value: + video_prefs["video_bucket"] = bucket_select.value + video_prefs["video_folder"] = folder_select.value or "" + video_prefs["video_name"] = video_select.value or "" + save_preferences(video_prefs) + # Enable Start button if model is also loaded if state["model_loaded"]: start_btn.enable() @@ -731,6 +793,15 @@ async def load_model(): status_label.text = f"Model loaded ✓ ({tracker_type} tracking)" ui.notify("Model loaded successfully", type="positive") + # Save model selection to preferences + model_prefs = {"model_source": model_source.value} + if model_source.value == "YOLO": + model_prefs["yolo_model_name"] = yolo_model_input.value + elif model_source.value == "Roboflow": + model_prefs["rf_project_id"] = rf_project_input.value + model_prefs["rf_version"] = rf_version_select.value + save_preferences(model_prefs) + # Enable Start button if video is also loaded if state["video_loaded"]: start_btn.enable() From 1f2144e173e193b785dec968adbe2227598f64d0 Mon Sep 17 00:00:00 2001 From: Dmitry Batenkov Date: Wed, 18 Feb 2026 12:27:40 -0500 Subject: [PATCH 11/21] gui fixes --- collab_env/tracking_studio/app.py | 64 ++++++++++++++++--------------- 1 file changed, 34 insertions(+), 30 deletions(-) diff --git a/collab_env/tracking_studio/app.py b/collab_env/tracking_studio/app.py index b8dfdaa4..45168c82 100644 --- a/collab_env/tracking_studio/app.py +++ b/collab_env/tracking_studio/app.py @@ -94,7 +94,6 @@ async def index(): "selected_model": None, "processing": False, "results": None, - "uploaded_video": None, "uploaded_model": None, # Uploaded model .pt file "stop_event": None, # Hard stop "pause_event": None, # Pause/resume @@ -161,8 +160,8 @@ async def update_video_list(e): logger.error(f"Failed to list videos: {error}") def enable_load_video_btn(e=None): - """Enable Load Video button when video is selected""" - if video_select.value or state.get("uploaded_video"): + """Enable Load Video button when a GCS video is selected""" + if video_select.value: load_video_btn.enable() bucket_select.on("update:model-value", update_folders) @@ -189,15 +188,14 @@ async def _restore_gcs_selection(): # Upload widget async def handle_upload(e): - """Handle user video upload""" + """Handle user video upload and auto-load it""" try: upload_path = Path(f"/tmp/uploads/{session_id}") upload_path.mkdir(parents=True, exist_ok=True) uploaded_file = upload_path / e.name uploaded_file.write_bytes(e.content.read()) - state["uploaded_video"] = uploaded_file ui.notify(f"Uploaded: {e.name}") - load_video_btn.enable() # Enable Load Video button + await load_video(local_video=uploaded_file) except Exception as error: logger.error(f"Upload failed: {error}") ui.notify(f"Upload failed: {error}", type="negative") @@ -252,6 +250,11 @@ async def list_rf_models(): return try: rf_list_btn.disable() + rf_version_select.options = {} + rf_version_select.value = None + rf_version_select.disable() + rf_detail_btn.visible = False + _rf_versions_raw.clear() versions = model_manager.list_roboflow_project_models(project_id) if versions: options = {} @@ -360,7 +363,7 @@ async def _restore_rf_version(): enable_load_model_btn() ui.timer(0.1, _restore_rf_version, once=True) elif prefs.get("model_source") == "YOLO": - enable_load_model_btn() + ui.timer(0.1, enable_load_model_btn, once=True) # Parameters card params_card = ui.card().classes("w-full shadow-md p-3") @@ -606,8 +609,13 @@ def reset_tracker_state(): # Full predictor reset — Ultralytics will create a fresh one on next call model.predictor = None - async def load_video(): - """Load and prepare video for viewing/tracking""" + async def load_video(local_video=None): + """Load and prepare video for viewing/tracking. + + Args: + local_video: Path to a local video file (e.g. from upload). + If None, downloads the video selected in the GCS dropdowns. + """ from nicegui import context try: @@ -617,26 +625,22 @@ async def load_video(): # Capture client context before threading client = context.client - # Get video (either download from GCS or use uploaded) - if state.get("uploaded_video"): - # Use uploaded video - local_video = state["uploaded_video"] - status_label.text = "Using uploaded video..." - elif gcs_browser and bucket_select.value and video_select.value: - # Download from GCS - status_label.text = "Downloading video..." - bucket = bucket_select.value - folder = folder_select.value or "" - video_name = video_select.value - gcs_path = f"{bucket}/{video_name}" - - local_video_dir = Path(f"/tmp/videos/{session_id}") - local_video_dir.mkdir(parents=True, exist_ok=True) - local_video = local_video_dir / Path(video_name).name - - await asyncio.to_thread(gcs_browser.download_video, gcs_path, str(local_video)) - else: - raise ValueError("No video selected. Please select or upload a video.") + if local_video is None: + # Download from GCS dropdowns + if gcs_browser and bucket_select.value and video_select.value: + status_label.text = "Downloading video..." + bucket = bucket_select.value + folder = folder_select.value or "" + video_name = video_select.value + gcs_path = f"{bucket}/{video_name}" + + local_video_dir = Path(f"/tmp/videos/{session_id}") + local_video_dir.mkdir(parents=True, exist_ok=True) + local_video = local_video_dir / Path(video_name).name + + await asyncio.to_thread(gcs_browser.download_video, gcs_path, str(local_video)) + else: + raise ValueError("No video selected. Please select a video from the dropdowns.") # Ensure browser-compatible H.264 MP4 if await asyncio.to_thread(needs_conversion, local_video): @@ -1009,7 +1013,7 @@ async def frame_callback(frame, detections, frame_idx, total_frames): params_card.style(remove="opacity: 0.5; pointer-events: none;") # Wire up buttons to event handlers (after functions are defined) - load_video_btn.on_click(load_video) + load_video_btn.on_click(lambda: load_video()) load_model_btn.on_click(load_model) start_btn.on_click(start_tracking) pause_btn.on_click(lambda: pause_tracking()) From 838d1c2edfaeadbd2f6e4016f2e29d874203a61a Mon Sep 17 00:00:00 2001 From: Dmitry Batenkov Date: Wed, 18 Feb 2026 13:56:54 -0500 Subject: [PATCH 12/21] removed MAP from the model dropdown, somewhat misleading --- collab_env/tracking_studio/app.py | 2 -- 1 file changed, 2 deletions(-) diff --git a/collab_env/tracking_studio/app.py b/collab_env/tracking_studio/app.py index 45168c82..7dbf3f5e 100644 --- a/collab_env/tracking_studio/app.py +++ b/collab_env/tracking_studio/app.py @@ -264,8 +264,6 @@ async def list_rf_models(): if v['name']: parts.append(v['name']) parts.append(f"{v['images']} imgs") - if v['map']: - parts.append(f"mAP {v['map']}") options[v['version']] = " | ".join(parts) _rf_versions_raw[v['version']] = v.get('raw', {}) rf_version_select.options = options From dad43a33e3be1410d5f6ea2cb6881bbeb0ff4475 Mon Sep 17 00:00:00 2001 From: Dmitry Batenkov Date: Thu, 19 Feb 2026 16:36:15 -0500 Subject: [PATCH 13/21] lint passes --- collab_env/dashboard/app.py | 8 +- collab_env/data/gcs_utils.py | 1 + collab_env/tracking/thermal_processing.py | 4 +- collab_env/tracking/visualization.py | 3 + collab_env/tracking_studio/app.py | 440 ++++++++++++------ collab_env/tracking_studio/gcs_browser.py | 11 +- collab_env/tracking_studio/model_manager.py | 134 ++++-- collab_env/tracking_studio/video_converter.py | 38 +- collab_env/tracking_studio/video_processor.py | 89 ++-- pyproject.toml | 3 + 10 files changed, 499 insertions(+), 232 deletions(-) diff --git a/collab_env/dashboard/app.py b/collab_env/dashboard/app.py index 5c75f716..0e017cb2 100644 --- a/collab_env/dashboard/app.py +++ b/collab_env/dashboard/app.py @@ -310,9 +310,7 @@ def _refresh_data_browser(self, event=None): except Exception as e: logger.error(f"Error refreshing data browser: {e}") - self.status_pane.object = ( - f"

Error refreshing: {e}

" - ) + self.status_pane.object = f"

Error refreshing: {e}

" def _on_session_change(self, event): """Handle session selection change.""" @@ -2153,7 +2151,9 @@ def create_layout(self): cache_controls = pn.Column(self.cache_info_pane, self.clear_cache_button) nav_panel = pn.Column( - pn.Row("## Data Browser", pn.Spacer(width=100), self.refresh_browser_button), + pn.Row( + "## Data Browser", pn.Spacer(width=100), self.refresh_browser_button + ), self.session_select, self.bucket_type_toggle, self.file_tree, diff --git a/collab_env/data/gcs_utils.py b/collab_env/data/gcs_utils.py index aa204783..e2818c16 100644 --- a/collab_env/data/gcs_utils.py +++ b/collab_env/data/gcs_utils.py @@ -33,6 +33,7 @@ def __init__( credentials_path = default_path self.project_id = project_id + self.credentials_path: Union[str, Path, None] if credentials_path and os.path.exists(str(credentials_path)): logger.info(f"Using service account credentials from {credentials_path}") diff --git a/collab_env/tracking/thermal_processing.py b/collab_env/tracking/thermal_processing.py index 21b75d31..acf86712 100644 --- a/collab_env/tracking/thermal_processing.py +++ b/collab_env/tracking/thermal_processing.py @@ -243,7 +243,7 @@ def _build_colorbar(self) -> np.ndarray: gradient = np.linspace(1.0, 0.0, self.frame_height, dtype=np.float32).reshape( self.frame_height, 1 ) - gradient = np.repeat(gradient, bar_width, axis=1) + gradient = np.repeat(gradient, bar_width, axis=1) # type: ignore[assignment] rgba = self.cmap(gradient) colorbar = np.clip(rgba[..., :3] * 255.0, 0, 255).astype(np.uint8) @@ -272,7 +272,7 @@ def render(self, frame: np.ndarray) -> np.ndarray: if self.vmax == self.vmin: normalized = np.zeros_like(frame, dtype=np.float32) else: - normalized = (frame - self.vmin) / (self.vmax - self.vmin) + normalized = (frame - self.vmin) / (self.vmax - self.vmin) # type: ignore[assignment] normalized = np.clip(normalized, 0.0, 1.0) rgba = self.cmap(normalized) rgb = np.clip(rgba[..., :3] * 255.0, 0, 255).astype(np.uint8) diff --git a/collab_env/tracking/visualization.py b/collab_env/tracking/visualization.py index 5e76edb2..7076a63f 100644 --- a/collab_env/tracking/visualization.py +++ b/collab_env/tracking/visualization.py @@ -135,6 +135,7 @@ def overlay_tracks_on_video( # Get frame size sample_frame = cv2.imread(str(frame_paths[0])) + assert sample_frame is not None, f"Failed to read frame: {frame_paths[0]}" h, w = sample_frame.shape[:2] writer = cv2.VideoWriter( str(output_video), @@ -146,6 +147,8 @@ def overlay_tracks_on_video( for frame_path in frame_paths: frame_idx = int(frame_path.stem.split("_")[-1]) frame = cv2.imread(str(frame_path)) + if frame is None: + continue frame_tracks = df[df["frame"] == frame_idx] for _, row in frame_tracks.iterrows(): diff --git a/collab_env/tracking_studio/app.py b/collab_env/tracking_studio/app.py index 7dbf3f5e..b739d872 100644 --- a/collab_env/tracking_studio/app.py +++ b/collab_env/tracking_studio/app.py @@ -8,12 +8,12 @@ - CSV output download """ -from nicegui import ui, app +from nicegui import ui import asyncio from pathlib import Path +from typing import Optional import uuid import os -import io import base64 from loguru import logger @@ -52,10 +52,12 @@ def save_preferences(prefs: dict): def load_bytetrack_params(): """Load ByteTrack parameter schema from JSON""" import json + params_file = Path(__file__).parent / "bytetrack_params.json" - with open(params_file, 'r') as f: + with open(params_file, "r") as f: return json.load(f) + bytetrack_params_schema = load_bytetrack_params() @@ -71,6 +73,7 @@ def get_credentials_path(): return None # GCSClient will use Application Default Credentials +gcs_browser: "Optional[GCSVideoBrowser]" try: gcs_browser = GCSVideoBrowser(credentials_path=get_credentials_path()) except Exception as e: @@ -88,6 +91,7 @@ async def index(): # State variables (stored in page context) import threading + state = { "selected_bucket": None, "selected_video_path": None, @@ -110,7 +114,9 @@ async def index(): with ui.row().classes("w-full items-center mb-2"): ui.label("🎯 Video Tracking Studio").classes("text-2xl font-bold") ui.space() - ui.label("Real-time object detection and tracking").classes("text-sm text-gray-600") + ui.label("Real-time object detection and tracking").classes( + "text-sm text-gray-600" + ) # Row 1: Video source selection (Bucket | Folder | Video | Upload) with ui.card().classes("w-full shadow-md p-2"): @@ -118,7 +124,8 @@ async def index(): # GCS selection if gcs_browser: bucket_select = ui.select( - label="Bucket", options=[], + label="Bucket", + options=[], ).style("width: 250px") try: buckets = gcs_browser.list_buckets() @@ -132,7 +139,9 @@ async def index(): label="Folder", options=[""], value="", clearable=True ).style("width: 350px") - video_select = ui.select(label="Video", options=[]).classes("flex-grow") + video_select = ui.select(label="Video", options=[]).classes( + "flex-grow" + ) async def update_folders(e): """Update folder list when bucket changes""" @@ -174,6 +183,7 @@ def enable_load_video_btn(e=None): bucket_select.value = saved_bucket if bucket_select.value: + async def _restore_gcs_selection(): await update_folders(None) saved_folder = prefs.get("video_folder", "") @@ -184,6 +194,7 @@ async def _restore_gcs_selection(): if saved_video and saved_video in video_select.options: video_select.value = saved_video load_video_btn.enable() + ui.timer(0.1, _restore_gcs_selection, once=True) # Upload widget @@ -200,15 +211,19 @@ async def handle_upload(e): logger.error(f"Upload failed: {error}") ui.notify(f"Upload failed: {error}", type="negative") - upload = ui.upload( + ui.upload( on_upload=handle_upload, auto_upload=True, - ).props("accept=video/mp4,video/quicktime,video/x-msvideo dense flat").props("label=Upload").style("width: 120px; height: 40px") + ).props( + "accept=video/mp4,video/quicktime,video/x-msvideo dense flat" + ).props("label=Upload").style("width: 120px; height: 40px") # Row 2: Model/Params (left) | Controls + Preview (right) with ui.row().classes("w-full gap-3"): # LEFT: Model + Parameters (stacked) - with ui.column().classes("gap-3").style("flex: 0 0 280px; min-width: 280px"): + with ( + ui.column().classes("gap-3").style("flex: 0 0 280px; min-width: 280px") + ): # Model card with ui.card().classes("w-full shadow-md p-3"): ui.label("Model").classes("text-sm font-semibold mb-2") @@ -221,22 +236,30 @@ async def handle_upload(e): # YOLO model selection yolo_container = ui.column().classes("w-full mt-2") - yolo_container.visible = (prefs.get("model_source", "Roboflow") == "YOLO") + yolo_container.visible = ( + prefs.get("model_source", "Roboflow") == "YOLO" + ) with yolo_container: - yolo_model_input = ui.input( - label="Model Name", - placeholder="e.g., yolo11n.pt", - value=prefs.get("yolo_model_name", "yolo11n.pt") - ).classes("w-full").tooltip("Enter any YOLO model name (will auto-download)") + yolo_model_input = ( + ui.input( + label="Model Name", + placeholder="e.g., yolo11n.pt", + value=prefs.get("yolo_model_name", "yolo11n.pt"), + ) + .classes("w-full") + .tooltip("Enter any YOLO model name (will auto-download)") + ) # Roboflow model selection rf_container = ui.column().classes("w-full mt-2 gap-2") - rf_container.visible = (prefs.get("model_source", "Roboflow") == "Roboflow") + rf_container.visible = ( + prefs.get("model_source", "Roboflow") == "Roboflow" + ) with rf_container: rf_project_input = ui.input( label="Project ID", placeholder="workspace/project", - value=prefs.get("rf_project_id", "") + value=prefs.get("rf_project_id", ""), ).classes("w-full") # Store raw version data for detail dialog @@ -255,22 +278,29 @@ async def list_rf_models(): rf_version_select.disable() rf_detail_btn.visible = False _rf_versions_raw.clear() - versions = model_manager.list_roboflow_project_models(project_id) + versions = model_manager.list_roboflow_project_models( + project_id + ) if versions: options = {} _rf_versions_raw.clear() for v in versions: parts = [f"v{v['version']}"] - if v['name']: - parts.append(v['name']) + if v["name"]: + parts.append(v["name"]) parts.append(f"{v['images']} imgs") - options[v['version']] = " | ".join(parts) - _rf_versions_raw[v['version']] = v.get('raw', {}) + options[v["version"]] = " | ".join(parts) + _rf_versions_raw[v["version"]] = v.get( + "raw", {} + ) rf_version_select.options = options - rf_version_select.value = versions[0]['version'] + rf_version_select.value = versions[0]["version"] rf_version_select.enable() rf_detail_btn.visible = True - ui.notify(f"Found {len(versions)} versions", type="positive") + ui.notify( + f"Found {len(versions)} versions", + type="positive", + ) else: ui.notify("No versions found", type="warning") except Exception as error: @@ -282,22 +312,34 @@ async def list_rf_models(): def show_version_detail(): """Show full JSON for the selected version in a dialog""" import json + ver = rf_version_select.value raw = _rf_versions_raw.get(ver, {}) if not raw: ui.notify("No version data available", type="warning") return - with ui.dialog() as dlg, ui.card().style("min-width: 500px; max-height: 80vh;"): - ui.label(f"Version {ver} Details").classes("text-sm font-semibold") + with ( + ui.dialog() as dlg, + ui.card().style("min-width: 500px; max-height: 80vh;"), + ): + ui.label(f"Version {ver} Details").classes( + "text-sm font-semibold" + ) ui.code(json.dumps(raw, indent=2, default=str)).classes( "w-full text-xs" ).style("max-height: 60vh; overflow: auto;") - ui.button("Close", on_click=dlg.close).props("size=sm flat") + ui.button("Close", on_click=dlg.close).props( + "size=sm flat" + ) dlg.open() with ui.row().classes("w-full gap-2 items-center"): - rf_list_btn = ui.button("List Models", on_click=list_rf_models).props("size=sm color=primary") - rf_detail_btn = ui.button("Details", on_click=show_version_detail).props("size=sm flat") + rf_list_btn = ui.button( + "List Models", on_click=list_rf_models + ).props("size=sm color=primary") + rf_detail_btn = ui.button( + "Details", on_click=show_version_detail + ).props("size=sm flat") rf_detail_btn.visible = False rf_version_select = ui.select( @@ -308,8 +350,11 @@ def show_version_detail(): # Custom model upload custom_container = ui.column().classes("w-full mt-2") - custom_container.visible = (prefs.get("model_source", "Roboflow") == "Custom") + custom_container.visible = ( + prefs.get("model_source", "Roboflow") == "Custom" + ) with custom_container: + async def handle_model_upload(e): """Handle model .pt file upload""" try: @@ -323,28 +368,36 @@ async def handle_model_upload(e): load_model_btn.enable() except Exception as error: logger.error(f"Model upload failed: {error}") - ui.notify(f"Model upload failed: {error}", type="negative") + ui.notify( + f"Model upload failed: {error}", type="negative" + ) ui.upload( on_upload=handle_model_upload, auto_upload=True, - ).props("accept=.pt dense flat").props("label=Upload Model (.pt)").classes("w-full") + ).props("accept=.pt dense flat").props( + "label=Upload Model (.pt)" + ).classes("w-full") # Toggle visibility based on model source def toggle_model_ui(e=None): value = model_source.value - yolo_container.visible = (value == "YOLO") - rf_container.visible = (value == "Roboflow") - custom_container.visible = (value == "Custom") + yolo_container.visible = value == "YOLO" + rf_container.visible = value == "Roboflow" + custom_container.visible = value == "Custom" enable_load_model_btn() def enable_load_model_btn(e=None): """Enable Load Model button when model is selected""" if model_source.value == "YOLO" and yolo_model_input.value: load_model_btn.enable() - elif model_source.value == "Roboflow" and rf_version_select.value: + elif ( + model_source.value == "Roboflow" and rf_version_select.value + ): load_model_btn.enable() - elif model_source.value == "Custom" and state.get("uploaded_model"): + elif model_source.value == "Custom" and state.get( + "uploaded_model" + ): load_model_btn.enable() model_source.on("update:model-value", toggle_model_ui) @@ -352,13 +405,19 @@ def enable_load_model_btn(e=None): rf_version_select.on("update:model-value", enable_load_model_btn) # Auto-fetch Roboflow versions if saved project exists - if prefs.get("model_source") == "Roboflow" and prefs.get("rf_project_id"): + if prefs.get("model_source") == "Roboflow" and prefs.get( + "rf_project_id" + ): + async def _restore_rf_version(): await list_rf_models() saved_ver = prefs.get("rf_version") - if saved_ver and saved_ver in (rf_version_select.options or {}): + if saved_ver and saved_ver in ( + rf_version_select.options or {} + ): rf_version_select.value = saved_ver enable_load_model_btn() + ui.timer(0.1, _restore_rf_version, once=True) elif prefs.get("model_source") == "YOLO": ui.timer(0.1, enable_load_model_btn, once=True) @@ -371,8 +430,15 @@ async def _restore_rf_version(): # Detection confidence (not in ByteTrack params) with ui.column().classes("w-full gap-1"): conf_label = ui.label("Confidence: 0.50").classes("text-xs") - conf_slider = ui.slider(min=0.1, max=0.9, step=0.05, value=0.5).classes("w-full").tooltip("Detection confidence threshold") - conf_slider.on("update:model-value", lambda e: conf_label.set_text(f"Confidence: {e.args:.2f}")) + conf_slider = ( + ui.slider(min=0.1, max=0.9, step=0.05, value=0.5) + .classes("w-full") + .tooltip("Detection confidence threshold") + ) + conf_slider.on( + "update:model-value", + lambda e: conf_label.set_text(f"Confidence: {e.args:.2f}"), + ) # Dynamic ByteTrack parameters from JSON param_widgets = {} # Store references to UI elements @@ -384,7 +450,9 @@ async def _restore_rf_version(): min_val, max_val = param_config["range"] # Create label with tooltip - param_label = ui.label(f"{param_name.replace('_', ' ').title()}: {default_val:.2f}").classes("text-xs") + param_label = ui.label( + f"{param_name.replace('_', ' ').title()}: {default_val:.2f}" + ).classes("text-xs") param_label.tooltip(param_config["description"]) # Create slider @@ -393,15 +461,17 @@ async def _restore_rf_version(): min=min_val, max=max_val, step=step, - value=default_val + value=default_val, ).classes("w-full") # Update label on change param_slider.on( "update:model-value", - lambda e, lbl=param_label, name=param_name: lbl.set_text( + lambda e, + lbl=param_label, + name=param_name: lbl.set_text( f"{name.replace('_', ' ').title()}: {e.args:.2f}" - ) + ), ) param_widgets[param_name] = param_slider @@ -409,31 +479,36 @@ async def _restore_rf_version(): # Int slider default_val = param_config["default"] min_val = param_config["range"][0] - max_val = param_config["range"][1] if param_config["range"][1] else 300 + max_val = ( + param_config["range"][1] + if param_config["range"][1] + else 300 + ) - param_label = ui.label(f"{param_name.replace('_', ' ').title()}: {default_val}").classes("text-xs") + param_label = ui.label( + f"{param_name.replace('_', ' ').title()}: {default_val}" + ).classes("text-xs") param_label.tooltip(param_config["description"]) param_slider = ui.slider( - min=min_val, - max=max_val, - step=1, - value=default_val + min=min_val, max=max_val, step=1, value=default_val ).classes("w-full") param_slider.on( "update:model-value", - lambda e, lbl=param_label, name=param_name: lbl.set_text( + lambda e, + lbl=param_label, + name=param_name: lbl.set_text( f"{name.replace('_', ' ').title()}: {int(e.args)}" - ) + ), ) param_widgets[param_name] = param_slider elif param_config["type"] == "bool": # Checkbox param_checkbox = ui.checkbox( - param_name.replace('_', ' ').title(), - value=param_config["default"] + param_name.replace("_", " ").title(), + value=param_config["default"], ).classes("text-xs") param_checkbox.tooltip(param_config["description"]) param_widgets[param_name] = param_checkbox @@ -441,9 +516,9 @@ async def _restore_rf_version(): elif param_config["type"] == "string": # Dropdown for options param_select = ui.select( - label=param_name.replace('_', ' ').title(), + label=param_name.replace("_", " ").title(), options=param_config["options"], - value=param_config["default"] + value=param_config["default"], ).classes("w-full text-xs") param_select.tooltip(param_config["description"]) param_widgets[param_name] = param_select @@ -453,25 +528,49 @@ async def _restore_rf_version(): "Detection only (no tracking)", value=False, ).classes("text-xs mt-2") - detection_only_checkbox.tooltip("Run detection without ByteTrack — shows raw detections per frame") + detection_only_checkbox.tooltip( + "Run detection without ByteTrack — shows raw detections per frame" + ) # Skip frames (for fast-forward, not a ByteTrack param) with ui.row().classes("w-full items-center gap-2 mt-2"): - skip_frames_label = ui.label("Skip: every frame").classes("text-xs") - skip_frames_slider = ui.slider(min=1, max=30, step=1, value=1).style("width: 100px") - skip_frames_slider.tooltip("Process every Nth frame (1 = all frames)") - skip_frames_slider.on("update:model-value", lambda e: skip_frames_label.set_text( - "Skip: every frame" if int(e.args) == 1 else f"Skip: every {int(e.args)} frames" - )) + skip_frames_label = ui.label("Skip: every frame").classes( + "text-xs" + ) + skip_frames_slider = ui.slider( + min=1, max=30, step=1, value=1 + ).style("width: 100px") + skip_frames_slider.tooltip( + "Process every Nth frame (1 = all frames)" + ) + skip_frames_slider.on( + "update:model-value", + lambda e: skip_frames_label.set_text( + "Skip: every frame" + if int(e.args) == 1 + else f"Skip: every {int(e.args)} frames" + ), + ) # GUI refresh rate (display updates, not a ByteTrack param) with ui.row().classes("w-full items-center gap-2 mt-1"): - display_update_label = ui.label("Display: every frame").classes("text-xs") - display_update_slider = ui.slider(min=1, max=30, step=1, value=1).style("width: 100px") - display_update_slider.tooltip("Update display every Nth frame (1 = every frame, higher = skip display frames)") - display_update_slider.on("update:model-value", lambda e: display_update_label.set_text( - "Display: every frame" if int(e.args) == 1 else f"Display: every {int(e.args)} frames" - )) + display_update_label = ui.label("Display: every frame").classes( + "text-xs" + ) + display_update_slider = ui.slider( + min=1, max=30, step=1, value=1 + ).style("width: 100px") + display_update_slider.tooltip( + "Update display every Nth frame (1 = every frame, higher = skip display frames)" + ) + display_update_slider.on( + "update:model-value", + lambda e: display_update_label.set_text( + "Display: every frame" + if int(e.args) == 1 + else f"Display: every {int(e.args)} frames" + ), + ) # RIGHT: Controls + Preview (stacked vertically) with ui.column().classes("flex-grow gap-3"): @@ -479,20 +578,28 @@ async def _restore_rf_version(): with ui.card().classes("w-full shadow-md p-3"): # Row 1: Load buttons with ui.row().classes("w-full items-center gap-2 mb-2"): - load_video_btn = ui.button("Load Video").props("color=primary icon=video_file") + load_video_btn = ui.button("Load Video").props( + "color=primary icon=video_file" + ) load_video_btn.disable() # Enabled when video selected - load_model_btn = ui.button("Load Model").props("color=primary icon=model_training") + load_model_btn = ui.button("Load Model").props( + "color=primary icon=model_training" + ) load_model_btn.disable() # Enabled when model selected ui.separator().props("vertical") with ui.column().classes("flex-grow gap-1"): - status_label = ui.label("Select video and model").classes("text-xs") + status_label = ui.label("Select video and model").classes( + "text-xs" + ) # Row 2: Playback controls with ui.row().classes("w-full items-center gap-2"): - start_btn = ui.button("Start Tracking").props("color=positive icon=play_arrow") + start_btn = ui.button("Start Tracking").props( + "color=positive icon=play_arrow" + ) start_btn.disable() # Enabled when both video and model loaded pause_btn = ui.button("Pause").props("color=warning icon=pause") @@ -503,12 +610,18 @@ async def _restore_rf_version(): ui.separator().props("vertical") - status_indicator = ui.label("Ready").classes("text-xs flex-grow") + status_indicator = ui.label("Ready").classes( + "text-xs flex-grow" + ) # Time slider for seeking with ui.column().classes("w-full gap-1 mt-2"): - time_label = ui.label("Frame: 0 / 0").classes("text-xs text-gray-600") - time_slider = ui.slider(min=0, max=100, value=0).classes("w-full") + time_label = ui.label("Frame: 0 / 0").classes( + "text-xs text-gray-600" + ) + time_slider = ui.slider(min=0, max=100, value=0).classes( + "w-full" + ) time_slider.disable() _preview_pending = [False] @@ -520,8 +633,9 @@ async def preview_frame_on_drag(e): _preview_pending[0] = True try: import cv2 + target_frame = int(e.args) - video_display.content = '' # Clear SVG overlay + video_display.content = "" # Clear SVG overlay if not state.get("processing"): reset_tracker_state() @@ -536,38 +650,60 @@ def read_frame(path, idx): read_frame, state["video_path"], target_frame ) if frame is not None: - _, buf = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, 50]) + _, buf = cv2.imencode( + ".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, 50] + ) b64 = base64.b64encode(buf).decode() - video_display.set_source(f'data:image/jpeg;base64,{b64}') + video_display.set_source( + f"data:image/jpeg;base64,{b64}" + ) total_frames = state.get("total_frames", target_frame) - time_label.text = f"Frame: {target_frame} / {total_frames}" + time_label.text = ( + f"Frame: {target_frame} / {total_frames}" + ) finally: _preview_pending[0] = False async def seek_to_frame(e): """Seek to specific frame when slider is released""" target_frame = int(e.args) - if state.get("processing") and state.get("skip_frames_event"): + if state.get("processing") and state.get( + "skip_frames_event" + ): current_frame = state.get("current_frame", 0) if target_frame != current_frame: - state["skip_frames_event"]["skip_amount"] = target_frame - current_frame - ui.notify(f"Seeking to frame {target_frame}...", type="info") + state["skip_frames_event"]["skip_amount"] = ( + target_frame - current_frame + ) + ui.notify( + f"Seeking to frame {target_frame}...", + type="info", + ) elif state.get("video_path"): # Not processing: show the frame at release position def read_frame(path, idx): import cv2 + cap = cv2.VideoCapture(str(path)) cap.set(cv2.CAP_PROP_POS_FRAMES, idx) ret, f = cap.read() cap.release() return f if ret else None - frame = await asyncio.to_thread(read_frame, state["video_path"], target_frame) + + frame = await asyncio.to_thread( + read_frame, state["video_path"], target_frame + ) if frame is not None: import cv2 - _, buf = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, 70]) + + _, buf = cv2.imencode( + ".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, 70] + ) b64 = base64.b64encode(buf).decode() - video_display.set_source(f'data:image/jpeg;base64,{b64}') + video_display.set_source( + f"data:image/jpeg;base64,{b64}" + ) # Live preview during drag time_slider.on("update:model-value", preview_frame_on_drag) @@ -577,19 +713,23 @@ def read_frame(path, idx): # Preview card with ui.card().classes("w-full shadow-md p-3"): ui.label("Live Preview").classes("text-sm font-semibold mb-2") - video_container = ui.element("div").classes("border-2 border-gray-200 rounded bg-gray-50").style( - "max-width: 100%; resize: horizontal; overflow: hidden;" + video_container = ( + ui.element("div") + .classes("border-2 border-gray-200 rounded bg-gray-50") + .style("max-width: 100%; resize: horizontal; overflow: hidden;") ) with video_container: - video_display = ui.interactive_image('').style( - "width: 100%;" - ) + video_display = ui.interactive_image("").style("width: 100%;") # Debug: actual parameters passed to detector/tracker debug_params_card = ui.card().classes("w-full shadow-md p-3 hidden") with debug_params_card: ui.label("Active Parameters").classes("text-xs font-semibold mb-1") - debug_params_label = ui.label("").classes("text-xs font-mono text-gray-600").style("white-space: pre-wrap;") + debug_params_label = ( + ui.label("") + .classes("text-xs font-mono text-gray-600") + .style("white-space: pre-wrap;") + ) # Results (initially hidden, separate row) results_container = ui.card().classes("w-full shadow-md p-3 hidden") @@ -597,13 +737,15 @@ def read_frame(path, idx): with ui.row().classes("w-full items-center gap-3"): ui.label("✅ Results").classes("text-sm font-semibold") stats_label = ui.label().classes("text-sm flex-grow") - download_track_btn = ui.button("Download CSV").props("color=primary icon=download size=sm") + download_track_btn = ui.button("Download CSV").props( + "color=primary icon=download size=sm" + ) # Event handlers def reset_tracker_state(): """Reset YOLO model's internal tracker so track IDs start fresh.""" model = state.get("loaded_model") - if model and hasattr(model, 'predictor') and model.predictor is not None: + if model and hasattr(model, "predictor") and model.predictor is not None: # Full predictor reset — Ultralytics will create a fresh one on next call model.predictor = None @@ -628,7 +770,6 @@ async def load_video(local_video=None): if gcs_browser and bucket_select.value and video_select.value: status_label.text = "Downloading video..." bucket = bucket_select.value - folder = folder_select.value or "" video_name = video_select.value gcs_path = f"{bucket}/{video_name}" @@ -636,21 +777,37 @@ async def load_video(local_video=None): local_video_dir.mkdir(parents=True, exist_ok=True) local_video = local_video_dir / Path(video_name).name - await asyncio.to_thread(gcs_browser.download_video, gcs_path, str(local_video)) + await asyncio.to_thread( + gcs_browser.download_video, gcs_path, str(local_video) + ) else: - raise ValueError("No video selected. Please select a video from the dropdowns.") + raise ValueError( + "No video selected. Please select a video from the dropdowns." + ) # Ensure browser-compatible H.264 MP4 if await asyncio.to_thread(needs_conversion, local_video): converted_video = local_video.parent / f"{local_video.stem}_h264.mp4" # Check if codec is already h264 (just needs container remux) import subprocess + try: probe = subprocess.run( - ["ffprobe", "-v", "error", "-select_streams", "v:0", - "-show_entries", "stream=codec_name", - "-of", "default=noprint_wrappers=1:nokey=1", str(local_video)], - capture_output=True, text=True, check=True + [ + "ffprobe", + "-v", + "error", + "-select_streams", + "v:0", + "-show_entries", + "stream=codec_name", + "-of", + "default=noprint_wrappers=1:nokey=1", + str(local_video), + ], + capture_output=True, + text=True, + check=True, ) is_h264 = probe.stdout.strip() == "h264" except Exception: @@ -658,10 +815,14 @@ async def load_video(local_video=None): if is_h264: status_label.text = "Remuxing to MP4..." - await asyncio.to_thread(convert_to_h264, local_video, converted_video, remux_only=True) + await asyncio.to_thread( + convert_to_h264, local_video, converted_video, remux_only=True + ) else: status_label.text = "Converting to H.264..." - await asyncio.to_thread(convert_to_h264, local_video, converted_video) + await asyncio.to_thread( + convert_to_h264, local_video, converted_video + ) local_video = converted_video with client: ui.notify("Video converted to H.264") @@ -674,6 +835,7 @@ async def load_video(local_video=None): # Read video metadata and first frame import cv2 + cap = cv2.VideoCapture(str(local_video)) total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) fps = cap.get(cv2.CAP_PROP_FPS) or 30 @@ -694,12 +856,14 @@ async def load_video(local_video=None): ) # Show first frame - video_display.content = '' + video_display.content = "" reset_tracker_state() if ret: - _, buf = cv2.imencode('.jpg', first_frame, [cv2.IMWRITE_JPEG_QUALITY, 80]) + _, buf = cv2.imencode( + ".jpg", first_frame, [cv2.IMWRITE_JPEG_QUALITY, 80] + ) b64 = base64.b64encode(buf).decode() - video_display.set_source(f'data:image/jpeg;base64,{b64}') + video_display.set_source(f"data:image/jpeg;base64,{b64}") logger.info(f"Displayed first frame ({w}x{h})") # Enable time slider for playback @@ -726,7 +890,7 @@ async def load_video(local_video=None): except Exception as e: logger.error(f"Failed to load video: {e}", exc_info=True) with client: - status_label.text = f"Error loading video" + status_label.text = "Error loading video" ui.notify(f"Error: {str(e)}", type="negative") finally: with client: @@ -753,10 +917,12 @@ async def load_model(): version = rf_version_select.value if not project_id or not version: - raise ValueError("Please select a Roboflow model (project ID and version)") + raise ValueError( + "Please select a Roboflow model (project ID and version)" + ) # Validate project ID format (should be workspace/project) - project_parts = project_id.split('/') + project_parts = project_id.split("/") if len(project_parts) != 2: raise ValueError( f"Invalid project ID: '{project_id}'\n" @@ -789,7 +955,10 @@ async def load_model(): # Detect tracker type for display from ultralytics import YOLO - tracker_type = "YOLO Native" if isinstance(model, YOLO) else "Supervision" + + tracker_type = ( + "YOLO Native" if isinstance(model, YOLO) else "Supervision" + ) state["tracker_type"] = tracker_type status_label.text = f"Model loaded ✓ ({tracker_type} tracking)" @@ -812,7 +981,7 @@ async def load_model(): logger.error(f"Failed to load model: {e}", exc_info=True) # Restore context for error UI updates with client: - status_label.text = f"Error loading model" + status_label.text = "Error loading model" ui.notify(f"Error: {str(e)}", type="negative") finally: with client: @@ -836,13 +1005,13 @@ def pause_tracking(): status_indicator.text = "Paused" ui.notify("Paused", type="warning") - def stop_tracking(): """Hard stop - terminates processing""" if state["stop_event"]: state["stop_event"].set() status_indicator.text = "Stopping..." ui.notify("Stopping tracking...", type="negative") + async def start_tracking(): """Start tracking on already-loaded video with already-loaded model""" if not state.get("video_loaded") or not state.get("model_loaded"): @@ -880,20 +1049,30 @@ async def start_tracking(): # Frame callback for real-time UI updates display_interval = int(display_update_slider.value) _track_colors = [ - '#00FF00', '#FF0000', '#0080FF', '#FFFF00', - '#FF00FF', '#00FFFF', '#FF8000', '#8000FF', - '#00FF80', '#FF0080', '#80FF00', '#0040FF', + "#00FF00", + "#FF0000", + "#0080FF", + "#FFFF00", + "#FF00FF", + "#00FFFF", + "#FF8000", + "#8000FF", + "#00FF80", + "#FF0080", + "#80FF00", + "#0040FF", ] async def frame_callback(frame, detections, frame_idx, total_frames): """Update UI: JPEG frame + SVG bbox overlay on every callback.""" import cv2 + state["current_frame"] = frame_idx # Update base JPEG image (rate controlled by Display Update slider) - _, buf = cv2.imencode('.jpg', frame, [cv2.IMWRITE_JPEG_QUALITY, 50]) + _, buf = cv2.imencode(".jpg", frame, [cv2.IMWRITE_JPEG_QUALITY, 50]) b64 = base64.b64encode(buf).decode() - video_display.set_source(f'data:image/jpeg;base64,{b64}') + video_display.set_source(f"data:image/jpeg;base64,{b64}") # Update SVG overlay with detection bboxes svg_rects = [] @@ -908,7 +1087,11 @@ async def frame_callback(frame, detections, frame_idx, total_frames): color = _track_colors[0] label = f"{conf:.2f}" else: - tid = int(detections.tracker_id[i]) if detections.tracker_id is not None else 0 + tid = ( + int(detections.tracker_id[i]) + if detections.tracker_id is not None + else 0 + ) color = _track_colors[tid % len(_track_colors)] label = f"#{tid} {conf:.2f}" @@ -920,7 +1103,7 @@ async def frame_callback(frame, detections, frame_idx, total_frames): f'stroke="black" stroke-width="0.3">{label}' ) - video_display.content = '\n'.join(svg_rects) + video_display.content = "\n".join(svg_rects) time_slider.set_value(frame_idx) time_label.text = f"Frame: {frame_idx} / {total_frames}" @@ -931,7 +1114,7 @@ async def frame_callback(frame, detections, frame_idx, total_frames): } # Add ByteTrack parameters from param_widgets for param_name, widget in param_widgets.items(): - if hasattr(widget, 'value'): + if hasattr(widget, "value"): tracker_config[param_name] = widget.value # Log and display actual parameters @@ -943,7 +1126,9 @@ async def frame_callback(frame, detections, frame_idx, total_frames): **tracker_config, } logger.info(f"Tracking params: {active_params}") - debug_params_label.text = " ".join(f"{k}={v}" for k, v in active_params.items()) + debug_params_label.text = " ".join( + f"{k}={v}" for k, v in active_params.items() + ) debug_params_card.classes(remove="hidden") # Initialize tracker with dynamic parameters @@ -988,15 +1173,8 @@ async def frame_callback(frame, detections, frame_idx, total_frames): except Exception as e: logger.error(f"Tracking failed: {e}", exc_info=True) - try: - status_indicator.text = f"Error: {str(e)}" - ui.notify(f"Error: {str(e)}", type="negative") - except Exception as notify_error: - logger.error(f"Failed to show error notification: {notify_error}") - try: - status_indicator.text = f"Error: {str(e)}" - except: - pass + status_indicator.text = f"Error: {str(e)}" + ui.notify(f"Error: {str(e)}", type="negative") finally: state["processing"] = False diff --git a/collab_env/tracking_studio/gcs_browser.py b/collab_env/tracking_studio/gcs_browser.py index 96288089..76a9f720 100644 --- a/collab_env/tracking_studio/gcs_browser.py +++ b/collab_env/tracking_studio/gcs_browser.py @@ -4,7 +4,6 @@ Provides interface for browsing and downloading videos from Google Cloud Storage. """ -from pathlib import Path from typing import List, Dict from loguru import logger @@ -72,7 +71,7 @@ def list_folders(self, bucket: str, prefix: str = "") -> List[str]: if prefix: if not rel_path.startswith(prefix): continue - rel_path = rel_path[len(prefix):] + rel_path = rel_path[len(prefix) :] # Get first directory component after prefix if "/" in rel_path: @@ -81,7 +80,9 @@ def list_folders(self, bucket: str, prefix: str = "") -> List[str]: unique_folders.add(folder) folder_list = sorted(list(unique_folders)) - logger.info(f"Found {len(folder_list)} folder prefixes in {bucket}/{prefix}") + logger.info( + f"Found {len(folder_list)} folder prefixes in {bucket}/{prefix}" + ) return folder_list except Exception as e: @@ -109,7 +110,9 @@ def list_videos(self, bucket: str, prefix: str = "") -> List[Dict[str, str]]: all_files = [] for ext in video_extensions: - pattern = f"{bucket}/{prefix}**/{ext}" if prefix else f"{bucket}/**/{ext}" + pattern = ( + f"{bucket}/{prefix}**/{ext}" if prefix else f"{bucket}/**/{ext}" + ) files = self.gcs.glob(pattern) all_files.extend(files) diff --git a/collab_env/tracking_studio/model_manager.py b/collab_env/tracking_studio/model_manager.py index 0cca7ab4..9b8650d7 100644 --- a/collab_env/tracking_studio/model_manager.py +++ b/collab_env/tracking_studio/model_manager.py @@ -6,7 +6,7 @@ import os from pathlib import Path -from typing import List +from typing import List, Optional from loguru import logger from ultralytics import YOLO @@ -15,7 +15,7 @@ class ModelManager: """Manager for detection models (YOLO and Roboflow)""" - def __init__(self, roboflow_api_key: str = None): + def __init__(self, roboflow_api_key: Optional[str] = None): """ Initialize model manager. @@ -105,7 +105,7 @@ def _validate_roboflow_model_id(self, model_id: str) -> str: Returns properly formatted model ID. """ - parts = model_id.split('/') + parts = model_id.split("/") if len(parts) == 2: # project/version format @@ -133,7 +133,11 @@ def load_roboflow_model(self, model_id: str): Loaded Roboflow model or YOLO model from local file """ # Check if model_id is a local file path - if model_id.startswith('/') or model_id.startswith('~') or model_id.endswith('.pt'): + if ( + model_id.startswith("/") + or model_id.startswith("~") + or model_id.endswith(".pt") + ): logger.info(f"Loading Roboflow model from local file: {model_id}") model_path = Path(model_id).expanduser() @@ -142,7 +146,9 @@ def load_roboflow_model(self, model_id: str): logger.info(f"Loading YOLO model from: {model_path}") model = YOLO(str(model_path)) - logger.info(f"Successfully loaded Roboflow model from local file: {model_id}") + logger.info( + f"Successfully loaded Roboflow model from local file: {model_id}" + ) return model if not self.roboflow_api_key: @@ -156,36 +162,52 @@ def load_roboflow_model(self, model_id: str): # Try downloading .pt file first (for YOLO native tracking) # This provides better performance and supports all ByteTrack parameters try: - logger.info(f"Downloading Roboflow model weights for YOLO native tracking: {model_id}") + logger.info( + f"Downloading Roboflow model weights for YOLO native tracking: {model_id}" + ) model = self._load_roboflow_with_pipeline(model_id) - logger.info(f"Successfully loaded Roboflow model with native tracking: {model_id}") + logger.info( + f"Successfully loaded Roboflow model with native tracking: {model_id}" + ) return model except Exception as download_error: # Fallback to get_model() (inference API) if download fails logger.warning(f"Download failed: {download_error}") - logger.info(f"Attempting fallback: loading with inference API (Supervision tracking)") + logger.info( + "Attempting fallback: loading with inference API (Supervision tracking)" + ) try: from inference import get_model # Extract project/version from workspace/project/version if needed - parts = model_id.split('/') + parts = model_id.split("/") if len(parts) == 3: # workspace/project/version -> project/version project_version = f"{parts[1]}/{parts[2]}" - logger.info(f"Trying to load Roboflow model with get_model(): {project_version}") - model = get_model(model_id=project_version, api_key=self.roboflow_api_key) + logger.info( + f"Trying to load Roboflow model with get_model(): {project_version}" + ) + model = get_model( + model_id=project_version, api_key=self.roboflow_api_key + ) else: # Already project/version format - logger.info(f"Trying to load Roboflow model with get_model(): {model_id}") + logger.info( + f"Trying to load Roboflow model with get_model(): {model_id}" + ) model = get_model(model_id=model_id, api_key=self.roboflow_api_key) - logger.info(f"Successfully loaded Roboflow model via inference API: {model_id}") + logger.info( + f"Successfully loaded Roboflow model via inference API: {model_id}" + ) return model except ImportError: - logger.error("inference library not installed. Install with: pip install inference") + logger.error( + "inference library not installed. Install with: pip install inference" + ) raise except Exception as inference_error: # Both methods failed @@ -212,10 +234,12 @@ def _load_roboflow_with_pipeline(self, model_id: str): """ import requests - logger.info(f"Downloading Roboflow model weights for local inference: {model_id}") + logger.info( + f"Downloading Roboflow model weights for local inference: {model_id}" + ) # Parse model ID to get workspace/project/version - parts = model_id.split('/') + parts = model_id.split("/") if len(parts) == 2: # project/version format - need workspace raise ValueError( @@ -245,38 +269,42 @@ def _load_roboflow_with_pipeline(self, model_id: str): try: # Call /ptFile endpoint to get signed download URL - ptfile_url = f"https://api.roboflow.com/{workspace}/{project}/{version}/ptFile" + ptfile_url = ( + f"https://api.roboflow.com/{workspace}/{project}/{version}/ptFile" + ) logger.info(f"Requesting weights URL from: {ptfile_url}") response = requests.get( - ptfile_url, - params={"api_key": self.roboflow_api_key}, - timeout=10 + ptfile_url, params={"api_key": self.roboflow_api_key}, timeout=10 ) response.raise_for_status() # Parse response to get weightsUrl data = response.json() - if 'weightsUrl' not in data: + if "weightsUrl" not in data: raise ValueError(f"No weightsUrl in response: {data}") - weights_url = data['weightsUrl'] - logger.info(f"Got weights URL, downloading...") + weights_url = data["weightsUrl"] + logger.info("Got weights URL, downloading...") # Download the .pt file from signed URL response = requests.get(weights_url, stream=True, timeout=120) response.raise_for_status() # Save to cache - with open(cached_model_path, 'wb') as f: + with open(cached_model_path, "wb") as f: for chunk in response.iter_content(chunk_size=8192): f.write(chunk) - logger.info(f"Downloaded model weights: {cached_model_path} ({cached_model_path.stat().st_size} bytes)") + logger.info( + f"Downloaded model weights: {cached_model_path} ({cached_model_path.stat().st_size} bytes)" + ) # Load with Ultralytics YOLO model = YOLO(str(cached_model_path)) - logger.info(f"Successfully loaded Roboflow model for local inference: {model_id}") + logger.info( + f"Successfully loaded Roboflow model for local inference: {model_id}" + ) return model except requests.exceptions.HTTPError as e: @@ -311,12 +339,11 @@ def list_roboflow_project_models(self, project_id: str) -> List[dict]: List of dicts with keys: version, name, images, map """ import requests - from datetime import datetime if not self.roboflow_api_key: raise ValueError("ROBOFLOW_API_KEY not set") - parts = project_id.split('/') + parts = project_id.split("/") if len(parts) != 2: raise ValueError("Project ID must be in format: workspace/project") @@ -327,45 +354,52 @@ def list_roboflow_project_models(self, project_id: str) -> List[dict]: logger.info(f"Querying Roboflow project models: {url}") response = requests.get( - url, - params={"api_key": self.roboflow_api_key}, - timeout=10 + url, params={"api_key": self.roboflow_api_key}, timeout=10 ) response.raise_for_status() data = response.json() versions = [] - if 'versions' in data: - for vd in data['versions']: - version_num = vd.get('id', '') - if isinstance(version_num, str) and '/' in version_num: - version_num = version_num.split('/')[-1] + if "versions" in data: + for vd in data["versions"]: + version_num = vd.get("id", "") + if isinstance(version_num, str) and "/" in version_num: + version_num = version_num.split("/")[-1] if not version_num: - version_num = vd.get('version') + version_num = vd.get("version") if not version_num: continue - map_val = vd.get('model', {}).get('map', '') - if map_val and str(map_val) != 'NaN': + map_val = vd.get("model", {}).get("map", "") + if map_val and str(map_val) != "NaN": map_str = f"{float(map_val):.1f}%" else: map_str = "" - versions.append({ - "version": str(version_num), - "name": vd.get('name', ''), - "images": vd.get('images', 0), - "map": map_str, - "raw": vd, - }) - - versions.sort(key=lambda x: int(x['version']) if x['version'].isdigit() else 0, reverse=True) - logger.info(f"Found {len(versions)} versions: {[v['version'] for v in versions]}") + versions.append( + { + "version": str(version_num), + "name": vd.get("name", ""), + "images": vd.get("images", 0), + "map": map_str, + "raw": vd, + } + ) + + versions.sort( + key=lambda x: int(x["version"]) if x["version"].isdigit() else 0, + reverse=True, + ) + logger.info( + f"Found {len(versions)} versions: {[v['version'] for v in versions]}" + ) return versions except requests.exceptions.HTTPError as e: - error_msg = f"Failed to query Roboflow project: HTTP {e.response.status_code}" + error_msg = ( + f"Failed to query Roboflow project: HTTP {e.response.status_code}" + ) logger.error(error_msg) raise ValueError(error_msg) from e except Exception as e: diff --git a/collab_env/tracking_studio/video_converter.py b/collab_env/tracking_studio/video_converter.py index aaaba55e..ed1acfb1 100644 --- a/collab_env/tracking_studio/video_converter.py +++ b/collab_env/tracking_studio/video_converter.py @@ -42,7 +42,9 @@ def needs_conversion(video_path: Path) -> bool: # H.264 in non-MP4 container (e.g. .mov) may not play in all browsers ext = Path(video_path).suffix.lower() if ext not in (".mp4", ".m4v"): - logger.info(f"H.264 in {ext} container — will remux to .mp4 for browser compatibility") + logger.info( + f"H.264 in {ext} container — will remux to .mp4 for browser compatibility" + ) return True return False @@ -55,7 +57,9 @@ def needs_conversion(video_path: Path) -> bool: raise -def convert_to_h264(input_path: Path, output_path: Path, remux_only: bool = False) -> Path: +def convert_to_h264( + input_path: Path, output_path: Path, remux_only: bool = False +) -> Path: """ Convert video to H.264 format using ffmpeg. @@ -72,9 +76,12 @@ def convert_to_h264(input_path: Path, output_path: Path, remux_only: bool = Fals logger.info(f"Remuxing {input_path} to MP4 container (no re-encoding)") cmd = [ "ffmpeg", - "-i", str(input_path), - "-c", "copy", # Copy all streams without re-encoding - "-movflags", "+faststart", + "-i", + str(input_path), + "-c", + "copy", # Copy all streams without re-encoding + "-movflags", + "+faststart", "-y", str(output_path), ] @@ -82,13 +89,20 @@ def convert_to_h264(input_path: Path, output_path: Path, remux_only: bool = Fals logger.info(f"Converting {input_path} to H.264 format") cmd = [ "ffmpeg", - "-i", str(input_path), - "-c:v", "libx264", - "-preset", "fast", - "-crf", "23", - "-c:a", "aac", - "-b:a", "128k", - "-movflags", "+faststart", + "-i", + str(input_path), + "-c:v", + "libx264", + "-preset", + "fast", + "-crf", + "23", + "-c:a", + "aac", + "-b:a", + "128k", + "-movflags", + "+faststart", "-y", str(output_path), ] diff --git a/collab_env/tracking_studio/video_processor.py b/collab_env/tracking_studio/video_processor.py index e1aa4a78..31d30f3f 100644 --- a/collab_env/tracking_studio/video_processor.py +++ b/collab_env/tracking_studio/video_processor.py @@ -11,8 +11,8 @@ from ultralytics import YOLO import pandas as pd from pathlib import Path -from typing import Callable, Dict, List, Union, Any -import numpy as np +from concurrent.futures import Future +from typing import Callable, Coroutine, Dict, Optional, Union, Any import tempfile import yaml from loguru import logger @@ -28,10 +28,10 @@ def __init__( confidence: float = 0.5, detection_only: bool = False, display_interval: int = 10, - frame_callback: Callable[[np.ndarray, int, int], None] = None, - stop_event: threading.Event = None, - pause_event: threading.Event = None, - skip_frames_event: Dict = None, + frame_callback: Optional[Callable[..., Coroutine[Any, Any, None]]] = None, + stop_event: Optional[threading.Event] = None, + pause_event: Optional[threading.Event] = None, + skip_frames_event: Optional[Dict] = None, ): """ Initialize video tracker. @@ -55,7 +55,7 @@ def __init__( self.stop_event = stop_event or threading.Event() self.pause_event = pause_event or threading.Event() self.skip_frames_event = skip_frames_event or {"skip_amount": 0} - self._pending_update = None # Track in-flight UI update + self._pending_update: Optional[Future[None]] = None # Store tracker config for use with model.track() self.tracker_config = tracker_config @@ -69,9 +69,13 @@ def __init__( self.tracker_yaml_path = None elif not self.use_native_tracking: # For Roboflow inference models (fallback), initialize supervision tracker - logger.info("Using supervision ByteTrack (Roboflow inference model fallback)") + logger.info( + "Using supervision ByteTrack (Roboflow inference model fallback)" + ) self.tracker = sv.ByteTrack( - track_activation_threshold=tracker_config.get("track_high_thresh", 0.25), + track_activation_threshold=tracker_config.get( + "track_high_thresh", 0.25 + ), lost_track_buffer=tracker_config.get("track_buffer", 30), minimum_matching_threshold=tracker_config.get("match_thresh", 0.8), minimum_consecutive_frames=1, @@ -85,7 +89,9 @@ def __init__( self.tracker_yaml_path = self._create_bytetrack_config(tracker_config) # Fast-forward: Skip frames for faster preview - self.skip_frames = tracker_config.get("skip_frames", 1) # 1 = process every frame + self.skip_frames = tracker_config.get( + "skip_frames", 1 + ) # 1 = process every frame logger.info( f"VideoTracker initialized (confidence: {self.confidence}, native_tracking: {self.use_native_tracking})" @@ -114,10 +120,7 @@ def _create_bytetrack_config(self, config: Dict) -> str: # Create temporary YAML file temp_file = tempfile.NamedTemporaryFile( - mode='w', - suffix='.yaml', - delete=False, - prefix='bytetrack_' + mode="w", suffix=".yaml", delete=False, prefix="bytetrack_" ) with temp_file as f: @@ -154,9 +157,7 @@ def _process_video_sync( width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) - logger.info( - f"Video info: {total_frames} frames, {fps} fps, {width}x{height}" - ) + logger.info(f"Video info: {total_frames} frames, {fps} fps, {width}x{height}") detections_list = [] tracking_list = [] @@ -174,6 +175,7 @@ def _process_video_sync( # Check if pause was requested while self.pause_event.is_set(): import time + time.sleep(0.1) # Wait while paused if self.stop_event.is_set(): break @@ -238,14 +240,21 @@ def _process_video_sync( logger.debug(f"Frame {frame_idx}: {len(detections)} detections") # Update tracker (adds track IDs via supervision ByteTrack) + assert self.tracker is not None tracked_detections = self.tracker.update_with_detections(detections) except Exception as e: - logger.error(f"Detection/tracking failed on frame {frame_idx}: {e}", exc_info=True) + logger.error( + f"Detection/tracking failed on frame {frame_idx}: {e}", + exc_info=True, + ) detections = sv.Detections.empty() tracked_detections = sv.Detections.empty() # 2. Save detections + if detections.confidence is None or detections.class_id is None: + frame_idx += 1 + continue for i, (bbox, conf, class_id) in enumerate( zip(detections.xyxy, detections.confidence, detections.class_id) ): @@ -262,12 +271,16 @@ def _process_video_sync( ) # 3. Save tracking data (with track IDs if tracking is enabled) - if not self.detection_only and tracked_detections.tracker_id is not None and len(tracked_detections) > 0: + if ( + not self.detection_only + and tracked_detections.tracker_id is not None + and len(tracked_detections) > 0 + ): for bbox, track_id, conf, class_id in zip( tracked_detections.xyxy, tracked_detections.tracker_id, - tracked_detections.confidence, - tracked_detections.class_id, + tracked_detections.confidence or [], + tracked_detections.class_id or [], ): tracking_list.append( { @@ -283,17 +296,17 @@ def _process_video_sync( ) # 4. Send frame + detections to UI for display - is_last = (frame_idx >= total_frames - 1) + is_last = frame_idx >= total_frames - 1 should_display = (frame_idx % self.display_interval == 0) or is_last - if should_display and self.frame_callback and event_loop: + if should_display and self.frame_callback is not None and event_loop: # Skip if previous UI update is still in-flight (prevents queue buildup) if self._pending_update is None or self._pending_update.done(): self._pending_update = asyncio.run_coroutine_threadsafe( self.frame_callback( frame, tracked_detections, frame_idx, total_frames ), - event_loop + event_loop, ) # Increment frame counter for next iteration @@ -305,12 +318,17 @@ def _process_video_sync( if self.tracker_yaml_path: try: import os + os.unlink(self.tracker_yaml_path) - logger.debug(f"Cleaned up temporary tracker config: {self.tracker_yaml_path}") + logger.debug( + f"Cleaned up temporary tracker config: {self.tracker_yaml_path}" + ) except Exception as e: logger.warning(f"Failed to cleanup tracker config: {e}") - unique_tracks = len(set(t['track_id'] for t in tracking_list)) if tracking_list else 0 + unique_tracks = ( + len(set(t["track_id"] for t in tracking_list)) if tracking_list else 0 + ) logger.info( f"Processing complete: {total_frames} frames, " f"{len(detections_list)} detections, " @@ -325,7 +343,9 @@ def _process_video_sync( output_csv = output_path / "detections.csv" if len(detections_list) > 0: det_df = pd.DataFrame(detections_list) - det_df = det_df[["frame", "x1", "y1", "x2", "y2", "confidence", "class"]] + det_df = det_df[ + ["frame", "x1", "y1", "x2", "y2", "confidence", "class"] + ] det_df.to_csv(output_csv, index=False) else: pd.DataFrame( @@ -343,7 +363,16 @@ def _process_video_sync( tracking_df.to_csv(output_csv, index=False) else: pd.DataFrame( - columns=["track_id", "frame", "x1", "y1", "x2", "y2", "confidence", "class"] + columns=[ + "track_id", + "frame", + "x1", + "y1", + "x2", + "y2", + "confidence", + "class", + ] ).to_csv(output_csv, index=False) logger.info(f"Saved tracking CSV to {output_csv}") @@ -378,7 +407,9 @@ async def process_video_realtime( loop = asyncio.get_running_loop() # Run processing in background thread - logger.info(f"Starting video processing in background thread (frame {start_frame})...") + logger.info( + f"Starting video processing in background thread (frame {start_frame})..." + ) result = await asyncio.to_thread( self._process_video_sync, video_path, output_dir, loop, start_frame ) diff --git a/pyproject.toml b/pyproject.toml index 0751b6db..866e6e8a 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -113,6 +113,9 @@ dev = [ requires = ["setuptools<70.0"] build-backend = "setuptools.build_meta" +[tool.ruff] +target-version = "py39" + [tool.flake8] max-line-length = 120 extend-ignore = ["E203", "E501"] From 54eb273697c022d0fd9cc5f30cac9e02d9bd313e Mon Sep 17 00:00:00 2001 From: Dmitry Batenkov Date: Mon, 23 Feb 2026 10:22:54 -0500 Subject: [PATCH 14/21] lint --- collab_env/tracking/thermal_processing.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/collab_env/tracking/thermal_processing.py b/collab_env/tracking/thermal_processing.py index acf86712..21b75d31 100644 --- a/collab_env/tracking/thermal_processing.py +++ b/collab_env/tracking/thermal_processing.py @@ -243,7 +243,7 @@ def _build_colorbar(self) -> np.ndarray: gradient = np.linspace(1.0, 0.0, self.frame_height, dtype=np.float32).reshape( self.frame_height, 1 ) - gradient = np.repeat(gradient, bar_width, axis=1) # type: ignore[assignment] + gradient = np.repeat(gradient, bar_width, axis=1) rgba = self.cmap(gradient) colorbar = np.clip(rgba[..., :3] * 255.0, 0, 255).astype(np.uint8) @@ -272,7 +272,7 @@ def render(self, frame: np.ndarray) -> np.ndarray: if self.vmax == self.vmin: normalized = np.zeros_like(frame, dtype=np.float32) else: - normalized = (frame - self.vmin) / (self.vmax - self.vmin) # type: ignore[assignment] + normalized = (frame - self.vmin) / (self.vmax - self.vmin) normalized = np.clip(normalized, 0.0, 1.0) rgba = self.cmap(normalized) rgb = np.clip(rgba[..., :3] * 255.0, 0, 255).astype(np.uint8) From c7ab2fdfb6763da3d078056cc523067864b6dfd2 Mon Sep 17 00:00:00 2001 From: Dmitry Batenkov Date: Mon, 23 Feb 2026 10:27:44 -0500 Subject: [PATCH 15/21] more linting --- collab_env/tracking_studio/app.py | 16 ++++++++-------- pyproject.toml | 2 +- 2 files changed, 9 insertions(+), 9 deletions(-) diff --git a/collab_env/tracking_studio/app.py b/collab_env/tracking_studio/app.py index b739d872..ab910e84 100644 --- a/collab_env/tracking_studio/app.py +++ b/collab_env/tracking_studio/app.py @@ -467,10 +467,10 @@ async def _restore_rf_version(): # Update label on change param_slider.on( "update:model-value", - lambda e, - lbl=param_label, - name=param_name: lbl.set_text( - f"{name.replace('_', ' ').title()}: {e.args:.2f}" + lambda e, lbl=param_label, name=param_name: ( + lbl.set_text( + f"{name.replace('_', ' ').title()}: {e.args:.2f}" + ) ), ) param_widgets[param_name] = param_slider @@ -496,10 +496,10 @@ async def _restore_rf_version(): param_slider.on( "update:model-value", - lambda e, - lbl=param_label, - name=param_name: lbl.set_text( - f"{name.replace('_', ' ').title()}: {int(e.args)}" + lambda e, lbl=param_label, name=param_name: ( + lbl.set_text( + f"{name.replace('_', ' ').title()}: {int(e.args)}" + ) ), ) param_widgets[param_name] = param_slider diff --git a/pyproject.toml b/pyproject.toml index f3c3c201..7938099f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -100,7 +100,7 @@ dev = [ "pytest-benchmark", "mypy==1.18.2", "mypy-extensions==1.1.0", - "ruff", + "ruff==0.15.2", "nbval", "nbqa", "black", From d96a54a84a9bce399f63001aa9532e6a9aec436a Mon Sep 17 00:00:00 2001 From: Dmitry Batenkov Date: Mon, 23 Feb 2026 11:22:10 -0500 Subject: [PATCH 16/21] more lint --- .gitignore | 26 +- collab_env/data/gcs_utils.py | 8 +- collab_env/tracking_studio/video_processor.py | 14 +- docs/data/db/basic_db.ipynb | 1233 +++++++++++++++++ docs/data/extended_properties.ipynb | 506 +++++++ scripts/lint.sh | 4 +- 6 files changed, 1769 insertions(+), 22 deletions(-) create mode 100644 docs/data/db/basic_db.ipynb create mode 100644 docs/data/extended_properties.ipynb diff --git a/.gitignore b/.gitignore index e09a4a89..7b610c48 100644 --- a/.gitignore +++ b/.gitignore @@ -91,19 +91,17 @@ venv/ .claude/ # Automatically generated files -docs/source/lightning_logs -docs/preconvert -docs/build -site/ -out/ +lightning_logs +/docs/preconvert +/docs/build *.matrix.gz -docs/source/.ipynb_checkpoints/ -config-local +/docs/source/.ipynb_checkpoints/ +/config-local *preview*.htm* -simulated_data/ -trained_models/ -logs/ -results/ -sim-output/ -good-runs/ -./data/ +/simulated_data/ +/trained_models/ +/logs/ +/results/ +/sim-output/ +/good-runs/ +/data/ diff --git a/collab_env/data/gcs_utils.py b/collab_env/data/gcs_utils.py index e2818c16..396d5660 100644 --- a/collab_env/data/gcs_utils.py +++ b/collab_env/data/gcs_utils.py @@ -26,9 +26,7 @@ def __init__( the default path, then fall back to Application Default Credentials. """ if credentials_path is None: - default_path = expand_path( - DEFAULT_GCS_CREDENTIALS_PATH, get_project_root() - ) + default_path = expand_path(DEFAULT_GCS_CREDENTIALS_PATH, get_project_root()) if os.path.exists(default_path): credentials_path = default_path @@ -48,7 +46,9 @@ def __init__( self.project_id, credentials=self.credentials ) else: - logger.info("Using Application Default Credentials (no credentials file found)") + logger.info( + "Using Application Default Credentials (no credentials file found)" + ) self.credentials_path = None self.credentials = None self._gcs = gcsfs.GCSFileSystem( diff --git a/collab_env/tracking_studio/video_processor.py b/collab_env/tracking_studio/video_processor.py index 31d30f3f..3df6eef0 100644 --- a/collab_env/tracking_studio/video_processor.py +++ b/collab_env/tracking_studio/video_processor.py @@ -276,11 +276,21 @@ def _process_video_sync( and tracked_detections.tracker_id is not None and len(tracked_detections) > 0 ): + confidences: Any = ( + tracked_detections.confidence + if tracked_detections.confidence is not None + else [] + ) + class_ids: Any = ( + tracked_detections.class_id + if tracked_detections.class_id is not None + else [] + ) for bbox, track_id, conf, class_id in zip( tracked_detections.xyxy, tracked_detections.tracker_id, - tracked_detections.confidence or [], - tracked_detections.class_id or [], + confidences, + class_ids, ): tracking_list.append( { diff --git a/docs/data/db/basic_db.ipynb b/docs/data/db/basic_db.ipynb new file mode 100644 index 00000000..cf413407 --- /dev/null +++ b/docs/data/db/basic_db.ipynb @@ -0,0 +1,1233 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Database Query & Insert Tutorial\n", + "\n", + "**Concise guide to querying and inserting data into the tracking analytics database.**\n", + "\n", + "This notebook covers:\n", + "- Connecting to DuckDB and PostgreSQL\n", + "- Querying sessions, episodes, and observations\n", + "- Querying extended properties\n", + "- Inserting data (DuckDB only)\n", + "- Using the high-level QueryBackend API" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "from pathlib import Path\n", + "\n", + "import numpy as np\n", + "import pandas as pd\n", + "\n", + "from collab_env.data.db.config import DBConfig, get_db_config\n", + "from collab_env.data.db.db_loader import Boids3DLoader, DatabaseConnection\n", + "from collab_env.data.db.query_backend import QueryBackend" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 1: Connecting to Databases\n", + "\n", + "The system supports both **DuckDB** (local file) and **PostgreSQL** (server)." + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "2025-11-13 17:02:58 | INFO | collab_env.data.db.db_loader:connect:113 - Connected to DuckDB: /tmp/test_tutorial.duckdb\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "✓ Connected to DuckDB: /tmp/test_tutorial.duckdb\n" + ] + } + ], + "source": [ + "# Connect to DuckDB (local file - for testing and insertions)\n", + "duckdb_path = \"/tmp/test_tutorial.duckdb\"\n", + "\n", + "# Set environment variable before creating config\n", + "os.environ[\"DUCKDB_PATH\"] = duckdb_path\n", + "\n", + "duckdb_config = get_db_config(backend=\"duckdb\")\n", + "\n", + "db_duckdb = DatabaseConnection(duckdb_config)\n", + "db_duckdb.connect()\n", + "print(f\"✓ Connected to DuckDB: {duckdb_path}\")" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "2025-11-13 17:03:02 | INFO | collab_env.data.db.db_loader:connect:111 - Connected to PostgreSQL: tracking_analytics\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "✓ Connected to PostgreSQL: tracking_analytics\n" + ] + } + ], + "source": [ + "# Connect to PostgreSQL (if available - for production queries)\n", + "# Skip this cell if you don't have PostgreSQL running\n", + "try:\n", + " # Set environment variables for PostgreSQL\n", + " os.environ[\"DB_BACKEND\"] = \"postgres\"\n", + " os.environ[\"POSTGRES_DB\"] = \"tracking_analytics\"\n", + " os.environ[\"POSTGRES_USER\"] = \"postgres\" # TODO: change to your username\n", + " os.environ[\"POSTGRES_PASSWORD\"] = \"password\" # TODO: change to your password\n", + "\n", + " postgres_config = get_db_config(backend=\"postgres\")\n", + "\n", + " db_postgres = DatabaseConnection(postgres_config)\n", + " db_postgres.connect()\n", + " print(\"✓ Connected to PostgreSQL: tracking_analytics\")\n", + "except Exception as e:\n", + " print(f\"⚠ PostgreSQL not available: {e}\")\n", + " db_postgres = None" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 2: Initializing Test Database\n", + "\n", + "Create tables and seed data in DuckDB for testing." + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "name": "stderr", + "output_type": "stream", + "text": [ + "\u001b[32m2025-11-13 17:03:08\u001b[0m | \u001b[32m\u001b[1mSUCCESS \u001b[0m | \u001b[32m\u001b[1mConnected to DuckDB: /tmp/test_tutorial.duckdb\u001b[0m\n", + "\u001b[32m2025-11-13 17:03:08\u001b[0m | \u001b[32m\u001b[1mSUCCESS \u001b[0m | \u001b[32m\u001b[1mExecuted 01_core_tables.sql\u001b[0m\n", + "\u001b[32m2025-11-13 17:03:08\u001b[0m | \u001b[32m\u001b[1mSUCCESS \u001b[0m | \u001b[32m\u001b[1mExecuted 02_extended_properties.sql\u001b[0m\n", + "\u001b[32m2025-11-13 17:03:08\u001b[0m | \u001b[32m\u001b[1mSUCCESS \u001b[0m | \u001b[32m\u001b[1mExecuted 03_seed_data.sql\u001b[0m\n" + ] + }, + { + "name": "stdout", + "output_type": "stream", + "text": [ + "Executing 01_core_tables.sql...\n", + "Executing 02_extended_properties.sql...\n", + "Executing 03_seed_data.sql...\n", + "✓ DuckDB schema initialized\n" + ] + } + ], + "source": [ + "# Initialize DuckDB with schema\n", + "from collab_env.data.db.init_database import DatabaseBackend, get_schema_files\n", + "from collab_env.data.file_utils import get_project_root\n", + "\n", + "# Get schema files\n", + "project_root = get_project_root()\n", + "schema_dir = project_root / \"schema\"\n", + "schema_files = get_schema_files(schema_dir)\n", + "\n", + "# Create backend and execute schema\n", + "backend = DatabaseBackend(duckdb_config)\n", + "backend.connect()\n", + "\n", + "for schema_file in schema_files:\n", + " print(f\"Executing {schema_file.name}...\")\n", + " backend.execute_file(schema_file)\n", + "\n", + "backend.close()\n", + "print(\"✓ DuckDB schema initialized\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 3: Inserting Data (DuckDB Only)\n", + "\n", + "Insert sample session, episode, and observations data." + ] + }, + { + "cell_type": "code", + "execution_count": 5, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "✓ Inserted session\n" + ] + } + ], + "source": [ + "# Insert a test session\n", + "import json\n", + "\n", + "session_data = {\n", + " \"session_id\": \"test-session-001\",\n", + " \"session_name\": \"Tutorial Example Session\",\n", + " \"category_id\": \"boids_3d\",\n", + " \"config\": json.dumps({\"num_agents\": 10, \"scene_size\": 480}),\n", + " \"metadata\": json.dumps({\"notes\": \"Created in tutorial notebook\"}),\n", + "}\n", + "\n", + "db_duckdb.execute(\n", + " \"\"\"\n", + " INSERT INTO sessions (session_id, session_name, category_id, config, metadata)\n", + " VALUES (:session_id, :session_name, :category_id, :config, :metadata)\n", + " \"\"\",\n", + " session_data,\n", + ")\n", + "print(\"✓ Inserted session\")" + ] + }, + { + "cell_type": "code", + "execution_count": 6, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "✓ Inserted episode\n" + ] + } + ], + "source": [ + "# Insert a test episode\n", + "episode_data = {\n", + " \"episode_id\": \"test-episode-001\",\n", + " \"session_id\": \"test-session-001\",\n", + " \"episode_number\": 0,\n", + " \"num_frames\": 100,\n", + " \"num_agents\": 10,\n", + " \"frame_rate\": 30.0,\n", + " \"file_path\": \"/tmp/test_episode.parquet\",\n", + "}\n", + "\n", + "db_duckdb.execute(\n", + " \"\"\"\n", + " INSERT INTO episodes (episode_id, session_id, episode_number, num_frames, num_agents, frame_rate, file_path)\n", + " VALUES (:episode_id, :session_id, :episode_number, :num_frames, :num_agents, :frame_rate, :file_path)\n", + " \"\"\",\n", + " episode_data,\n", + ")\n", + "print(\"✓ Inserted episode\")" + ] + }, + { + "cell_type": "code", + "execution_count": 7, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "✓ Inserted 1000 observations\n" + ] + } + ], + "source": [ + "# Insert test observations using pandas (bulk insert)\n", + "num_agents = 10\n", + "num_frames = 100\n", + "\n", + "# Generate synthetic trajectory data\n", + "observations = []\n", + "for time_idx in range(num_frames):\n", + " for agent_id in range(num_agents):\n", + " # Simple circular motion\n", + " angle = 2 * np.pi * time_idx / num_frames + agent_id * 0.2\n", + " radius = 100 + agent_id * 10\n", + "\n", + " x = 240 + radius * np.cos(angle)\n", + " y = 240 + radius * np.sin(angle)\n", + " z = 50 + 20 * np.sin(angle * 2)\n", + "\n", + " v_x = -radius * np.sin(angle) * 2 * np.pi / num_frames\n", + " v_y = radius * np.cos(angle) * 2 * np.pi / num_frames\n", + " v_z = 40 * np.cos(angle * 2) * 2 * np.pi / num_frames\n", + "\n", + " observations.append(\n", + " {\n", + " \"episode_id\": \"test-episode-001\",\n", + " \"time_index\": time_idx,\n", + " \"agent_id\": agent_id,\n", + " \"agent_type_id\": \"agent\",\n", + " \"x\": x,\n", + " \"y\": y,\n", + " \"z\": z,\n", + " \"v_x\": v_x,\n", + " \"v_y\": v_y,\n", + " \"v_z\": v_z,\n", + " }\n", + " )\n", + "\n", + "obs_df = pd.DataFrame(observations)\n", + "db_duckdb.insert_dataframe(obs_df, \"observations\", if_exists=\"append\")\n", + "print(f\"✓ Inserted {len(obs_df)} observations\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Insert extended properties (distance to target)\n", + "# First, get observation IDs\n", + "obs_ids = db_duckdb.fetch_all(\n", + " \"\"\"\n", + " SELECT observation_id, time_index, agent_id\n", + " FROM observations\n", + " WHERE episode_id = :episode_id\n", + " ORDER BY time_index, agent_id\n", + " \"\"\",\n", + " {\"episode_id\": \"test-episode-001\"},\n", + ")\n", + "\n", + "# Compute synthetic distance to target center\n", + "target_center = np.array([240, 240, 50])\n", + "extended_props = []\n", + "\n", + "for obs_id, time_idx, agent_id in obs_ids:\n", + " # Get position from observations\n", + " obs = obs_df[\n", + " (obs_df[\"time_index\"] == time_idx) & (obs_df[\"agent_id\"] == agent_id)\n", + " ].iloc[0]\n", + " pos = np.array([obs[\"x\"], obs[\"y\"], obs[\"z\"]])\n", + " distance = np.linalg.norm(pos - target_center)\n", + "\n", + " extended_props.append(\n", + " {\n", + " \"observation_id\": obs_id,\n", + " \"property_id\": \"distance_to_target_center\",\n", + " \"value_float\": distance,\n", + " \"value_text\": None,\n", + " }\n", + " )\n", + "\n", + "ext_df = pd.DataFrame(extended_props)\n", + "db_duckdb.insert_dataframe(ext_df, \"extended_properties\", if_exists=\"append\")\n", + "print(f\"✓ Inserted {len(ext_df)} extended properties\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 4: Basic Queries\n", + "\n", + "Query the database using low-level SQL." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Query 1: List all sessions\n", + "sessions = db_duckdb.fetch_all(\"SELECT * FROM sessions\")\n", + "print(\"Sessions:\")\n", + "for session in sessions:\n", + " print(f\" - {session[0]}: {session[1]} ({session[2]})\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Query 2: Get episodes for a session\n", + "episodes = db_duckdb.fetch_all(\n", + " \"\"\"\n", + " SELECT episode_id, episode_number, num_frames, num_agents, frame_rate\n", + " FROM episodes\n", + " WHERE session_id = :session_id\n", + " ORDER BY episode_number\n", + " \"\"\",\n", + " {\"session_id\": \"test-session-001\"},\n", + ")\n", + "\n", + "print(\"\\nEpisodes for test-session-001:\")\n", + "for ep in episodes:\n", + " print(f\" - {ep[0]}: {ep[2]} frames, {ep[3]} agents @ {ep[4]} fps\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Query 3: Get observations with computed speed\n", + "obs_query = \"\"\"\n", + "SELECT \n", + " time_index,\n", + " agent_id,\n", + " x, y, z,\n", + " v_x, v_y, v_z,\n", + " sqrt(v_x*v_x + v_y*v_y + v_z*v_z) as speed\n", + "FROM observations\n", + "WHERE episode_id = :episode_id\n", + " AND time_index < 5\n", + "ORDER BY time_index, agent_id\n", + "\"\"\"\n", + "\n", + "from sqlalchemy import text\n", + "\n", + "with db_duckdb.engine.connect() as conn:\n", + " result = conn.execute(text(obs_query), {\"episode_id\": \"test-episode-001\"})\n", + " obs_df_query = pd.DataFrame(result.fetchall(), columns=result.keys())\n", + "\n", + "print(\"\\nFirst 5 frames of observations:\")\n", + "print(\n", + " obs_df_query[[\"time_index\", \"agent_id\", \"x\", \"y\", \"z\", \"speed\"]].to_string(\n", + " index=False\n", + " )\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Query 4: Get observations with extended properties\n", + "extended_query = \"\"\"\n", + "SELECT \n", + " o.time_index,\n", + " o.agent_id,\n", + " o.x, o.y, o.z,\n", + " pd.property_name,\n", + " ep.value_float\n", + "FROM observations o\n", + "JOIN extended_properties ep ON o.observation_id = ep.observation_id\n", + "JOIN property_definitions pd ON ep.property_id = pd.property_id\n", + "WHERE o.episode_id = :episode_id\n", + " AND o.time_index < 5\n", + "ORDER BY o.time_index, o.agent_id\n", + "\"\"\"\n", + "\n", + "with db_duckdb.engine.connect() as conn:\n", + " result = conn.execute(text(extended_query), {\"episode_id\": \"test-episode-001\"})\n", + " ext_query_df = pd.DataFrame(result.fetchall(), columns=result.keys())\n", + "\n", + "print(\"\\nObservations with extended properties (first 5 frames):\")\n", + "print(ext_query_df.to_string(index=False))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Query 5: Aggregate statistics\n", + "stats_query = \"\"\"\n", + "SELECT \n", + " COUNT(*) as total_observations,\n", + " COUNT(DISTINCT agent_id) as num_agents,\n", + " COUNT(DISTINCT time_index) as num_frames,\n", + " AVG(sqrt(v_x*v_x + v_y*v_y + v_z*v_z)) as avg_speed,\n", + " MAX(sqrt(v_x*v_x + v_y*v_y + v_z*v_z)) as max_speed\n", + "FROM observations\n", + "WHERE episode_id = :episode_id\n", + "\"\"\"\n", + "\n", + "stats = db_duckdb.fetch_one(stats_query, {\"episode_id\": \"test-episode-001\"})\n", + "print(\"\\nEpisode Statistics:\")\n", + "print(f\" Total observations: {stats[0]}\")\n", + "print(f\" Num agents: {stats[1]}\")\n", + "print(f\" Num frames: {stats[2]}\")\n", + "print(f\" Avg speed: {stats[3]:.2f}\")\n", + "print(f\" Max speed: {stats[4]:.2f}\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 5: QueryBackend API (Dashboard Pattern)\n", + "\n", + "The **QueryBackend** provides high-level methods for common queries. This is how the dashboard uses it." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Initialize QueryBackend with DuckDB\n", + "query = QueryBackend(config=duckdb_config)\n", + "print(\"✓ QueryBackend initialized\")" + ] + }, + { + "cell_type": "markdown", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "### 5.1 Session and Episode Discovery" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Get all categories\n", + "categories = query.get_categories()\n", + "print(\"Categories:\")\n", + "print(categories[[\"category_id\", \"category_name\"]].to_string(index=False))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Get sessions by category\n", + "sessions = query.get_sessions(category_id=\"boids_3d\")\n", + "print(\"\\nBoids 3D Sessions:\")\n", + "print(sessions[[\"session_id\", \"session_name\", \"category_id\"]].to_string(index=False))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Get episodes for a session\n", + "episodes = query.get_episodes(\"test-session-001\")\n", + "print(\"\\nEpisodes:\")\n", + "print(\n", + " episodes[[\"episode_id\", \"num_frames\", \"num_agents\", \"frame_rate\"]].to_string(\n", + " index=False\n", + " )\n", + ")" + ] + }, + { + "cell_type": "markdown", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "### 5.2 Spatial Analysis" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Get spatial heatmap (binned positions)\n", + "heatmap = query.get_spatial_heatmap(\n", + " episode_id=\"test-episode-001\", bin_size=50.0, agent_type=\"agent\"\n", + ")\n", + "print(\"\\nSpatial Heatmap (top 10 bins by density):\")\n", + "top_bins = heatmap.nlargest(10, \"density\")[[\"x_bin\", \"y_bin\", \"z_bin\", \"density\"]]\n", + "print(top_bins.to_string(index=False))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Get episode tracks for visualization\n", + "tracks = query.get_episode_tracks(\n", + " episode_id=\"test-episode-001\", start_time=0, end_time=10\n", + ")\n", + "print(\"\\nTracks (first 10 frames):\")\n", + "print(\n", + " tracks[[\"agent_id\", \"time_index\", \"x\", \"y\", \"z\", \"speed\"]]\n", + " .head(20)\n", + " .to_string(index=False)\n", + ")" + ] + }, + { + "cell_type": "markdown", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "### 5.3 Episode Tracks (for visualization)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Get available extended properties\n", + "props = query.get_available_properties(\"test-episode-001\")\n", + "print(\"\\nAvailable Extended Properties:\")\n", + "print(\n", + " props[[\"property_id\", \"property_name\", \"data_type\", \"unit\"]].to_string(index=False)\n", + ")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Get property distributions for histogram\n", + "dist = query.get_property_distributions(\n", + " episode_id=\"test-episode-001\", property_ids=[\"distance_to_target_center\"]\n", + ")\n", + "print(\"\\nDistance to Target Distribution:\")\n", + "print(f\" Count: {len(dist)}\")\n", + "print(f\" Mean: {dist['value_float'].mean():.2f}\")\n", + "print(f\" Std: {dist['value_float'].std():.2f}\")\n", + "print(f\" Min: {dist['value_float'].min():.2f}\")\n", + "print(f\" Max: {dist['value_float'].max():.2f}\")" + ] + }, + { + "cell_type": "markdown", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "### 5.4 Extended Properties" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Get extended properties time series (windowed)\n", + "timeseries = query.get_extended_properties_timeseries(\n", + " episode_id=\"test-episode-001\",\n", + " window_size=20,\n", + " property_ids=[\"distance_to_target_center\"],\n", + ")\n", + "print(\"\\nExtended Properties Time Series (20-frame windows):\")\n", + "print(\n", + " timeseries[[\"time_window\", \"property_id\", \"avg_value\", \"std_value\"]].to_string(\n", + " index=False\n", + " )\n", + ")" + ] + }, + { + "cell_type": "markdown", + "source": [ + "#", + "#", + "#", + " ", + "5", + ".", + "6", + " ", + "D", + "a", + "s", + "h", + "b", + "o", + "a", + "r", + "d", + " ", + "P", + "a", + "t", + "t", + "e", + "r", + "n", + ":", + " ", + "U", + "s", + "i", + "n", + "g", + " ", + "A", + "n", + "a", + "l", + "y", + "s", + "i", + "s", + "C", + "o", + "n", + "t", + "e", + "x", + "t", + "\n", + "\n", + "T", + "h", + "e", + " ", + "d", + "a", + "s", + "h", + "b", + "o", + "a", + "r", + "d", + " ", + "u", + "s", + "e", + "s", + " ", + "`", + "A", + "n", + "a", + "l", + "y", + "s", + "i", + "s", + "C", + "o", + "n", + "t", + "e", + "x", + "t", + "`", + " ", + "t", + "o", + " ", + "s", + "h", + "a", + "r", + "e", + " ", + "q", + "u", + "e", + "r", + "y", + " ", + "p", + "a", + "r", + "a", + "m", + "e", + "t", + "e", + "r", + "s", + " ", + "a", + "c", + "r", + "o", + "s", + "s", + " ", + "w", + "i", + "d", + "g", + "e", + "t", + "s", + ".", + " ", + "T", + "h", + "i", + "s", + " ", + "p", + "a", + "t", + "t", + "e", + "r", + "n", + " ", + "e", + "n", + "a", + "b", + "l", + "e", + "s", + ":", + "\n", + "-", + " ", + "C", + "o", + "n", + "s", + "i", + "s", + "t", + "e", + "n", + "t", + " ", + "p", + "a", + "r", + "a", + "m", + "e", + "t", + "e", + "r", + "s", + " ", + "a", + "c", + "r", + "o", + "s", + "s", + " ", + "m", + "u", + "l", + "t", + "i", + "p", + "l", + "e", + " ", + "a", + "n", + "a", + "l", + "y", + "s", + "e", + "s", + "\n", + "-", + " ", + "E", + "a", + "s", + "y", + " ", + "p", + "a", + "r", + "a", + "m", + "e", + "t", + "e", + "r", + " ", + "o", + "v", + "e", + "r", + "r", + "i", + "d", + "e", + "s", + " ", + "f", + "o", + "r", + " ", + "w", + "i", + "d", + "g", + "e", + "t", + "-", + "s", + "p", + "e", + "c", + "i", + "f", + "i", + "c", + " ", + "c", + "u", + "s", + "t", + "o", + "m", + "i", + "z", + "a", + "t", + "i", + "o", + "n", + "\n", + "-", + " ", + "C", + "e", + "n", + "t", + "r", + "a", + "l", + "i", + "z", + "e", + "d", + " ", + "s", + "c", + "o", + "p", + "e", + " ", + "m", + "a", + "n", + "a", + "g", + "e", + "m", + "e", + "n", + "t", + " ", + "(", + "e", + "p", + "i", + "s", + "o", + "d", + "e", + "/", + "s", + "e", + "s", + "s", + "i", + "o", + "n", + " ", + "l", + "e", + "v", + "e", + "l", + ")" + ], + "metadata": {} + }, + { + "cell_type": "code", + "source": [ + "# Import context classes (dashboard pattern)\n", + "from collab_env.dashboard.widgets import AnalysisContext, QueryScope, ScopeType\n", + "\n", + "# Create a query scope for an episode\n", + "scope = QueryScope(\n", + " scope_type=ScopeType.EPISODE,\n", + " episode_id=\"test-episode-001\",\n", + " session_id=\"test-session-001\",\n", + " start_time=0,\n", + " end_time=50,\n", + " agent_type=\"agent\",\n", + ")\n", + "\n", + "# Create analysis context with shared parameters\n", + "context = AnalysisContext(\n", + " query_backend=query,\n", + " scope=scope,\n", + " spatial_bin_size=20.0, # Shared spatial discretization\n", + " temporal_window_size=10, # Shared time window\n", + " min_samples=10, # Shared minimum sample threshold\n", + " on_loading=lambda msg: print(f\"⏳ {msg}\"),\n", + " on_success=lambda msg: print(f\"✓ {msg}\"),\n", + " on_error=lambda msg: print(f\"✗ {msg}\"),\n", + ")\n", + "\n", + "print(\"✓ Created AnalysisContext\")\n", + "print(f\" Scope: {scope.scope_type.value}\")\n", + "print(f\" Episode: {scope.episode_id}\")\n", + "print(f\" Time range: {scope.start_time}-{scope.end_time}\")\n", + "print(f\" Spatial bin: {context.spatial_bin_size}\")\n", + "print(f\" Time window: {context.temporal_window_size}\")" + ], + "metadata": {}, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Use context to get merged query parameters\n", + "params = context.get_query_params()\n", + "print(\"\\nMerged Query Parameters:\")\n", + "for key, value in params.items():\n", + " print(f\" {key}: {value}\")\n", + "\n", + "# Query using merged parameters (dashboard pattern)\n", + "context.report_loading(\"Loading spatial heatmap...\")\n", + "heatmap = query.get_spatial_heatmap(**params)\n", + "context.report_success(f\"Loaded {len(heatmap)} bins\")\n", + "\n", + "print(\"\\nHeatmap with context parameters:\")\n", + "print(heatmap[[\"x_bin\", \"y_bin\", \"z_bin\", \"density\"]].head().to_string(index=False))" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Override specific parameters (widget-specific customization)\n", + "custom_params = context.get_query_params(\n", + " bin_size=10.0, min_count=5 # Override spatial bin size # Override minimum count\n", + ")\n", + "\n", + "print(\"\\nCustom Parameters (with overrides):\")\n", + "print(f\" bin_size: {custom_params['bin_size']} (was {params['bin_size']})\")\n", + "print(f\" min_count: {custom_params['min_count']} (was {params.get('min_count', 1)})\")\n", + "\n", + "# Query with custom parameters\n", + "custom_heatmap = query.get_spatial_heatmap(**custom_params)\n", + "print(f\"\\nCustom heatmap: {len(custom_heatmap)} bins (finer resolution)\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 6: Querying PostgreSQL (if available)\n", + "\n", + "Same queries work on PostgreSQL with production data." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "if db_postgres is not None:\n", + " # Initialize QueryBackend for PostgreSQL\n", + " postgres_config = get_db_config(backend=\"postgres\")\n", + " query_pg = QueryBackend(config=postgres_config)\n", + "\n", + " # Get sessions\n", + " sessions_pg = query_pg.get_sessions(category_id=\"boids_3d\")\n", + " print(\"PostgreSQL - Boids 3D Sessions:\")\n", + " print(f\" Found {len(sessions_pg)} sessions\")\n", + "\n", + " if len(sessions_pg) > 0:\n", + " # Get first session's episodes\n", + " session_id = sessions_pg.iloc[0][\"session_id\"]\n", + " episodes_pg = query_pg.get_episodes(session_id)\n", + " print(f\"\\n Episodes for {session_id}:\")\n", + " print(f\" Found {len(episodes_pg)} episodes\")\n", + "\n", + " if len(episodes_pg) > 0:\n", + " # Get spatial heatmap\n", + " episode_id = episodes_pg.iloc[0][\"episode_id\"]\n", + " heatmap_pg = query_pg.get_spatial_heatmap(episode_id, bin_size=20.0)\n", + " print(f\"\\n Heatmap for {episode_id}:\")\n", + " print(f\" Generated {len(heatmap_pg)} bins\")\n", + "\n", + " query_pg.close()\n", + "else:\n", + " print(\"PostgreSQL not available - skipping\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Part 7: Loading Real Data\n", + "\n", + "Load actual simulation data using data loaders." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Example: Load 3D boids simulation (if data exists)\n", + "# Uncomment and modify path as needed\n", + "\n", + "# sim_dir = Path(\"simulated_data/hackathon/hackathon-boid-small-200-sim_run-started-20250926-220926\")\n", + "# if sim_dir.exists():\n", + "# loader = Boids3DLoader(db_duckdb, max_episodes=2) # Load only 2 episodes\n", + "# loader.load_simulation(sim_dir)\n", + "# print(\"✓ Loaded simulation data\")\n", + "# else:\n", + "# print(\"Simulation directory not found\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Cleanup" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Close connections\n", + "query.close()\n", + "db_duckdb.close()\n", + "if db_postgres is not None:\n", + " db_postgres.close()\n", + "\n", + "print(\"✓ All connections closed\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Summary\n", + "\n", + "### Key Takeaways\n", + "\n", + "1. **Connecting**:\n", + " - Set environment variables (`DUCKDB_PATH`, `DB_BACKEND`, etc.)\n", + " - Use `get_db_config()` to create config from environment\n", + " - Use `DatabaseConnection` for low-level access\n", + " - Use `QueryBackend()` for high-level queries (no args needed!)\n", + "\n", + "2. **Inserting** (DuckDB only):\n", + " - `db.execute()` for single inserts with named parameters\n", + " - `db.insert_dataframe()` for bulk inserts (much faster)\n", + " - Use transactions for multi-operation consistency\n", + "\n", + "3. **Querying with QueryBackend**:\n", + " - **Discovery**: `get_categories()`, `get_sessions()`, `get_episodes()`\n", + " - **Spatial**: `get_spatial_heatmap()`, `get_episode_tracks()`\n", + " - **Properties**: `get_available_properties()`, `get_property_distributions()`, `get_extended_properties_timeseries()`\n", + " - All methods return pandas DataFrames\n", + " - All methods support optional time/agent filtering\n", + "\n", + "4. **Advanced Pattern (Dashboard)**:\n", + " - Create `QueryScope` to define what data to analyze\n", + " - Create `AnalysisContext` with shared parameters\n", + " - Use `context.get_query_params()` to merge scope + shared + custom params\n", + " - Pass merged params to QueryBackend methods\n", + " - Enables consistent parameters across multiple widgets/analyses\n", + "\n", + "5. **Best Practices**:\n", + " - Use DuckDB for testing and local development\n", + " - Use PostgreSQL for production and Grafana\n", + " - Use QueryBackend for cleaner, higher-level code\n", + " - Use AnalysisContext for multi-widget applications\n", + " - Always filter by episode_id for performance\n", + "\n", + "### Next Steps\n", + "\n", + "- [docs/data/db/README.md](README.md) - Complete database documentation\n", + "- [schema/README.md](../../../schema/README.md) - Database schema details\n", + "- [collab_env/dashboard/widgets/](../../dashboard/widgets/) - Widget examples using AnalysisContext\n", + "- [collab_env/data/db/queries/](../../data/db/queries/) - SQL query library" + ] + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv-310", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.16" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/docs/data/extended_properties.ipynb b/docs/data/extended_properties.ipynb new file mode 100644 index 00000000..766ee9b3 --- /dev/null +++ b/docs/data/extended_properties.ipynb @@ -0,0 +1,506 @@ +{ + "cells": [ + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "# Extended Properties Viewer - Session-Scope Analysis\n", + "\n", + "**Minimal example of querying session-level extended properties and visualizing distributions.**\n", + "\n", + "This notebook demonstrates:\n", + "- Using the high-level QueryBackend API (same as dashboard widgets)\n", + "- Querying session-scope extended properties\n", + "- Visualizing property distributions with seaborn violin plots" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": {}, + "outputs": [], + "source": [ + "import os\n", + "\n", + "import matplotlib.pyplot as plt\n", + "import pandas as pd\n", + "import seaborn as sns\n", + "\n", + "from collab_env.dashboard.widgets.analysis_context import AnalysisContext\n", + "from collab_env.dashboard.widgets.query_scope import QueryScope, ScopeType\n", + "from collab_env.data.db.query_backend import QueryBackend" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Setup: Initialize QueryBackend and Session Scope" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "### Recommended: Connect to Cloud SQL via Auth Proxy\n\n**Google's recommended approach** for connecting to Cloud SQL is using the Cloud SQL Auth Proxy, which provides:\n- ✅ Automatic SSL/TLS encryption\n- ✅ IAM-based authentication\n- ✅ No need to manage authorized networks\n- ✅ Works with both public and private IP\n\n**Setup (one-time):**\n```bash\n# Terminal 1: Start Cloud SQL Auth Proxy\n./cloud-sql-proxy PROJECT_ID:REGION:INSTANCE_NAME --port 5433\n\n# Keep this running while using the notebook\n```\n\n**In the notebook:** Set environment variables to connect via proxy on localhost:\n\n```python\nimport os\nos.environ['DB_BACKEND'] = 'postgres'\nos.environ['POSTGRES_HOST'] = 'localhost'\nos.environ['POSTGRES_PORT'] = '5433' # Proxy port\nos.environ['POSTGRES_DB'] = 'tracking_analytics'\nos.environ['POSTGRES_USER'] = 'postgres'\nos.environ['POSTGRES_PASSWORD'] = 'your-password' # Or fetch from Secret Manager\n```\n\nThen skip to **Cell 8** to use `QueryBackend()` with environment variables.\n\n---\n\n### Alternative (Not Recommended): Direct Public IP Connection\n\n⚠️ **Warning**: Direct connections bypass the proxy's security features. Google recommends using the Auth Proxy instead.\n\nIf you must connect directly (e.g., testing, special circumstances):" + }, + { + "cell_type": "code", + "source": [ + "# RECOMMENDED: Configure environment for Cloud SQL Auth Proxy connection\n", + "# Run this cell if you have the proxy running in another terminal\n", + "\n", + "import os\n", + "import subprocess\n", + "\n", + "# Fetch password from Secret Manager (recommended)\n", + "\n", + "\n", + "def get_password_from_secret_manager(secret_name=\"postgres-password\"):\n", + " \"\"\"Fetch password from Google Cloud Secret Manager.\"\"\"\n", + " try:\n", + " result = subprocess.run(\n", + " [\n", + " \"gcloud\",\n", + " \"secrets\",\n", + " \"versions\",\n", + " \"access\",\n", + " \"latest\",\n", + " \"--secret\",\n", + " secret_name,\n", + " ],\n", + " capture_output=True,\n", + " text=True,\n", + " check=True,\n", + " )\n", + " return result.stdout.strip()\n", + " except subprocess.CalledProcessError as e:\n", + " print(f\"⚠ Failed to fetch secret: {e}\")\n", + " return None\n", + "\n", + "\n", + "# Configure environment variables for proxy connection\n", + "os.environ[\"DB_BACKEND\"] = \"postgres\"\n", + "os.environ[\"POSTGRES_HOST\"] = \"localhost\" # Connect via proxy on localhost\n", + "os.environ[\"POSTGRES_PORT\"] = (\n", + " \"5433\" # Proxy port (use 5433 if local postgres uses 5432)\n", + ")\n", + "os.environ[\"POSTGRES_DB\"] = \"tracking_analytics\"\n", + "os.environ[\"POSTGRES_USER\"] = \"postgres\"\n", + "os.environ[\"POSTGRES_PASSWORD\"] = (\n", + " get_password_from_secret_manager() or \"your-password-here\"\n", + ")\n", + "\n", + "print(\"✓ Environment configured for Cloud SQL Auth Proxy connection\")\n", + "print(f\" Proxy endpoint: localhost:{os.environ['POSTGRES_PORT']}\")\n", + "print(f\" Database: {os.environ['POSTGRES_DB']}\")\n", + "print(f\" User: {os.environ['POSTGRES_USER']}\")\n", + "\n", + "# Now skip to Cell 8 to initialize QueryBackend() with these settings" + ], + "metadata": {}, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": "---\n\n**If using Cloud SQL Auth Proxy (recommended), skip cells 6-9 below and jump directly to Cell 10.**\n\n---", + "metadata": {} + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "✓ Password retrieved successfully\n" + ] + } + ], + "source": [ + "# Option 1: Fetch password from Google Cloud Secret Manager (recommended)\n", + "# Requires: gcloud CLI installed and authenticated\n", + "import subprocess\n", + "\n", + "\n", + "def get_password_from_secret_manager(secret_name=\"postgres-password\"):\n", + " \"\"\"Fetch password from Google Cloud Secret Manager.\"\"\"\n", + " try:\n", + " result = subprocess.run(\n", + " [\n", + " \"gcloud\",\n", + " \"secrets\",\n", + " \"versions\",\n", + " \"access\",\n", + " \"latest\",\n", + " \"--secret\",\n", + " secret_name,\n", + " ],\n", + " capture_output=True,\n", + " text=True,\n", + " check=True,\n", + " )\n", + " return result.stdout.strip()\n", + " except subprocess.CalledProcessError as e:\n", + " print(f\"⚠ Failed to fetch secret: {e}\")\n", + " return None\n", + "\n", + "\n", + "# Option 2: Specify password directly (not recommended for shared notebooks)\n", + "# db_password = 'your-password-here'\n", + "\n", + "\n", + "# Try Secret Manager first, fallback to environment variable\n", + "db_password = get_password_from_secret_manager() or os.getenv(\"POSTGRES_PASSWORD\")\n", + "\n", + "if db_password:\n", + " print(\"✓ Password retrieved successfully\")\n", + "else:\n", + " print(\"⚠ No password found - will fail to connect\")" + ] + }, + { + "cell_type": "code", + "execution_count": 3, + "metadata": {}, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "✓ Database configuration created:\n", + " Host: 34.67.80.127\n", + " Database: tracking_analytics\n", + " User: postgres\n", + " Password set: True\n" + ] + } + ], + "source": [ + "# Construct database configuration for direct Cloud SQL connection\n", + "from collab_env.data.db.config import DBConfig, PostgresConfig\n", + "\n", + "# Cloud SQL instance details (modify with your values)\n", + "CLOUD_SQL_PUBLIC_IP = \"34.67.80.127\" # Get from: gcloud sql instances describe INSTANCE_NAME --format=\"value(ipAddresses[0].ipAddress)\"\n", + "DB_NAME = \"tracking_analytics\"\n", + "DB_USER = \"postgres\"\n", + "DB_PORT = 5432 # Default PostgreSQL port\n", + "\n", + "# Create PostgresConfig with direct connection parameters\n", + "postgres_config = PostgresConfig(\n", + " host=CLOUD_SQL_PUBLIC_IP,\n", + " port=DB_PORT,\n", + " dbname=DB_NAME,\n", + " user=DB_USER,\n", + " password=db_password,\n", + ")\n", + "\n", + "# Create DBConfig with postgres backend\n", + "db_config = DBConfig(backend=\"postgres\")\n", + "db_config.postgres = postgres_config # Override with custom config\n", + "\n", + "print(f\"✓ Database configuration created:\")\n", + "print(f\" Host: {postgres_config.host}\")\n", + "print(f\" Database: {postgres_config.dbname}\")\n", + "print(f\" User: {postgres_config.user}\")\n", + "print(f\" Password set: {postgres_config.password is not None}\")" + ] + }, + { + "cell_type": "code", + "execution_count": 4, + "metadata": {}, + "outputs": [ + { + "ename": "OperationalError", + "evalue": "(psycopg2.OperationalError) connection to server at \"34.67.80.127\", port 5432 failed: Operation timed out\n\tIs the server running on that host and accepting TCP/IP connections?\n\n(Background on this error at: https://sqlalche.me/e/20/e3q8)", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mOperationalError\u001b[0m Traceback (most recent call last)", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/engine/base.py:143\u001b[0m, in \u001b[0;36mConnection.__init__\u001b[0;34m(self, engine, connection, _has_events, _allow_revalidate, _allow_autobegin)\u001b[0m\n\u001b[1;32m 142\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 143\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_dbapi_connection \u001b[38;5;241m=\u001b[39m \u001b[43mengine\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mraw_connection\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 144\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m dialect\u001b[38;5;241m.\u001b[39mloaded_dbapi\u001b[38;5;241m.\u001b[39mError \u001b[38;5;28;01mas\u001b[39;00m err:\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/engine/base.py:3301\u001b[0m, in \u001b[0;36mEngine.raw_connection\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 3280\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Return a \"raw\" DBAPI connection from the connection pool.\u001b[39;00m\n\u001b[1;32m 3281\u001b[0m \n\u001b[1;32m 3282\u001b[0m \u001b[38;5;124;03mThe returned object is a proxied version of the DBAPI\u001b[39;00m\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 3299\u001b[0m \n\u001b[1;32m 3300\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[0;32m-> 3301\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mpool\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconnect\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/base.py:447\u001b[0m, in \u001b[0;36mPool.connect\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 440\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Return a DBAPI connection from the pool.\u001b[39;00m\n\u001b[1;32m 441\u001b[0m \n\u001b[1;32m 442\u001b[0m \u001b[38;5;124;03mThe connection is instrumented such that when its\u001b[39;00m\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 445\u001b[0m \n\u001b[1;32m 446\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[0;32m--> 447\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_ConnectionFairy\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_checkout\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/base.py:1264\u001b[0m, in \u001b[0;36m_ConnectionFairy._checkout\u001b[0;34m(cls, pool, threadconns, fairy)\u001b[0m\n\u001b[1;32m 1263\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m fairy:\n\u001b[0;32m-> 1264\u001b[0m fairy \u001b[38;5;241m=\u001b[39m \u001b[43m_ConnectionRecord\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcheckout\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpool\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1266\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m threadconns \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/base.py:711\u001b[0m, in \u001b[0;36m_ConnectionRecord.checkout\u001b[0;34m(cls, pool)\u001b[0m\n\u001b[1;32m 710\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 711\u001b[0m rec \u001b[38;5;241m=\u001b[39m \u001b[43mpool\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_do_get\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 713\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/impl.py:177\u001b[0m, in \u001b[0;36mQueuePool._do_get\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 176\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m:\n\u001b[0;32m--> 177\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m util\u001b[38;5;241m.\u001b[39msafe_reraise():\n\u001b[1;32m 178\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_dec_overflow()\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/util/langhelpers.py:224\u001b[0m, in \u001b[0;36msafe_reraise.__exit__\u001b[0;34m(self, type_, value, traceback)\u001b[0m\n\u001b[1;32m 223\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_exc_info \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;66;03m# remove potential circular references\u001b[39;00m\n\u001b[0;32m--> 224\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m exc_value\u001b[38;5;241m.\u001b[39mwith_traceback(exc_tb)\n\u001b[1;32m 225\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/impl.py:175\u001b[0m, in \u001b[0;36mQueuePool._do_get\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 174\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 175\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_create_connection\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 176\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m:\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/base.py:388\u001b[0m, in \u001b[0;36mPool._create_connection\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 386\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Called by subclasses to create a new ConnectionRecord.\"\"\"\u001b[39;00m\n\u001b[0;32m--> 388\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_ConnectionRecord\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/base.py:673\u001b[0m, in \u001b[0;36m_ConnectionRecord.__init__\u001b[0;34m(self, pool, connect)\u001b[0m\n\u001b[1;32m 672\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m connect:\n\u001b[0;32m--> 673\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m__connect\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 674\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfinalize_callback \u001b[38;5;241m=\u001b[39m deque()\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/base.py:899\u001b[0m, in \u001b[0;36m_ConnectionRecord.__connect\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 898\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mBaseException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[0;32m--> 899\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m util\u001b[38;5;241m.\u001b[39msafe_reraise():\n\u001b[1;32m 900\u001b[0m pool\u001b[38;5;241m.\u001b[39mlogger\u001b[38;5;241m.\u001b[39mdebug(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mError on connect(): \u001b[39m\u001b[38;5;132;01m%s\u001b[39;00m\u001b[38;5;124m\"\u001b[39m, e)\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/util/langhelpers.py:224\u001b[0m, in \u001b[0;36msafe_reraise.__exit__\u001b[0;34m(self, type_, value, traceback)\u001b[0m\n\u001b[1;32m 223\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_exc_info \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;66;03m# remove potential circular references\u001b[39;00m\n\u001b[0;32m--> 224\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m exc_value\u001b[38;5;241m.\u001b[39mwith_traceback(exc_tb)\n\u001b[1;32m 225\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/base.py:895\u001b[0m, in \u001b[0;36m_ConnectionRecord.__connect\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 894\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstarttime \u001b[38;5;241m=\u001b[39m time\u001b[38;5;241m.\u001b[39mtime()\n\u001b[0;32m--> 895\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdbapi_connection \u001b[38;5;241m=\u001b[39m connection \u001b[38;5;241m=\u001b[39m \u001b[43mpool\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_invoke_creator\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m 896\u001b[0m pool\u001b[38;5;241m.\u001b[39mlogger\u001b[38;5;241m.\u001b[39mdebug(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mCreated new connection \u001b[39m\u001b[38;5;132;01m%r\u001b[39;00m\u001b[38;5;124m\"\u001b[39m, connection)\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/engine/create.py:661\u001b[0m, in \u001b[0;36mcreate_engine..connect\u001b[0;34m(connection_record)\u001b[0m\n\u001b[1;32m 659\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m connection\n\u001b[0;32m--> 661\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mdialect\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconnect\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mcargs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mcparams\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/engine/default.py:629\u001b[0m, in \u001b[0;36mDefaultDialect.connect\u001b[0;34m(self, *cargs, **cparams)\u001b[0m\n\u001b[1;32m 627\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21mconnect\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39mcargs: Any, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mcparams: Any) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m DBAPIConnection:\n\u001b[1;32m 628\u001b[0m \u001b[38;5;66;03m# inherits the docstring from interfaces.Dialect.connect\u001b[39;00m\n\u001b[0;32m--> 629\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mloaded_dbapi\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconnect\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mcargs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mcparams\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/psycopg2/__init__.py:122\u001b[0m, in \u001b[0;36mconnect\u001b[0;34m(dsn, connection_factory, cursor_factory, **kwargs)\u001b[0m\n\u001b[1;32m 121\u001b[0m dsn \u001b[38;5;241m=\u001b[39m _ext\u001b[38;5;241m.\u001b[39mmake_dsn(dsn, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m--> 122\u001b[0m conn \u001b[38;5;241m=\u001b[39m \u001b[43m_connect\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdsn\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mconnection_factory\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mconnection_factory\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwasync\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 123\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m cursor_factory \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n", + "\u001b[0;31mOperationalError\u001b[0m: connection to server at \"34.67.80.127\", port 5432 failed: Operation timed out\n\tIs the server running on that host and accepting TCP/IP connections?\n", + "\nThe above exception was the direct cause of the following exception:\n", + "\u001b[0;31mOperationalError\u001b[0m Traceback (most recent call last)", + "Cell \u001b[0;32mIn[4], line 2\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[38;5;66;03m# Initialize QueryBackend with custom Cloud SQL config\u001b[39;00m\n\u001b[0;32m----> 2\u001b[0m query_backend \u001b[38;5;241m=\u001b[39m \u001b[43mQueryBackend\u001b[49m\u001b[43m(\u001b[49m\u001b[43mconfig\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdb_config\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 3\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m✓ QueryBackend initialized with direct Cloud SQL connection\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m 5\u001b[0m \u001b[38;5;66;03m# Test connection by listing sessions\u001b[39;00m\n", + "File \u001b[0;32m~/git/collab-environment/collab_env/data/db/query_backend.py:71\u001b[0m, in \u001b[0;36mQueryBackend.__init__\u001b[0;34m(self, config, backend)\u001b[0m\n\u001b[1;32m 69\u001b[0m \u001b[38;5;66;03m# Initialize database connection\u001b[39;00m\n\u001b[1;32m 70\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdb \u001b[38;5;241m=\u001b[39m DatabaseConnection(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mconfig)\n\u001b[0;32m---> 71\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdb\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconnect\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 73\u001b[0m \u001b[38;5;66;03m# Load SQL queries using aiosql with driver-specific adapter\u001b[39;00m\n\u001b[1;32m 74\u001b[0m queries_dir \u001b[38;5;241m=\u001b[39m Path(\u001b[38;5;18m__file__\u001b[39m)\u001b[38;5;241m.\u001b[39mparent \u001b[38;5;241m/\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mqueries\u001b[39m\u001b[38;5;124m\"\u001b[39m\n", + "File \u001b[0;32m~/git/collab-environment/collab_env/data/db/db_loader.py:158\u001b[0m, in \u001b[0;36mDatabaseConnection.connect\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 155\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mengine \u001b[38;5;241m=\u001b[39m create_engine(url, echo\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m)\n\u001b[1;32m 157\u001b[0m \u001b[38;5;66;03m# Test connection\u001b[39;00m\n\u001b[0;32m--> 158\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mengine\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconnect\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;28;01mas\u001b[39;00m conn:\n\u001b[1;32m 159\u001b[0m conn\u001b[38;5;241m.\u001b[39mexecute(text(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mSELECT 1\u001b[39m\u001b[38;5;124m\"\u001b[39m))\n\u001b[1;32m 161\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mconfig\u001b[38;5;241m.\u001b[39mbackend \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mpostgres\u001b[39m\u001b[38;5;124m'\u001b[39m:\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/engine/base.py:3277\u001b[0m, in \u001b[0;36mEngine.connect\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 3254\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21mconnect\u001b[39m(\u001b[38;5;28mself\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Connection:\n\u001b[1;32m 3255\u001b[0m \u001b[38;5;250m \u001b[39m\u001b[38;5;124;03m\"\"\"Return a new :class:`_engine.Connection` object.\u001b[39;00m\n\u001b[1;32m 3256\u001b[0m \n\u001b[1;32m 3257\u001b[0m \u001b[38;5;124;03m The :class:`_engine.Connection` acts as a Python context manager, so\u001b[39;00m\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 3274\u001b[0m \n\u001b[1;32m 3275\u001b[0m \u001b[38;5;124;03m \"\"\"\u001b[39;00m\n\u001b[0;32m-> 3277\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_connection_cls\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/engine/base.py:145\u001b[0m, in \u001b[0;36mConnection.__init__\u001b[0;34m(self, engine, connection, _has_events, _allow_revalidate, _allow_autobegin)\u001b[0m\n\u001b[1;32m 143\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_dbapi_connection \u001b[38;5;241m=\u001b[39m engine\u001b[38;5;241m.\u001b[39mraw_connection()\n\u001b[1;32m 144\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m dialect\u001b[38;5;241m.\u001b[39mloaded_dbapi\u001b[38;5;241m.\u001b[39mError \u001b[38;5;28;01mas\u001b[39;00m err:\n\u001b[0;32m--> 145\u001b[0m \u001b[43mConnection\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_handle_dbapi_exception_noconnection\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 146\u001b[0m \u001b[43m \u001b[49m\u001b[43merr\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdialect\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mengine\u001b[49m\n\u001b[1;32m 147\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 148\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m\n\u001b[1;32m 149\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/engine/base.py:2440\u001b[0m, in \u001b[0;36mConnection._handle_dbapi_exception_noconnection\u001b[0;34m(cls, e, dialect, engine, is_disconnect, invalidate_pool_on_disconnect, is_pre_ping)\u001b[0m\n\u001b[1;32m 2438\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m should_wrap:\n\u001b[1;32m 2439\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m sqlalchemy_exception \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m-> 2440\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m sqlalchemy_exception\u001b[38;5;241m.\u001b[39mwith_traceback(exc_info[\u001b[38;5;241m2\u001b[39m]) \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01me\u001b[39;00m\n\u001b[1;32m 2441\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m 2442\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m exc_info[\u001b[38;5;241m1\u001b[39m] \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/engine/base.py:143\u001b[0m, in \u001b[0;36mConnection.__init__\u001b[0;34m(self, engine, connection, _has_events, _allow_revalidate, _allow_autobegin)\u001b[0m\n\u001b[1;32m 141\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m connection \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m 142\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 143\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_dbapi_connection \u001b[38;5;241m=\u001b[39m \u001b[43mengine\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mraw_connection\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 144\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m dialect\u001b[38;5;241m.\u001b[39mloaded_dbapi\u001b[38;5;241m.\u001b[39mError \u001b[38;5;28;01mas\u001b[39;00m err:\n\u001b[1;32m 145\u001b[0m Connection\u001b[38;5;241m.\u001b[39m_handle_dbapi_exception_noconnection(\n\u001b[1;32m 146\u001b[0m err, dialect, engine\n\u001b[1;32m 147\u001b[0m )\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/engine/base.py:3301\u001b[0m, in \u001b[0;36mEngine.raw_connection\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 3279\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21mraw_connection\u001b[39m(\u001b[38;5;28mself\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m PoolProxiedConnection:\n\u001b[1;32m 3280\u001b[0m \u001b[38;5;250m \u001b[39m\u001b[38;5;124;03m\"\"\"Return a \"raw\" DBAPI connection from the connection pool.\u001b[39;00m\n\u001b[1;32m 3281\u001b[0m \n\u001b[1;32m 3282\u001b[0m \u001b[38;5;124;03m The returned object is a proxied version of the DBAPI\u001b[39;00m\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 3299\u001b[0m \n\u001b[1;32m 3300\u001b[0m \u001b[38;5;124;03m \"\"\"\u001b[39;00m\n\u001b[0;32m-> 3301\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mpool\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconnect\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/base.py:447\u001b[0m, in \u001b[0;36mPool.connect\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 439\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21mconnect\u001b[39m(\u001b[38;5;28mself\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m PoolProxiedConnection:\n\u001b[1;32m 440\u001b[0m \u001b[38;5;250m \u001b[39m\u001b[38;5;124;03m\"\"\"Return a DBAPI connection from the pool.\u001b[39;00m\n\u001b[1;32m 441\u001b[0m \n\u001b[1;32m 442\u001b[0m \u001b[38;5;124;03m The connection is instrumented such that when its\u001b[39;00m\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 445\u001b[0m \n\u001b[1;32m 446\u001b[0m \u001b[38;5;124;03m \"\"\"\u001b[39;00m\n\u001b[0;32m--> 447\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_ConnectionFairy\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_checkout\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/base.py:1264\u001b[0m, in \u001b[0;36m_ConnectionFairy._checkout\u001b[0;34m(cls, pool, threadconns, fairy)\u001b[0m\n\u001b[1;32m 1256\u001b[0m \u001b[38;5;129m@classmethod\u001b[39m\n\u001b[1;32m 1257\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21m_checkout\u001b[39m(\n\u001b[1;32m 1258\u001b[0m \u001b[38;5;28mcls\u001b[39m,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 1261\u001b[0m fairy: Optional[_ConnectionFairy] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m 1262\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m _ConnectionFairy:\n\u001b[1;32m 1263\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m fairy:\n\u001b[0;32m-> 1264\u001b[0m fairy \u001b[38;5;241m=\u001b[39m \u001b[43m_ConnectionRecord\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcheckout\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpool\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1266\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m threadconns \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m 1267\u001b[0m threadconns\u001b[38;5;241m.\u001b[39mcurrent \u001b[38;5;241m=\u001b[39m weakref\u001b[38;5;241m.\u001b[39mref(fairy)\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/base.py:711\u001b[0m, in \u001b[0;36m_ConnectionRecord.checkout\u001b[0;34m(cls, pool)\u001b[0m\n\u001b[1;32m 709\u001b[0m rec \u001b[38;5;241m=\u001b[39m cast(_ConnectionRecord, pool\u001b[38;5;241m.\u001b[39m_do_get())\n\u001b[1;32m 710\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 711\u001b[0m rec \u001b[38;5;241m=\u001b[39m \u001b[43mpool\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_do_get\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 713\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m 714\u001b[0m dbapi_connection \u001b[38;5;241m=\u001b[39m rec\u001b[38;5;241m.\u001b[39mget_connection()\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/impl.py:177\u001b[0m, in \u001b[0;36mQueuePool._do_get\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 175\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_create_connection()\n\u001b[1;32m 176\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m:\n\u001b[0;32m--> 177\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m util\u001b[38;5;241m.\u001b[39msafe_reraise():\n\u001b[1;32m 178\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_dec_overflow()\n\u001b[1;32m 179\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/util/langhelpers.py:224\u001b[0m, in \u001b[0;36msafe_reraise.__exit__\u001b[0;34m(self, type_, value, traceback)\u001b[0m\n\u001b[1;32m 222\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m exc_value \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m 223\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_exc_info \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;66;03m# remove potential circular references\u001b[39;00m\n\u001b[0;32m--> 224\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m exc_value\u001b[38;5;241m.\u001b[39mwith_traceback(exc_tb)\n\u001b[1;32m 225\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m 226\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_exc_info \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;66;03m# remove potential circular references\u001b[39;00m\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/impl.py:175\u001b[0m, in \u001b[0;36mQueuePool._do_get\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 173\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_inc_overflow():\n\u001b[1;32m 174\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 175\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_create_connection\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 176\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m:\n\u001b[1;32m 177\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m util\u001b[38;5;241m.\u001b[39msafe_reraise():\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/base.py:388\u001b[0m, in \u001b[0;36mPool._create_connection\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 385\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21m_create_connection\u001b[39m(\u001b[38;5;28mself\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m ConnectionPoolEntry:\n\u001b[1;32m 386\u001b[0m \u001b[38;5;250m \u001b[39m\u001b[38;5;124;03m\"\"\"Called by subclasses to create a new ConnectionRecord.\"\"\"\u001b[39;00m\n\u001b[0;32m--> 388\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_ConnectionRecord\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/base.py:673\u001b[0m, in \u001b[0;36m_ConnectionRecord.__init__\u001b[0;34m(self, pool, connect)\u001b[0m\n\u001b[1;32m 671\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m__pool \u001b[38;5;241m=\u001b[39m pool\n\u001b[1;32m 672\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m connect:\n\u001b[0;32m--> 673\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m__connect\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 674\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfinalize_callback \u001b[38;5;241m=\u001b[39m deque()\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/base.py:899\u001b[0m, in \u001b[0;36m_ConnectionRecord.__connect\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 897\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfresh \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mTrue\u001b[39;00m\n\u001b[1;32m 898\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mBaseException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[0;32m--> 899\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m util\u001b[38;5;241m.\u001b[39msafe_reraise():\n\u001b[1;32m 900\u001b[0m pool\u001b[38;5;241m.\u001b[39mlogger\u001b[38;5;241m.\u001b[39mdebug(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mError on connect(): \u001b[39m\u001b[38;5;132;01m%s\u001b[39;00m\u001b[38;5;124m\"\u001b[39m, e)\n\u001b[1;32m 901\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m 902\u001b[0m \u001b[38;5;66;03m# in SQLAlchemy 1.4 the first_connect event is not used by\u001b[39;00m\n\u001b[1;32m 903\u001b[0m \u001b[38;5;66;03m# the engine, so this will usually not be set\u001b[39;00m\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/util/langhelpers.py:224\u001b[0m, in \u001b[0;36msafe_reraise.__exit__\u001b[0;34m(self, type_, value, traceback)\u001b[0m\n\u001b[1;32m 222\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m exc_value \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m 223\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_exc_info \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;66;03m# remove potential circular references\u001b[39;00m\n\u001b[0;32m--> 224\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m exc_value\u001b[38;5;241m.\u001b[39mwith_traceback(exc_tb)\n\u001b[1;32m 225\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m 226\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_exc_info \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;66;03m# remove potential circular references\u001b[39;00m\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/base.py:895\u001b[0m, in \u001b[0;36m_ConnectionRecord.__connect\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 893\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m 894\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstarttime \u001b[38;5;241m=\u001b[39m time\u001b[38;5;241m.\u001b[39mtime()\n\u001b[0;32m--> 895\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdbapi_connection \u001b[38;5;241m=\u001b[39m connection \u001b[38;5;241m=\u001b[39m \u001b[43mpool\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_invoke_creator\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m 896\u001b[0m pool\u001b[38;5;241m.\u001b[39mlogger\u001b[38;5;241m.\u001b[39mdebug(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mCreated new connection \u001b[39m\u001b[38;5;132;01m%r\u001b[39;00m\u001b[38;5;124m\"\u001b[39m, connection)\n\u001b[1;32m 897\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfresh \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mTrue\u001b[39;00m\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/engine/create.py:661\u001b[0m, in \u001b[0;36mcreate_engine..connect\u001b[0;34m(connection_record)\u001b[0m\n\u001b[1;32m 658\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m connection \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m 659\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m connection\n\u001b[0;32m--> 661\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mdialect\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconnect\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mcargs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mcparams\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/engine/default.py:629\u001b[0m, in \u001b[0;36mDefaultDialect.connect\u001b[0;34m(self, *cargs, **cparams)\u001b[0m\n\u001b[1;32m 627\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21mconnect\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39mcargs: Any, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mcparams: Any) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m DBAPIConnection:\n\u001b[1;32m 628\u001b[0m \u001b[38;5;66;03m# inherits the docstring from interfaces.Dialect.connect\u001b[39;00m\n\u001b[0;32m--> 629\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mloaded_dbapi\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconnect\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mcargs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mcparams\u001b[49m\u001b[43m)\u001b[49m\n", + "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/psycopg2/__init__.py:122\u001b[0m, in \u001b[0;36mconnect\u001b[0;34m(dsn, connection_factory, cursor_factory, **kwargs)\u001b[0m\n\u001b[1;32m 119\u001b[0m kwasync[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124masync_\u001b[39m\u001b[38;5;124m'\u001b[39m] \u001b[38;5;241m=\u001b[39m kwargs\u001b[38;5;241m.\u001b[39mpop(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124masync_\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m 121\u001b[0m dsn \u001b[38;5;241m=\u001b[39m _ext\u001b[38;5;241m.\u001b[39mmake_dsn(dsn, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m--> 122\u001b[0m conn \u001b[38;5;241m=\u001b[39m \u001b[43m_connect\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdsn\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mconnection_factory\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mconnection_factory\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwasync\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 123\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m cursor_factory \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m 124\u001b[0m conn\u001b[38;5;241m.\u001b[39mcursor_factory \u001b[38;5;241m=\u001b[39m cursor_factory\n", + "\u001b[0;31mOperationalError\u001b[0m: (psycopg2.OperationalError) connection to server at \"34.67.80.127\", port 5432 failed: Operation timed out\n\tIs the server running on that host and accepting TCP/IP connections?\n\n(Background on this error at: https://sqlalche.me/e/20/e3q8)" + ] + } + ], + "source": [ + "# Initialize QueryBackend with custom Cloud SQL config\n", + "query_backend = QueryBackend(config=db_config)\n", + "print(\"✓ QueryBackend initialized with direct Cloud SQL connection\")\n", + "\n", + "# Test connection by listing sessions\n", + "try:\n", + " sessions = query_backend.get_sessions()\n", + " print(f\"\\n✓ Connection successful! Found {len(sessions)} sessions\")\n", + "except Exception as e:\n", + " print(f\"\\n✗ Connection failed: {e}\")\n", + " print(\"\\nTroubleshooting:\")\n", + " print(\" 1. Verify Cloud SQL public IP is correct\")\n", + " print(\" 2. Ensure your IP is in authorized networks:\")\n", + " print(\n", + " \" gcloud sql instances patch INSTANCE_NAME --authorized-networks=YOUR_IP/32\"\n", + " )\n", + " print(\" 3. Check password is correct\")\n", + " print(\" 4. Verify database exists and user has permissions\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "**Notes on Direct Public IP Connections:**\n\n⚠️ **Google Cloud Recommendation**: Use the Cloud SQL Auth Proxy instead of direct connections for better security.\n\nIf using direct connections:\n1. **Security**: Requires SSL/TLS configuration and authorized networks management\n2. **Public IP**: Get your instance's public IP with:\n ```bash\n gcloud sql instances describe INSTANCE_NAME --format=\"value(ipAddresses[0].ipAddress)\"\n ```\n3. **Firewall**: Must authorize your IP address:\n ```bash\n gcloud sql instances patch INSTANCE_NAME --authorized-networks=YOUR_IP/32\n ```\n4. **SSL**: Should configure SSL certificates for encrypted connections\n\n**Recommended**: See [docs/dashboard/CLOUD_SETUP.md](../dashboard/CLOUD_SETUP.md) for Cloud SQL Auth Proxy setup.\n\n---\n\n**If cells 3-6 above worked, skip cell 8 below and proceed directly to cell 9.**\n\n---\n\n### Default: Connect Using Environment Variables (Recommended with Proxy)" + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Initialize QueryBackend (reads from environment variables)\n", + "# Set DUCKDB_PATH or DB_BACKEND as needed before running\n", + "# Example: os.environ['DUCKDB_PATH'] = '/path/to/your.duckdb'\n", + "\n", + "query_backend = QueryBackend()\n", + "print(\"✓ QueryBackend initialized\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# List available sessions\n", + "sessions = query_backend.get_sessions(\"boids_2d_rollout\")\n", + "print(f\"Found {len(sessions)} sessions:\")\n", + "sessions[[\"session_id\"]]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Select a session to analyze (modify as needed)\n", + "session_id = (\n", + " \"rollout-boid_food_basic_vpluspplus_a_n0_h1_vr0.1_s0_rollout5_selfloops-unknown\"\n", + ")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Query Extended Properties at Session Scope\n", + "\n", + "This uses the same high-level API as the dashboard's `ExtendedPropertiesViewerWidget`." + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Create session-scope QueryScope\n", + "scope = QueryScope.from_session(\n", + " session_id=session_id, agent_type=\"agent\" # Filter to 'agent' type (not 'target')\n", + ")\n", + "\n", + "# Create AnalysisContext (mimics dashboard widget pattern)\n", + "context = AnalysisContext(query_backend=query_backend, scope=scope)\n", + "\n", + "print(f\"✓ Created context with scope: {scope}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Get available properties for this session\n", + "params = context.get_query_params()\n", + "available_props = query_backend.get_available_properties(**params)\n", + "\n", + "print(f\"\\nAvailable properties ({len(available_props)} total):\")\n", + "available_props[[\"property_id\", \"property_name\", \"data_type\", \"unit\"]]" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Query property distributions (aggregated across all episodes in session)\n", + "distributions = query_backend.get_property_distributions(**params)\n", + "\n", + "print(f\"\\n✓ Loaded {len(distributions)} property observations\")\n", + "print(f\" Properties: {distributions['property_id'].unique().tolist()}\")" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "distributions.head()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Visualize Selected Properties with Violin Plots" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Select properties to visualize (modify as needed)\n", + "# Example: attention weight properties from GNN models\n", + "selected_properties = [\"attn_weight_food\", \"attn_weight_boid\", \"attn_weight_self\"]\n", + "\n", + "# Filter distributions to selected properties\n", + "filtered_df = distributions[\n", + " distributions[\"property_id\"].isin(selected_properties)\n", + "].copy()\n", + "\n", + "if len(filtered_df) == 0:\n", + " print(f\"⚠ No data found for properties: {selected_properties}\")\n", + " print(f\" Available properties: {distributions['property_id'].unique().tolist()}\")\n", + "else:\n", + " print(\n", + " f\"✓ Filtered to {len(filtered_df)} observations across {len(selected_properties)} properties\"\n", + " )" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Create violin plot using seaborn\n", + "if len(filtered_df) > 0:\n", + " # Set figure size and style\n", + " plt.figure(figsize=(12, 6))\n", + " sns.set_style(\"whitegrid\")\n", + "\n", + " # Create violin plot\n", + " sns.violinplot(data=filtered_df, x=\"property_id\", y=\"value_float\", palette=\"Set2\")\n", + "\n", + " # Customize plot\n", + " plt.title(\n", + " f\"Extended Property Distributions - Session: {session_id}\",\n", + " fontsize=14,\n", + " fontweight=\"bold\",\n", + " )\n", + " plt.xlabel(\"Property\", fontsize=12)\n", + " plt.ylabel(\"Value\", fontsize=12)\n", + " plt.xticks(rotation=45, ha=\"right\")\n", + " plt.tight_layout()\n", + "\n", + " plt.show()\n", + "\n", + " # Print summary statistics\n", + " print(\"\\nSummary Statistics:\")\n", + " summary = filtered_df.groupby(\"property_id\")[\"value_float\"].agg(\n", + " [\"count\", \"mean\", \"std\", \"min\", \"max\"]\n", + " )\n", + " print(summary)" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "## Cleanup" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Close database connection\n", + "query_backend.close()\n", + "print(\"✓ Connection closed\")" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": "## Summary\n\nThis notebook demonstrates:\n\n1. **High-Level API**: Uses the same `QueryBackend` and `AnalysisContext` pattern as dashboard widgets\n2. **Session Scope**: Queries property distributions aggregated across all episodes in a session\n3. **Property Agnostic**: Works with any extended property (attention weights, distances, speeds, etc.)\n4. **Visualization**: Creates violin plots using seaborn for distribution analysis\n5. **Cloud SQL Integration**: Multiple connection approaches with security considerations\n\n### Key Components\n\n- `QueryBackend`: High-level database query interface\n- `QueryScope.from_session()`: Define session-scope analysis\n- `AnalysisContext`: Merges scope + shared parameters\n- `get_property_distributions()`: Query raw property values for distributions\n- `sns.violinplot()`: Visualize property distributions\n\n### Database Connection Options (Priority Order)\n\n**Option 1: Cloud SQL Auth Proxy + Environment Variables (RECOMMENDED)**\n- Start proxy: `./cloud-sql-proxy PROJECT:REGION:INSTANCE --port 5433`\n- Set environment variables to connect via localhost:5433\n- Initialize `QueryBackend()` with no arguments\n- **Benefits**: SSL/TLS encryption, IAM auth, no IP management\n- **Best for**: Production notebooks, shared environments, secure connections\n- See Cell 3 and [docs/dashboard/CLOUD_SETUP.md](../dashboard/CLOUD_SETUP.md)\n\n**Option 2: Local DuckDB (Cell 8)**\n- Initialize `QueryBackend()` with `DUCKDB_PATH` environment variable\n- **Best for**: Local development, testing, offline work\n\n**Option 3: Direct Public IP (Cells 4-6, NOT RECOMMENDED)**\n- ⚠️ Bypasses proxy security features (no automatic SSL/TLS, IAM auth)\n- Requires firewall configuration and authorized networks management\n- Construct `PostgresConfig` with public IP and credentials\n- Initialize `QueryBackend(config=db_config)`\n- **Only for**: Special testing scenarios, legacy systems\n\n### Customization\n\nModify these cells to customize the analysis:\n- Cell 3: Configure Cloud SQL Auth Proxy connection\n- Cells 4-6: Configure direct connection (if absolutely necessary)\n- Cell 9/10: Select different session\n- Cell 17: Change `selected_properties` to analyze different properties\n- Cell 18: Customize plot style, colors, or add statistical overlays\n\n### References\n\n- **Recommended Setup**: [docs/dashboard/CLOUD_SETUP.md](../dashboard/CLOUD_SETUP.md)\n- **Google Cloud SQL Auth Proxy**: https://cloud.google.com/sql/docs/postgres/connect-auth-proxy\n- **Dashboard Widgets**: [collab_env/dashboard/widgets/](../../collab_env/dashboard/widgets/)" + } + ], + "metadata": { + "kernelspec": { + "display_name": ".venv-310", + "language": "python", + "name": "python3" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 3 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython3", + "version": "3.10.16" + } + }, + "nbformat": 4, + "nbformat_minor": 4 +} \ No newline at end of file diff --git a/scripts/lint.sh b/scripts/lint.sh index fafd9d14..03168126 100755 --- a/scripts/lint.sh +++ b/scripts/lint.sh @@ -1,8 +1,8 @@ #!/bin/bash set -euxo pipefail -SRC="tests/ collab_env/" +SRC="collab_env tests" mypy $SRC ruff check $SRC -ruff format --diff $SRC +ruff format --check $SRC From 987314708789cbf07bd3570147a46ea8f4a90adc Mon Sep 17 00:00:00 2001 From: Dmitry Batenkov Date: Mon, 23 Feb 2026 11:39:28 -0500 Subject: [PATCH 17/21] not sure how these got added --- docs/data/db/basic_db.ipynb | 1233 --------------------------- docs/data/extended_properties.ipynb | 506 ----------- 2 files changed, 1739 deletions(-) delete mode 100644 docs/data/db/basic_db.ipynb delete mode 100644 docs/data/extended_properties.ipynb diff --git a/docs/data/db/basic_db.ipynb b/docs/data/db/basic_db.ipynb deleted file mode 100644 index cf413407..00000000 --- a/docs/data/db/basic_db.ipynb +++ /dev/null @@ -1,1233 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Database Query & Insert Tutorial\n", - "\n", - "**Concise guide to querying and inserting data into the tracking analytics database.**\n", - "\n", - "This notebook covers:\n", - "- Connecting to DuckDB and PostgreSQL\n", - "- Querying sessions, episodes, and observations\n", - "- Querying extended properties\n", - "- Inserting data (DuckDB only)\n", - "- Using the high-level QueryBackend API" - ] - }, - { - "cell_type": "code", - "execution_count": 1, - "metadata": {}, - "outputs": [], - "source": [ - "import os\n", - "from pathlib import Path\n", - "\n", - "import numpy as np\n", - "import pandas as pd\n", - "\n", - "from collab_env.data.db.config import DBConfig, get_db_config\n", - "from collab_env.data.db.db_loader import Boids3DLoader, DatabaseConnection\n", - "from collab_env.data.db.query_backend import QueryBackend" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Part 1: Connecting to Databases\n", - "\n", - "The system supports both **DuckDB** (local file) and **PostgreSQL** (server)." - ] - }, - { - "cell_type": "code", - "execution_count": 2, - "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "2025-11-13 17:02:58 | INFO | collab_env.data.db.db_loader:connect:113 - Connected to DuckDB: /tmp/test_tutorial.duckdb\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "✓ Connected to DuckDB: /tmp/test_tutorial.duckdb\n" - ] - } - ], - "source": [ - "# Connect to DuckDB (local file - for testing and insertions)\n", - "duckdb_path = \"/tmp/test_tutorial.duckdb\"\n", - "\n", - "# Set environment variable before creating config\n", - "os.environ[\"DUCKDB_PATH\"] = duckdb_path\n", - "\n", - "duckdb_config = get_db_config(backend=\"duckdb\")\n", - "\n", - "db_duckdb = DatabaseConnection(duckdb_config)\n", - "db_duckdb.connect()\n", - "print(f\"✓ Connected to DuckDB: {duckdb_path}\")" - ] - }, - { - "cell_type": "code", - "execution_count": 3, - "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "2025-11-13 17:03:02 | INFO | collab_env.data.db.db_loader:connect:111 - Connected to PostgreSQL: tracking_analytics\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "✓ Connected to PostgreSQL: tracking_analytics\n" - ] - } - ], - "source": [ - "# Connect to PostgreSQL (if available - for production queries)\n", - "# Skip this cell if you don't have PostgreSQL running\n", - "try:\n", - " # Set environment variables for PostgreSQL\n", - " os.environ[\"DB_BACKEND\"] = \"postgres\"\n", - " os.environ[\"POSTGRES_DB\"] = \"tracking_analytics\"\n", - " os.environ[\"POSTGRES_USER\"] = \"postgres\" # TODO: change to your username\n", - " os.environ[\"POSTGRES_PASSWORD\"] = \"password\" # TODO: change to your password\n", - "\n", - " postgres_config = get_db_config(backend=\"postgres\")\n", - "\n", - " db_postgres = DatabaseConnection(postgres_config)\n", - " db_postgres.connect()\n", - " print(\"✓ Connected to PostgreSQL: tracking_analytics\")\n", - "except Exception as e:\n", - " print(f\"⚠ PostgreSQL not available: {e}\")\n", - " db_postgres = None" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Part 2: Initializing Test Database\n", - "\n", - "Create tables and seed data in DuckDB for testing." - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "metadata": {}, - "outputs": [ - { - "name": "stderr", - "output_type": "stream", - "text": [ - "\u001b[32m2025-11-13 17:03:08\u001b[0m | \u001b[32m\u001b[1mSUCCESS \u001b[0m | \u001b[32m\u001b[1mConnected to DuckDB: /tmp/test_tutorial.duckdb\u001b[0m\n", - "\u001b[32m2025-11-13 17:03:08\u001b[0m | \u001b[32m\u001b[1mSUCCESS \u001b[0m | \u001b[32m\u001b[1mExecuted 01_core_tables.sql\u001b[0m\n", - "\u001b[32m2025-11-13 17:03:08\u001b[0m | \u001b[32m\u001b[1mSUCCESS \u001b[0m | \u001b[32m\u001b[1mExecuted 02_extended_properties.sql\u001b[0m\n", - "\u001b[32m2025-11-13 17:03:08\u001b[0m | \u001b[32m\u001b[1mSUCCESS \u001b[0m | \u001b[32m\u001b[1mExecuted 03_seed_data.sql\u001b[0m\n" - ] - }, - { - "name": "stdout", - "output_type": "stream", - "text": [ - "Executing 01_core_tables.sql...\n", - "Executing 02_extended_properties.sql...\n", - "Executing 03_seed_data.sql...\n", - "✓ DuckDB schema initialized\n" - ] - } - ], - "source": [ - "# Initialize DuckDB with schema\n", - "from collab_env.data.db.init_database import DatabaseBackend, get_schema_files\n", - "from collab_env.data.file_utils import get_project_root\n", - "\n", - "# Get schema files\n", - "project_root = get_project_root()\n", - "schema_dir = project_root / \"schema\"\n", - "schema_files = get_schema_files(schema_dir)\n", - "\n", - "# Create backend and execute schema\n", - "backend = DatabaseBackend(duckdb_config)\n", - "backend.connect()\n", - "\n", - "for schema_file in schema_files:\n", - " print(f\"Executing {schema_file.name}...\")\n", - " backend.execute_file(schema_file)\n", - "\n", - "backend.close()\n", - "print(\"✓ DuckDB schema initialized\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Part 3: Inserting Data (DuckDB Only)\n", - "\n", - "Insert sample session, episode, and observations data." - ] - }, - { - "cell_type": "code", - "execution_count": 5, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "✓ Inserted session\n" - ] - } - ], - "source": [ - "# Insert a test session\n", - "import json\n", - "\n", - "session_data = {\n", - " \"session_id\": \"test-session-001\",\n", - " \"session_name\": \"Tutorial Example Session\",\n", - " \"category_id\": \"boids_3d\",\n", - " \"config\": json.dumps({\"num_agents\": 10, \"scene_size\": 480}),\n", - " \"metadata\": json.dumps({\"notes\": \"Created in tutorial notebook\"}),\n", - "}\n", - "\n", - "db_duckdb.execute(\n", - " \"\"\"\n", - " INSERT INTO sessions (session_id, session_name, category_id, config, metadata)\n", - " VALUES (:session_id, :session_name, :category_id, :config, :metadata)\n", - " \"\"\",\n", - " session_data,\n", - ")\n", - "print(\"✓ Inserted session\")" - ] - }, - { - "cell_type": "code", - "execution_count": 6, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "✓ Inserted episode\n" - ] - } - ], - "source": [ - "# Insert a test episode\n", - "episode_data = {\n", - " \"episode_id\": \"test-episode-001\",\n", - " \"session_id\": \"test-session-001\",\n", - " \"episode_number\": 0,\n", - " \"num_frames\": 100,\n", - " \"num_agents\": 10,\n", - " \"frame_rate\": 30.0,\n", - " \"file_path\": \"/tmp/test_episode.parquet\",\n", - "}\n", - "\n", - "db_duckdb.execute(\n", - " \"\"\"\n", - " INSERT INTO episodes (episode_id, session_id, episode_number, num_frames, num_agents, frame_rate, file_path)\n", - " VALUES (:episode_id, :session_id, :episode_number, :num_frames, :num_agents, :frame_rate, :file_path)\n", - " \"\"\",\n", - " episode_data,\n", - ")\n", - "print(\"✓ Inserted episode\")" - ] - }, - { - "cell_type": "code", - "execution_count": 7, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "✓ Inserted 1000 observations\n" - ] - } - ], - "source": [ - "# Insert test observations using pandas (bulk insert)\n", - "num_agents = 10\n", - "num_frames = 100\n", - "\n", - "# Generate synthetic trajectory data\n", - "observations = []\n", - "for time_idx in range(num_frames):\n", - " for agent_id in range(num_agents):\n", - " # Simple circular motion\n", - " angle = 2 * np.pi * time_idx / num_frames + agent_id * 0.2\n", - " radius = 100 + agent_id * 10\n", - "\n", - " x = 240 + radius * np.cos(angle)\n", - " y = 240 + radius * np.sin(angle)\n", - " z = 50 + 20 * np.sin(angle * 2)\n", - "\n", - " v_x = -radius * np.sin(angle) * 2 * np.pi / num_frames\n", - " v_y = radius * np.cos(angle) * 2 * np.pi / num_frames\n", - " v_z = 40 * np.cos(angle * 2) * 2 * np.pi / num_frames\n", - "\n", - " observations.append(\n", - " {\n", - " \"episode_id\": \"test-episode-001\",\n", - " \"time_index\": time_idx,\n", - " \"agent_id\": agent_id,\n", - " \"agent_type_id\": \"agent\",\n", - " \"x\": x,\n", - " \"y\": y,\n", - " \"z\": z,\n", - " \"v_x\": v_x,\n", - " \"v_y\": v_y,\n", - " \"v_z\": v_z,\n", - " }\n", - " )\n", - "\n", - "obs_df = pd.DataFrame(observations)\n", - "db_duckdb.insert_dataframe(obs_df, \"observations\", if_exists=\"append\")\n", - "print(f\"✓ Inserted {len(obs_df)} observations\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Insert extended properties (distance to target)\n", - "# First, get observation IDs\n", - "obs_ids = db_duckdb.fetch_all(\n", - " \"\"\"\n", - " SELECT observation_id, time_index, agent_id\n", - " FROM observations\n", - " WHERE episode_id = :episode_id\n", - " ORDER BY time_index, agent_id\n", - " \"\"\",\n", - " {\"episode_id\": \"test-episode-001\"},\n", - ")\n", - "\n", - "# Compute synthetic distance to target center\n", - "target_center = np.array([240, 240, 50])\n", - "extended_props = []\n", - "\n", - "for obs_id, time_idx, agent_id in obs_ids:\n", - " # Get position from observations\n", - " obs = obs_df[\n", - " (obs_df[\"time_index\"] == time_idx) & (obs_df[\"agent_id\"] == agent_id)\n", - " ].iloc[0]\n", - " pos = np.array([obs[\"x\"], obs[\"y\"], obs[\"z\"]])\n", - " distance = np.linalg.norm(pos - target_center)\n", - "\n", - " extended_props.append(\n", - " {\n", - " \"observation_id\": obs_id,\n", - " \"property_id\": \"distance_to_target_center\",\n", - " \"value_float\": distance,\n", - " \"value_text\": None,\n", - " }\n", - " )\n", - "\n", - "ext_df = pd.DataFrame(extended_props)\n", - "db_duckdb.insert_dataframe(ext_df, \"extended_properties\", if_exists=\"append\")\n", - "print(f\"✓ Inserted {len(ext_df)} extended properties\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Part 4: Basic Queries\n", - "\n", - "Query the database using low-level SQL." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Query 1: List all sessions\n", - "sessions = db_duckdb.fetch_all(\"SELECT * FROM sessions\")\n", - "print(\"Sessions:\")\n", - "for session in sessions:\n", - " print(f\" - {session[0]}: {session[1]} ({session[2]})\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Query 2: Get episodes for a session\n", - "episodes = db_duckdb.fetch_all(\n", - " \"\"\"\n", - " SELECT episode_id, episode_number, num_frames, num_agents, frame_rate\n", - " FROM episodes\n", - " WHERE session_id = :session_id\n", - " ORDER BY episode_number\n", - " \"\"\",\n", - " {\"session_id\": \"test-session-001\"},\n", - ")\n", - "\n", - "print(\"\\nEpisodes for test-session-001:\")\n", - "for ep in episodes:\n", - " print(f\" - {ep[0]}: {ep[2]} frames, {ep[3]} agents @ {ep[4]} fps\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Query 3: Get observations with computed speed\n", - "obs_query = \"\"\"\n", - "SELECT \n", - " time_index,\n", - " agent_id,\n", - " x, y, z,\n", - " v_x, v_y, v_z,\n", - " sqrt(v_x*v_x + v_y*v_y + v_z*v_z) as speed\n", - "FROM observations\n", - "WHERE episode_id = :episode_id\n", - " AND time_index < 5\n", - "ORDER BY time_index, agent_id\n", - "\"\"\"\n", - "\n", - "from sqlalchemy import text\n", - "\n", - "with db_duckdb.engine.connect() as conn:\n", - " result = conn.execute(text(obs_query), {\"episode_id\": \"test-episode-001\"})\n", - " obs_df_query = pd.DataFrame(result.fetchall(), columns=result.keys())\n", - "\n", - "print(\"\\nFirst 5 frames of observations:\")\n", - "print(\n", - " obs_df_query[[\"time_index\", \"agent_id\", \"x\", \"y\", \"z\", \"speed\"]].to_string(\n", - " index=False\n", - " )\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Query 4: Get observations with extended properties\n", - "extended_query = \"\"\"\n", - "SELECT \n", - " o.time_index,\n", - " o.agent_id,\n", - " o.x, o.y, o.z,\n", - " pd.property_name,\n", - " ep.value_float\n", - "FROM observations o\n", - "JOIN extended_properties ep ON o.observation_id = ep.observation_id\n", - "JOIN property_definitions pd ON ep.property_id = pd.property_id\n", - "WHERE o.episode_id = :episode_id\n", - " AND o.time_index < 5\n", - "ORDER BY o.time_index, o.agent_id\n", - "\"\"\"\n", - "\n", - "with db_duckdb.engine.connect() as conn:\n", - " result = conn.execute(text(extended_query), {\"episode_id\": \"test-episode-001\"})\n", - " ext_query_df = pd.DataFrame(result.fetchall(), columns=result.keys())\n", - "\n", - "print(\"\\nObservations with extended properties (first 5 frames):\")\n", - "print(ext_query_df.to_string(index=False))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Query 5: Aggregate statistics\n", - "stats_query = \"\"\"\n", - "SELECT \n", - " COUNT(*) as total_observations,\n", - " COUNT(DISTINCT agent_id) as num_agents,\n", - " COUNT(DISTINCT time_index) as num_frames,\n", - " AVG(sqrt(v_x*v_x + v_y*v_y + v_z*v_z)) as avg_speed,\n", - " MAX(sqrt(v_x*v_x + v_y*v_y + v_z*v_z)) as max_speed\n", - "FROM observations\n", - "WHERE episode_id = :episode_id\n", - "\"\"\"\n", - "\n", - "stats = db_duckdb.fetch_one(stats_query, {\"episode_id\": \"test-episode-001\"})\n", - "print(\"\\nEpisode Statistics:\")\n", - "print(f\" Total observations: {stats[0]}\")\n", - "print(f\" Num agents: {stats[1]}\")\n", - "print(f\" Num frames: {stats[2]}\")\n", - "print(f\" Avg speed: {stats[3]:.2f}\")\n", - "print(f\" Max speed: {stats[4]:.2f}\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Part 5: QueryBackend API (Dashboard Pattern)\n", - "\n", - "The **QueryBackend** provides high-level methods for common queries. This is how the dashboard uses it." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Initialize QueryBackend with DuckDB\n", - "query = QueryBackend(config=duckdb_config)\n", - "print(\"✓ QueryBackend initialized\")" - ] - }, - { - "cell_type": "markdown", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "### 5.1 Session and Episode Discovery" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Get all categories\n", - "categories = query.get_categories()\n", - "print(\"Categories:\")\n", - "print(categories[[\"category_id\", \"category_name\"]].to_string(index=False))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Get sessions by category\n", - "sessions = query.get_sessions(category_id=\"boids_3d\")\n", - "print(\"\\nBoids 3D Sessions:\")\n", - "print(sessions[[\"session_id\", \"session_name\", \"category_id\"]].to_string(index=False))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Get episodes for a session\n", - "episodes = query.get_episodes(\"test-session-001\")\n", - "print(\"\\nEpisodes:\")\n", - "print(\n", - " episodes[[\"episode_id\", \"num_frames\", \"num_agents\", \"frame_rate\"]].to_string(\n", - " index=False\n", - " )\n", - ")" - ] - }, - { - "cell_type": "markdown", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "### 5.2 Spatial Analysis" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Get spatial heatmap (binned positions)\n", - "heatmap = query.get_spatial_heatmap(\n", - " episode_id=\"test-episode-001\", bin_size=50.0, agent_type=\"agent\"\n", - ")\n", - "print(\"\\nSpatial Heatmap (top 10 bins by density):\")\n", - "top_bins = heatmap.nlargest(10, \"density\")[[\"x_bin\", \"y_bin\", \"z_bin\", \"density\"]]\n", - "print(top_bins.to_string(index=False))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Get episode tracks for visualization\n", - "tracks = query.get_episode_tracks(\n", - " episode_id=\"test-episode-001\", start_time=0, end_time=10\n", - ")\n", - "print(\"\\nTracks (first 10 frames):\")\n", - "print(\n", - " tracks[[\"agent_id\", \"time_index\", \"x\", \"y\", \"z\", \"speed\"]]\n", - " .head(20)\n", - " .to_string(index=False)\n", - ")" - ] - }, - { - "cell_type": "markdown", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "### 5.3 Episode Tracks (for visualization)" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Get available extended properties\n", - "props = query.get_available_properties(\"test-episode-001\")\n", - "print(\"\\nAvailable Extended Properties:\")\n", - "print(\n", - " props[[\"property_id\", \"property_name\", \"data_type\", \"unit\"]].to_string(index=False)\n", - ")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Get property distributions for histogram\n", - "dist = query.get_property_distributions(\n", - " episode_id=\"test-episode-001\", property_ids=[\"distance_to_target_center\"]\n", - ")\n", - "print(\"\\nDistance to Target Distribution:\")\n", - "print(f\" Count: {len(dist)}\")\n", - "print(f\" Mean: {dist['value_float'].mean():.2f}\")\n", - "print(f\" Std: {dist['value_float'].std():.2f}\")\n", - "print(f\" Min: {dist['value_float'].min():.2f}\")\n", - "print(f\" Max: {dist['value_float'].max():.2f}\")" - ] - }, - { - "cell_type": "markdown", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "### 5.4 Extended Properties" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Get extended properties time series (windowed)\n", - "timeseries = query.get_extended_properties_timeseries(\n", - " episode_id=\"test-episode-001\",\n", - " window_size=20,\n", - " property_ids=[\"distance_to_target_center\"],\n", - ")\n", - "print(\"\\nExtended Properties Time Series (20-frame windows):\")\n", - "print(\n", - " timeseries[[\"time_window\", \"property_id\", \"avg_value\", \"std_value\"]].to_string(\n", - " index=False\n", - " )\n", - ")" - ] - }, - { - "cell_type": "markdown", - "source": [ - "#", - "#", - "#", - " ", - "5", - ".", - "6", - " ", - "D", - "a", - "s", - "h", - "b", - "o", - "a", - "r", - "d", - " ", - "P", - "a", - "t", - "t", - "e", - "r", - "n", - ":", - " ", - "U", - "s", - "i", - "n", - "g", - " ", - "A", - "n", - "a", - "l", - "y", - "s", - "i", - "s", - "C", - "o", - "n", - "t", - "e", - "x", - "t", - "\n", - "\n", - "T", - "h", - "e", - " ", - "d", - "a", - "s", - "h", - "b", - "o", - "a", - "r", - "d", - " ", - "u", - "s", - "e", - "s", - " ", - "`", - "A", - "n", - "a", - "l", - "y", - "s", - "i", - "s", - "C", - "o", - "n", - "t", - "e", - "x", - "t", - "`", - " ", - "t", - "o", - " ", - "s", - "h", - "a", - "r", - "e", - " ", - "q", - "u", - "e", - "r", - "y", - " ", - "p", - "a", - "r", - "a", - "m", - "e", - "t", - "e", - "r", - "s", - " ", - "a", - "c", - "r", - "o", - "s", - "s", - " ", - "w", - "i", - "d", - "g", - "e", - "t", - "s", - ".", - " ", - "T", - "h", - "i", - "s", - " ", - "p", - "a", - "t", - "t", - "e", - "r", - "n", - " ", - "e", - "n", - "a", - "b", - "l", - "e", - "s", - ":", - "\n", - "-", - " ", - "C", - "o", - "n", - "s", - "i", - "s", - "t", - "e", - "n", - "t", - " ", - "p", - "a", - "r", - "a", - "m", - "e", - "t", - "e", - "r", - "s", - " ", - "a", - "c", - "r", - "o", - "s", - "s", - " ", - "m", - "u", - "l", - "t", - "i", - "p", - "l", - "e", - " ", - "a", - "n", - "a", - "l", - "y", - "s", - "e", - "s", - "\n", - "-", - " ", - "E", - "a", - "s", - "y", - " ", - "p", - "a", - "r", - "a", - "m", - "e", - "t", - "e", - "r", - " ", - "o", - "v", - "e", - "r", - "r", - "i", - "d", - "e", - "s", - " ", - "f", - "o", - "r", - " ", - "w", - "i", - "d", - "g", - "e", - "t", - "-", - "s", - "p", - "e", - "c", - "i", - "f", - "i", - "c", - " ", - "c", - "u", - "s", - "t", - "o", - "m", - "i", - "z", - "a", - "t", - "i", - "o", - "n", - "\n", - "-", - " ", - "C", - "e", - "n", - "t", - "r", - "a", - "l", - "i", - "z", - "e", - "d", - " ", - "s", - "c", - "o", - "p", - "e", - " ", - "m", - "a", - "n", - "a", - "g", - "e", - "m", - "e", - "n", - "t", - " ", - "(", - "e", - "p", - "i", - "s", - "o", - "d", - "e", - "/", - "s", - "e", - "s", - "s", - "i", - "o", - "n", - " ", - "l", - "e", - "v", - "e", - "l", - ")" - ], - "metadata": {} - }, - { - "cell_type": "code", - "source": [ - "# Import context classes (dashboard pattern)\n", - "from collab_env.dashboard.widgets import AnalysisContext, QueryScope, ScopeType\n", - "\n", - "# Create a query scope for an episode\n", - "scope = QueryScope(\n", - " scope_type=ScopeType.EPISODE,\n", - " episode_id=\"test-episode-001\",\n", - " session_id=\"test-session-001\",\n", - " start_time=0,\n", - " end_time=50,\n", - " agent_type=\"agent\",\n", - ")\n", - "\n", - "# Create analysis context with shared parameters\n", - "context = AnalysisContext(\n", - " query_backend=query,\n", - " scope=scope,\n", - " spatial_bin_size=20.0, # Shared spatial discretization\n", - " temporal_window_size=10, # Shared time window\n", - " min_samples=10, # Shared minimum sample threshold\n", - " on_loading=lambda msg: print(f\"⏳ {msg}\"),\n", - " on_success=lambda msg: print(f\"✓ {msg}\"),\n", - " on_error=lambda msg: print(f\"✗ {msg}\"),\n", - ")\n", - "\n", - "print(\"✓ Created AnalysisContext\")\n", - "print(f\" Scope: {scope.scope_type.value}\")\n", - "print(f\" Episode: {scope.episode_id}\")\n", - "print(f\" Time range: {scope.start_time}-{scope.end_time}\")\n", - "print(f\" Spatial bin: {context.spatial_bin_size}\")\n", - "print(f\" Time window: {context.temporal_window_size}\")" - ], - "metadata": {}, - "execution_count": null, - "outputs": [] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Use context to get merged query parameters\n", - "params = context.get_query_params()\n", - "print(\"\\nMerged Query Parameters:\")\n", - "for key, value in params.items():\n", - " print(f\" {key}: {value}\")\n", - "\n", - "# Query using merged parameters (dashboard pattern)\n", - "context.report_loading(\"Loading spatial heatmap...\")\n", - "heatmap = query.get_spatial_heatmap(**params)\n", - "context.report_success(f\"Loaded {len(heatmap)} bins\")\n", - "\n", - "print(\"\\nHeatmap with context parameters:\")\n", - "print(heatmap[[\"x_bin\", \"y_bin\", \"z_bin\", \"density\"]].head().to_string(index=False))" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Override specific parameters (widget-specific customization)\n", - "custom_params = context.get_query_params(\n", - " bin_size=10.0, min_count=5 # Override spatial bin size # Override minimum count\n", - ")\n", - "\n", - "print(\"\\nCustom Parameters (with overrides):\")\n", - "print(f\" bin_size: {custom_params['bin_size']} (was {params['bin_size']})\")\n", - "print(f\" min_count: {custom_params['min_count']} (was {params.get('min_count', 1)})\")\n", - "\n", - "# Query with custom parameters\n", - "custom_heatmap = query.get_spatial_heatmap(**custom_params)\n", - "print(f\"\\nCustom heatmap: {len(custom_heatmap)} bins (finer resolution)\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Part 6: Querying PostgreSQL (if available)\n", - "\n", - "Same queries work on PostgreSQL with production data." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "if db_postgres is not None:\n", - " # Initialize QueryBackend for PostgreSQL\n", - " postgres_config = get_db_config(backend=\"postgres\")\n", - " query_pg = QueryBackend(config=postgres_config)\n", - "\n", - " # Get sessions\n", - " sessions_pg = query_pg.get_sessions(category_id=\"boids_3d\")\n", - " print(\"PostgreSQL - Boids 3D Sessions:\")\n", - " print(f\" Found {len(sessions_pg)} sessions\")\n", - "\n", - " if len(sessions_pg) > 0:\n", - " # Get first session's episodes\n", - " session_id = sessions_pg.iloc[0][\"session_id\"]\n", - " episodes_pg = query_pg.get_episodes(session_id)\n", - " print(f\"\\n Episodes for {session_id}:\")\n", - " print(f\" Found {len(episodes_pg)} episodes\")\n", - "\n", - " if len(episodes_pg) > 0:\n", - " # Get spatial heatmap\n", - " episode_id = episodes_pg.iloc[0][\"episode_id\"]\n", - " heatmap_pg = query_pg.get_spatial_heatmap(episode_id, bin_size=20.0)\n", - " print(f\"\\n Heatmap for {episode_id}:\")\n", - " print(f\" Generated {len(heatmap_pg)} bins\")\n", - "\n", - " query_pg.close()\n", - "else:\n", - " print(\"PostgreSQL not available - skipping\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Part 7: Loading Real Data\n", - "\n", - "Load actual simulation data using data loaders." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Example: Load 3D boids simulation (if data exists)\n", - "# Uncomment and modify path as needed\n", - "\n", - "# sim_dir = Path(\"simulated_data/hackathon/hackathon-boid-small-200-sim_run-started-20250926-220926\")\n", - "# if sim_dir.exists():\n", - "# loader = Boids3DLoader(db_duckdb, max_episodes=2) # Load only 2 episodes\n", - "# loader.load_simulation(sim_dir)\n", - "# print(\"✓ Loaded simulation data\")\n", - "# else:\n", - "# print(\"Simulation directory not found\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Cleanup" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Close connections\n", - "query.close()\n", - "db_duckdb.close()\n", - "if db_postgres is not None:\n", - " db_postgres.close()\n", - "\n", - "print(\"✓ All connections closed\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Summary\n", - "\n", - "### Key Takeaways\n", - "\n", - "1. **Connecting**:\n", - " - Set environment variables (`DUCKDB_PATH`, `DB_BACKEND`, etc.)\n", - " - Use `get_db_config()` to create config from environment\n", - " - Use `DatabaseConnection` for low-level access\n", - " - Use `QueryBackend()` for high-level queries (no args needed!)\n", - "\n", - "2. **Inserting** (DuckDB only):\n", - " - `db.execute()` for single inserts with named parameters\n", - " - `db.insert_dataframe()` for bulk inserts (much faster)\n", - " - Use transactions for multi-operation consistency\n", - "\n", - "3. **Querying with QueryBackend**:\n", - " - **Discovery**: `get_categories()`, `get_sessions()`, `get_episodes()`\n", - " - **Spatial**: `get_spatial_heatmap()`, `get_episode_tracks()`\n", - " - **Properties**: `get_available_properties()`, `get_property_distributions()`, `get_extended_properties_timeseries()`\n", - " - All methods return pandas DataFrames\n", - " - All methods support optional time/agent filtering\n", - "\n", - "4. **Advanced Pattern (Dashboard)**:\n", - " - Create `QueryScope` to define what data to analyze\n", - " - Create `AnalysisContext` with shared parameters\n", - " - Use `context.get_query_params()` to merge scope + shared + custom params\n", - " - Pass merged params to QueryBackend methods\n", - " - Enables consistent parameters across multiple widgets/analyses\n", - "\n", - "5. **Best Practices**:\n", - " - Use DuckDB for testing and local development\n", - " - Use PostgreSQL for production and Grafana\n", - " - Use QueryBackend for cleaner, higher-level code\n", - " - Use AnalysisContext for multi-widget applications\n", - " - Always filter by episode_id for performance\n", - "\n", - "### Next Steps\n", - "\n", - "- [docs/data/db/README.md](README.md) - Complete database documentation\n", - "- [schema/README.md](../../../schema/README.md) - Database schema details\n", - "- [collab_env/dashboard/widgets/](../../dashboard/widgets/) - Widget examples using AnalysisContext\n", - "- [collab_env/data/db/queries/](../../data/db/queries/) - SQL query library" - ] - } - ], - "metadata": { - "kernelspec": { - "display_name": ".venv-310", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.10.16" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} \ No newline at end of file diff --git a/docs/data/extended_properties.ipynb b/docs/data/extended_properties.ipynb deleted file mode 100644 index 766ee9b3..00000000 --- a/docs/data/extended_properties.ipynb +++ /dev/null @@ -1,506 +0,0 @@ -{ - "cells": [ - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "# Extended Properties Viewer - Session-Scope Analysis\n", - "\n", - "**Minimal example of querying session-level extended properties and visualizing distributions.**\n", - "\n", - "This notebook demonstrates:\n", - "- Using the high-level QueryBackend API (same as dashboard widgets)\n", - "- Querying session-scope extended properties\n", - "- Visualizing property distributions with seaborn violin plots" - ] - }, - { - "cell_type": "code", - "execution_count": 1, - "metadata": {}, - "outputs": [], - "source": [ - "import os\n", - "\n", - "import matplotlib.pyplot as plt\n", - "import pandas as pd\n", - "import seaborn as sns\n", - "\n", - "from collab_env.dashboard.widgets.analysis_context import AnalysisContext\n", - "from collab_env.dashboard.widgets.query_scope import QueryScope, ScopeType\n", - "from collab_env.data.db.query_backend import QueryBackend" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Setup: Initialize QueryBackend and Session Scope" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": "### Recommended: Connect to Cloud SQL via Auth Proxy\n\n**Google's recommended approach** for connecting to Cloud SQL is using the Cloud SQL Auth Proxy, which provides:\n- ✅ Automatic SSL/TLS encryption\n- ✅ IAM-based authentication\n- ✅ No need to manage authorized networks\n- ✅ Works with both public and private IP\n\n**Setup (one-time):**\n```bash\n# Terminal 1: Start Cloud SQL Auth Proxy\n./cloud-sql-proxy PROJECT_ID:REGION:INSTANCE_NAME --port 5433\n\n# Keep this running while using the notebook\n```\n\n**In the notebook:** Set environment variables to connect via proxy on localhost:\n\n```python\nimport os\nos.environ['DB_BACKEND'] = 'postgres'\nos.environ['POSTGRES_HOST'] = 'localhost'\nos.environ['POSTGRES_PORT'] = '5433' # Proxy port\nos.environ['POSTGRES_DB'] = 'tracking_analytics'\nos.environ['POSTGRES_USER'] = 'postgres'\nos.environ['POSTGRES_PASSWORD'] = 'your-password' # Or fetch from Secret Manager\n```\n\nThen skip to **Cell 8** to use `QueryBackend()` with environment variables.\n\n---\n\n### Alternative (Not Recommended): Direct Public IP Connection\n\n⚠️ **Warning**: Direct connections bypass the proxy's security features. Google recommends using the Auth Proxy instead.\n\nIf you must connect directly (e.g., testing, special circumstances):" - }, - { - "cell_type": "code", - "source": [ - "# RECOMMENDED: Configure environment for Cloud SQL Auth Proxy connection\n", - "# Run this cell if you have the proxy running in another terminal\n", - "\n", - "import os\n", - "import subprocess\n", - "\n", - "# Fetch password from Secret Manager (recommended)\n", - "\n", - "\n", - "def get_password_from_secret_manager(secret_name=\"postgres-password\"):\n", - " \"\"\"Fetch password from Google Cloud Secret Manager.\"\"\"\n", - " try:\n", - " result = subprocess.run(\n", - " [\n", - " \"gcloud\",\n", - " \"secrets\",\n", - " \"versions\",\n", - " \"access\",\n", - " \"latest\",\n", - " \"--secret\",\n", - " secret_name,\n", - " ],\n", - " capture_output=True,\n", - " text=True,\n", - " check=True,\n", - " )\n", - " return result.stdout.strip()\n", - " except subprocess.CalledProcessError as e:\n", - " print(f\"⚠ Failed to fetch secret: {e}\")\n", - " return None\n", - "\n", - "\n", - "# Configure environment variables for proxy connection\n", - "os.environ[\"DB_BACKEND\"] = \"postgres\"\n", - "os.environ[\"POSTGRES_HOST\"] = \"localhost\" # Connect via proxy on localhost\n", - "os.environ[\"POSTGRES_PORT\"] = (\n", - " \"5433\" # Proxy port (use 5433 if local postgres uses 5432)\n", - ")\n", - "os.environ[\"POSTGRES_DB\"] = \"tracking_analytics\"\n", - "os.environ[\"POSTGRES_USER\"] = \"postgres\"\n", - "os.environ[\"POSTGRES_PASSWORD\"] = (\n", - " get_password_from_secret_manager() or \"your-password-here\"\n", - ")\n", - "\n", - "print(\"✓ Environment configured for Cloud SQL Auth Proxy connection\")\n", - "print(f\" Proxy endpoint: localhost:{os.environ['POSTGRES_PORT']}\")\n", - "print(f\" Database: {os.environ['POSTGRES_DB']}\")\n", - "print(f\" User: {os.environ['POSTGRES_USER']}\")\n", - "\n", - "# Now skip to Cell 8 to initialize QueryBackend() with these settings" - ], - "metadata": {}, - "execution_count": null, - "outputs": [] - }, - { - "cell_type": "markdown", - "source": "---\n\n**If using Cloud SQL Auth Proxy (recommended), skip cells 6-9 below and jump directly to Cell 10.**\n\n---", - "metadata": {} - }, - { - "cell_type": "code", - "execution_count": 2, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "✓ Password retrieved successfully\n" - ] - } - ], - "source": [ - "# Option 1: Fetch password from Google Cloud Secret Manager (recommended)\n", - "# Requires: gcloud CLI installed and authenticated\n", - "import subprocess\n", - "\n", - "\n", - "def get_password_from_secret_manager(secret_name=\"postgres-password\"):\n", - " \"\"\"Fetch password from Google Cloud Secret Manager.\"\"\"\n", - " try:\n", - " result = subprocess.run(\n", - " [\n", - " \"gcloud\",\n", - " \"secrets\",\n", - " \"versions\",\n", - " \"access\",\n", - " \"latest\",\n", - " \"--secret\",\n", - " secret_name,\n", - " ],\n", - " capture_output=True,\n", - " text=True,\n", - " check=True,\n", - " )\n", - " return result.stdout.strip()\n", - " except subprocess.CalledProcessError as e:\n", - " print(f\"⚠ Failed to fetch secret: {e}\")\n", - " return None\n", - "\n", - "\n", - "# Option 2: Specify password directly (not recommended for shared notebooks)\n", - "# db_password = 'your-password-here'\n", - "\n", - "\n", - "# Try Secret Manager first, fallback to environment variable\n", - "db_password = get_password_from_secret_manager() or os.getenv(\"POSTGRES_PASSWORD\")\n", - "\n", - "if db_password:\n", - " print(\"✓ Password retrieved successfully\")\n", - "else:\n", - " print(\"⚠ No password found - will fail to connect\")" - ] - }, - { - "cell_type": "code", - "execution_count": 3, - "metadata": {}, - "outputs": [ - { - "name": "stdout", - "output_type": "stream", - "text": [ - "✓ Database configuration created:\n", - " Host: 34.67.80.127\n", - " Database: tracking_analytics\n", - " User: postgres\n", - " Password set: True\n" - ] - } - ], - "source": [ - "# Construct database configuration for direct Cloud SQL connection\n", - "from collab_env.data.db.config import DBConfig, PostgresConfig\n", - "\n", - "# Cloud SQL instance details (modify with your values)\n", - "CLOUD_SQL_PUBLIC_IP = \"34.67.80.127\" # Get from: gcloud sql instances describe INSTANCE_NAME --format=\"value(ipAddresses[0].ipAddress)\"\n", - "DB_NAME = \"tracking_analytics\"\n", - "DB_USER = \"postgres\"\n", - "DB_PORT = 5432 # Default PostgreSQL port\n", - "\n", - "# Create PostgresConfig with direct connection parameters\n", - "postgres_config = PostgresConfig(\n", - " host=CLOUD_SQL_PUBLIC_IP,\n", - " port=DB_PORT,\n", - " dbname=DB_NAME,\n", - " user=DB_USER,\n", - " password=db_password,\n", - ")\n", - "\n", - "# Create DBConfig with postgres backend\n", - "db_config = DBConfig(backend=\"postgres\")\n", - "db_config.postgres = postgres_config # Override with custom config\n", - "\n", - "print(f\"✓ Database configuration created:\")\n", - "print(f\" Host: {postgres_config.host}\")\n", - "print(f\" Database: {postgres_config.dbname}\")\n", - "print(f\" User: {postgres_config.user}\")\n", - "print(f\" Password set: {postgres_config.password is not None}\")" - ] - }, - { - "cell_type": "code", - "execution_count": 4, - "metadata": {}, - "outputs": [ - { - "ename": "OperationalError", - "evalue": "(psycopg2.OperationalError) connection to server at \"34.67.80.127\", port 5432 failed: Operation timed out\n\tIs the server running on that host and accepting TCP/IP connections?\n\n(Background on this error at: https://sqlalche.me/e/20/e3q8)", - "output_type": "error", - "traceback": [ - "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", - "\u001b[0;31mOperationalError\u001b[0m Traceback (most recent call last)", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/engine/base.py:143\u001b[0m, in \u001b[0;36mConnection.__init__\u001b[0;34m(self, engine, connection, _has_events, _allow_revalidate, _allow_autobegin)\u001b[0m\n\u001b[1;32m 142\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 143\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_dbapi_connection \u001b[38;5;241m=\u001b[39m \u001b[43mengine\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mraw_connection\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 144\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m dialect\u001b[38;5;241m.\u001b[39mloaded_dbapi\u001b[38;5;241m.\u001b[39mError \u001b[38;5;28;01mas\u001b[39;00m err:\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/engine/base.py:3301\u001b[0m, in \u001b[0;36mEngine.raw_connection\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 3280\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Return a \"raw\" DBAPI connection from the connection pool.\u001b[39;00m\n\u001b[1;32m 3281\u001b[0m \n\u001b[1;32m 3282\u001b[0m \u001b[38;5;124;03mThe returned object is a proxied version of the DBAPI\u001b[39;00m\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 3299\u001b[0m \n\u001b[1;32m 3300\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[0;32m-> 3301\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mpool\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconnect\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/base.py:447\u001b[0m, in \u001b[0;36mPool.connect\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 440\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Return a DBAPI connection from the pool.\u001b[39;00m\n\u001b[1;32m 441\u001b[0m \n\u001b[1;32m 442\u001b[0m \u001b[38;5;124;03mThe connection is instrumented such that when its\u001b[39;00m\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 445\u001b[0m \n\u001b[1;32m 446\u001b[0m \u001b[38;5;124;03m\"\"\"\u001b[39;00m\n\u001b[0;32m--> 447\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_ConnectionFairy\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_checkout\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/base.py:1264\u001b[0m, in \u001b[0;36m_ConnectionFairy._checkout\u001b[0;34m(cls, pool, threadconns, fairy)\u001b[0m\n\u001b[1;32m 1263\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m fairy:\n\u001b[0;32m-> 1264\u001b[0m fairy \u001b[38;5;241m=\u001b[39m \u001b[43m_ConnectionRecord\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcheckout\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpool\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1266\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m threadconns \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/base.py:711\u001b[0m, in \u001b[0;36m_ConnectionRecord.checkout\u001b[0;34m(cls, pool)\u001b[0m\n\u001b[1;32m 710\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 711\u001b[0m rec \u001b[38;5;241m=\u001b[39m \u001b[43mpool\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_do_get\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 713\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/impl.py:177\u001b[0m, in \u001b[0;36mQueuePool._do_get\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 176\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m:\n\u001b[0;32m--> 177\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m util\u001b[38;5;241m.\u001b[39msafe_reraise():\n\u001b[1;32m 178\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_dec_overflow()\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/util/langhelpers.py:224\u001b[0m, in \u001b[0;36msafe_reraise.__exit__\u001b[0;34m(self, type_, value, traceback)\u001b[0m\n\u001b[1;32m 223\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_exc_info \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;66;03m# remove potential circular references\u001b[39;00m\n\u001b[0;32m--> 224\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m exc_value\u001b[38;5;241m.\u001b[39mwith_traceback(exc_tb)\n\u001b[1;32m 225\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/impl.py:175\u001b[0m, in \u001b[0;36mQueuePool._do_get\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 174\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 175\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_create_connection\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 176\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m:\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/base.py:388\u001b[0m, in \u001b[0;36mPool._create_connection\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 386\u001b[0m \u001b[38;5;250m\u001b[39m\u001b[38;5;124;03m\"\"\"Called by subclasses to create a new ConnectionRecord.\"\"\"\u001b[39;00m\n\u001b[0;32m--> 388\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_ConnectionRecord\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/base.py:673\u001b[0m, in \u001b[0;36m_ConnectionRecord.__init__\u001b[0;34m(self, pool, connect)\u001b[0m\n\u001b[1;32m 672\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m connect:\n\u001b[0;32m--> 673\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m__connect\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 674\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfinalize_callback \u001b[38;5;241m=\u001b[39m deque()\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/base.py:899\u001b[0m, in \u001b[0;36m_ConnectionRecord.__connect\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 898\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mBaseException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[0;32m--> 899\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m util\u001b[38;5;241m.\u001b[39msafe_reraise():\n\u001b[1;32m 900\u001b[0m pool\u001b[38;5;241m.\u001b[39mlogger\u001b[38;5;241m.\u001b[39mdebug(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mError on connect(): \u001b[39m\u001b[38;5;132;01m%s\u001b[39;00m\u001b[38;5;124m\"\u001b[39m, e)\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/util/langhelpers.py:224\u001b[0m, in \u001b[0;36msafe_reraise.__exit__\u001b[0;34m(self, type_, value, traceback)\u001b[0m\n\u001b[1;32m 223\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_exc_info \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;66;03m# remove potential circular references\u001b[39;00m\n\u001b[0;32m--> 224\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m exc_value\u001b[38;5;241m.\u001b[39mwith_traceback(exc_tb)\n\u001b[1;32m 225\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/base.py:895\u001b[0m, in \u001b[0;36m_ConnectionRecord.__connect\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 894\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstarttime \u001b[38;5;241m=\u001b[39m time\u001b[38;5;241m.\u001b[39mtime()\n\u001b[0;32m--> 895\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdbapi_connection \u001b[38;5;241m=\u001b[39m connection \u001b[38;5;241m=\u001b[39m \u001b[43mpool\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_invoke_creator\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m 896\u001b[0m pool\u001b[38;5;241m.\u001b[39mlogger\u001b[38;5;241m.\u001b[39mdebug(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mCreated new connection \u001b[39m\u001b[38;5;132;01m%r\u001b[39;00m\u001b[38;5;124m\"\u001b[39m, connection)\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/engine/create.py:661\u001b[0m, in \u001b[0;36mcreate_engine..connect\u001b[0;34m(connection_record)\u001b[0m\n\u001b[1;32m 659\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m connection\n\u001b[0;32m--> 661\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mdialect\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconnect\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mcargs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mcparams\u001b[49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/engine/default.py:629\u001b[0m, in \u001b[0;36mDefaultDialect.connect\u001b[0;34m(self, *cargs, **cparams)\u001b[0m\n\u001b[1;32m 627\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21mconnect\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39mcargs: Any, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mcparams: Any) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m DBAPIConnection:\n\u001b[1;32m 628\u001b[0m \u001b[38;5;66;03m# inherits the docstring from interfaces.Dialect.connect\u001b[39;00m\n\u001b[0;32m--> 629\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mloaded_dbapi\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconnect\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mcargs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mcparams\u001b[49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/psycopg2/__init__.py:122\u001b[0m, in \u001b[0;36mconnect\u001b[0;34m(dsn, connection_factory, cursor_factory, **kwargs)\u001b[0m\n\u001b[1;32m 121\u001b[0m dsn \u001b[38;5;241m=\u001b[39m _ext\u001b[38;5;241m.\u001b[39mmake_dsn(dsn, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m--> 122\u001b[0m conn \u001b[38;5;241m=\u001b[39m \u001b[43m_connect\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdsn\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mconnection_factory\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mconnection_factory\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwasync\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 123\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m cursor_factory \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n", - "\u001b[0;31mOperationalError\u001b[0m: connection to server at \"34.67.80.127\", port 5432 failed: Operation timed out\n\tIs the server running on that host and accepting TCP/IP connections?\n", - "\nThe above exception was the direct cause of the following exception:\n", - "\u001b[0;31mOperationalError\u001b[0m Traceback (most recent call last)", - "Cell \u001b[0;32mIn[4], line 2\u001b[0m\n\u001b[1;32m 1\u001b[0m \u001b[38;5;66;03m# Initialize QueryBackend with custom Cloud SQL config\u001b[39;00m\n\u001b[0;32m----> 2\u001b[0m query_backend \u001b[38;5;241m=\u001b[39m \u001b[43mQueryBackend\u001b[49m\u001b[43m(\u001b[49m\u001b[43mconfig\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mdb_config\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 3\u001b[0m \u001b[38;5;28mprint\u001b[39m(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124m✓ QueryBackend initialized with direct Cloud SQL connection\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[1;32m 5\u001b[0m \u001b[38;5;66;03m# Test connection by listing sessions\u001b[39;00m\n", - "File \u001b[0;32m~/git/collab-environment/collab_env/data/db/query_backend.py:71\u001b[0m, in \u001b[0;36mQueryBackend.__init__\u001b[0;34m(self, config, backend)\u001b[0m\n\u001b[1;32m 69\u001b[0m \u001b[38;5;66;03m# Initialize database connection\u001b[39;00m\n\u001b[1;32m 70\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdb \u001b[38;5;241m=\u001b[39m DatabaseConnection(\u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mconfig)\n\u001b[0;32m---> 71\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mdb\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconnect\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 73\u001b[0m \u001b[38;5;66;03m# Load SQL queries using aiosql with driver-specific adapter\u001b[39;00m\n\u001b[1;32m 74\u001b[0m queries_dir \u001b[38;5;241m=\u001b[39m Path(\u001b[38;5;18m__file__\u001b[39m)\u001b[38;5;241m.\u001b[39mparent \u001b[38;5;241m/\u001b[39m \u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mqueries\u001b[39m\u001b[38;5;124m\"\u001b[39m\n", - "File \u001b[0;32m~/git/collab-environment/collab_env/data/db/db_loader.py:158\u001b[0m, in \u001b[0;36mDatabaseConnection.connect\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 155\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mengine \u001b[38;5;241m=\u001b[39m create_engine(url, echo\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mFalse\u001b[39;00m)\n\u001b[1;32m 157\u001b[0m \u001b[38;5;66;03m# Test connection\u001b[39;00m\n\u001b[0;32m--> 158\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mengine\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconnect\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m \u001b[38;5;28;01mas\u001b[39;00m conn:\n\u001b[1;32m 159\u001b[0m conn\u001b[38;5;241m.\u001b[39mexecute(text(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mSELECT 1\u001b[39m\u001b[38;5;124m\"\u001b[39m))\n\u001b[1;32m 161\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mconfig\u001b[38;5;241m.\u001b[39mbackend \u001b[38;5;241m==\u001b[39m \u001b[38;5;124m'\u001b[39m\u001b[38;5;124mpostgres\u001b[39m\u001b[38;5;124m'\u001b[39m:\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/engine/base.py:3277\u001b[0m, in \u001b[0;36mEngine.connect\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 3254\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21mconnect\u001b[39m(\u001b[38;5;28mself\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m Connection:\n\u001b[1;32m 3255\u001b[0m \u001b[38;5;250m \u001b[39m\u001b[38;5;124;03m\"\"\"Return a new :class:`_engine.Connection` object.\u001b[39;00m\n\u001b[1;32m 3256\u001b[0m \n\u001b[1;32m 3257\u001b[0m \u001b[38;5;124;03m The :class:`_engine.Connection` acts as a Python context manager, so\u001b[39;00m\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 3274\u001b[0m \n\u001b[1;32m 3275\u001b[0m \u001b[38;5;124;03m \"\"\"\u001b[39;00m\n\u001b[0;32m-> 3277\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_connection_cls\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/engine/base.py:145\u001b[0m, in \u001b[0;36mConnection.__init__\u001b[0;34m(self, engine, connection, _has_events, _allow_revalidate, _allow_autobegin)\u001b[0m\n\u001b[1;32m 143\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_dbapi_connection \u001b[38;5;241m=\u001b[39m engine\u001b[38;5;241m.\u001b[39mraw_connection()\n\u001b[1;32m 144\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m dialect\u001b[38;5;241m.\u001b[39mloaded_dbapi\u001b[38;5;241m.\u001b[39mError \u001b[38;5;28;01mas\u001b[39;00m err:\n\u001b[0;32m--> 145\u001b[0m \u001b[43mConnection\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_handle_dbapi_exception_noconnection\u001b[49m\u001b[43m(\u001b[49m\n\u001b[1;32m 146\u001b[0m \u001b[43m \u001b[49m\u001b[43merr\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mdialect\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mengine\u001b[49m\n\u001b[1;32m 147\u001b[0m \u001b[43m \u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 148\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m\n\u001b[1;32m 149\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/engine/base.py:2440\u001b[0m, in \u001b[0;36mConnection._handle_dbapi_exception_noconnection\u001b[0;34m(cls, e, dialect, engine, is_disconnect, invalidate_pool_on_disconnect, is_pre_ping)\u001b[0m\n\u001b[1;32m 2438\u001b[0m \u001b[38;5;28;01melif\u001b[39;00m should_wrap:\n\u001b[1;32m 2439\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m sqlalchemy_exception \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[0;32m-> 2440\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m sqlalchemy_exception\u001b[38;5;241m.\u001b[39mwith_traceback(exc_info[\u001b[38;5;241m2\u001b[39m]) \u001b[38;5;28;01mfrom\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21;01me\u001b[39;00m\n\u001b[1;32m 2441\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m 2442\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m exc_info[\u001b[38;5;241m1\u001b[39m] \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/engine/base.py:143\u001b[0m, in \u001b[0;36mConnection.__init__\u001b[0;34m(self, engine, connection, _has_events, _allow_revalidate, _allow_autobegin)\u001b[0m\n\u001b[1;32m 141\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m connection \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m 142\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 143\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_dbapi_connection \u001b[38;5;241m=\u001b[39m \u001b[43mengine\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mraw_connection\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 144\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m dialect\u001b[38;5;241m.\u001b[39mloaded_dbapi\u001b[38;5;241m.\u001b[39mError \u001b[38;5;28;01mas\u001b[39;00m err:\n\u001b[1;32m 145\u001b[0m Connection\u001b[38;5;241m.\u001b[39m_handle_dbapi_exception_noconnection(\n\u001b[1;32m 146\u001b[0m err, dialect, engine\n\u001b[1;32m 147\u001b[0m )\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/engine/base.py:3301\u001b[0m, in \u001b[0;36mEngine.raw_connection\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 3279\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21mraw_connection\u001b[39m(\u001b[38;5;28mself\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m PoolProxiedConnection:\n\u001b[1;32m 3280\u001b[0m \u001b[38;5;250m \u001b[39m\u001b[38;5;124;03m\"\"\"Return a \"raw\" DBAPI connection from the connection pool.\u001b[39;00m\n\u001b[1;32m 3281\u001b[0m \n\u001b[1;32m 3282\u001b[0m \u001b[38;5;124;03m The returned object is a proxied version of the DBAPI\u001b[39;00m\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 3299\u001b[0m \n\u001b[1;32m 3300\u001b[0m \u001b[38;5;124;03m \"\"\"\u001b[39;00m\n\u001b[0;32m-> 3301\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mpool\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconnect\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/base.py:447\u001b[0m, in \u001b[0;36mPool.connect\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 439\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21mconnect\u001b[39m(\u001b[38;5;28mself\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m PoolProxiedConnection:\n\u001b[1;32m 440\u001b[0m \u001b[38;5;250m \u001b[39m\u001b[38;5;124;03m\"\"\"Return a DBAPI connection from the pool.\u001b[39;00m\n\u001b[1;32m 441\u001b[0m \n\u001b[1;32m 442\u001b[0m \u001b[38;5;124;03m The connection is instrumented such that when its\u001b[39;00m\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 445\u001b[0m \n\u001b[1;32m 446\u001b[0m \u001b[38;5;124;03m \"\"\"\u001b[39;00m\n\u001b[0;32m--> 447\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_ConnectionFairy\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_checkout\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/base.py:1264\u001b[0m, in \u001b[0;36m_ConnectionFairy._checkout\u001b[0;34m(cls, pool, threadconns, fairy)\u001b[0m\n\u001b[1;32m 1256\u001b[0m \u001b[38;5;129m@classmethod\u001b[39m\n\u001b[1;32m 1257\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21m_checkout\u001b[39m(\n\u001b[1;32m 1258\u001b[0m \u001b[38;5;28mcls\u001b[39m,\n\u001b[0;32m (...)\u001b[0m\n\u001b[1;32m 1261\u001b[0m fairy: Optional[_ConnectionFairy] \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[1;32m 1262\u001b[0m ) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m _ConnectionFairy:\n\u001b[1;32m 1263\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m fairy:\n\u001b[0;32m-> 1264\u001b[0m fairy \u001b[38;5;241m=\u001b[39m \u001b[43m_ConnectionRecord\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mcheckout\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpool\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 1266\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m threadconns \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m 1267\u001b[0m threadconns\u001b[38;5;241m.\u001b[39mcurrent \u001b[38;5;241m=\u001b[39m weakref\u001b[38;5;241m.\u001b[39mref(fairy)\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/base.py:711\u001b[0m, in \u001b[0;36m_ConnectionRecord.checkout\u001b[0;34m(cls, pool)\u001b[0m\n\u001b[1;32m 709\u001b[0m rec \u001b[38;5;241m=\u001b[39m cast(_ConnectionRecord, pool\u001b[38;5;241m.\u001b[39m_do_get())\n\u001b[1;32m 710\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[0;32m--> 711\u001b[0m rec \u001b[38;5;241m=\u001b[39m \u001b[43mpool\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_do_get\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 713\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m 714\u001b[0m dbapi_connection \u001b[38;5;241m=\u001b[39m rec\u001b[38;5;241m.\u001b[39mget_connection()\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/impl.py:177\u001b[0m, in \u001b[0;36mQueuePool._do_get\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 175\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_create_connection()\n\u001b[1;32m 176\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m:\n\u001b[0;32m--> 177\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m util\u001b[38;5;241m.\u001b[39msafe_reraise():\n\u001b[1;32m 178\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_dec_overflow()\n\u001b[1;32m 179\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/util/langhelpers.py:224\u001b[0m, in \u001b[0;36msafe_reraise.__exit__\u001b[0;34m(self, type_, value, traceback)\u001b[0m\n\u001b[1;32m 222\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m exc_value \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m 223\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_exc_info \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;66;03m# remove potential circular references\u001b[39;00m\n\u001b[0;32m--> 224\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m exc_value\u001b[38;5;241m.\u001b[39mwith_traceback(exc_tb)\n\u001b[1;32m 225\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m 226\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_exc_info \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;66;03m# remove potential circular references\u001b[39;00m\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/impl.py:175\u001b[0m, in \u001b[0;36mQueuePool._do_get\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 173\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_inc_overflow():\n\u001b[1;32m 174\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[0;32m--> 175\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_create_connection\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 176\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m:\n\u001b[1;32m 177\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m util\u001b[38;5;241m.\u001b[39msafe_reraise():\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/base.py:388\u001b[0m, in \u001b[0;36mPool._create_connection\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 385\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21m_create_connection\u001b[39m(\u001b[38;5;28mself\u001b[39m) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m ConnectionPoolEntry:\n\u001b[1;32m 386\u001b[0m \u001b[38;5;250m \u001b[39m\u001b[38;5;124;03m\"\"\"Called by subclasses to create a new ConnectionRecord.\"\"\"\u001b[39;00m\n\u001b[0;32m--> 388\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43m_ConnectionRecord\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/base.py:673\u001b[0m, in \u001b[0;36m_ConnectionRecord.__init__\u001b[0;34m(self, pool, connect)\u001b[0m\n\u001b[1;32m 671\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m__pool \u001b[38;5;241m=\u001b[39m pool\n\u001b[1;32m 672\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m connect:\n\u001b[0;32m--> 673\u001b[0m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m__connect\u001b[49m\u001b[43m(\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 674\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfinalize_callback \u001b[38;5;241m=\u001b[39m deque()\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/base.py:899\u001b[0m, in \u001b[0;36m_ConnectionRecord.__connect\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 897\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfresh \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mTrue\u001b[39;00m\n\u001b[1;32m 898\u001b[0m \u001b[38;5;28;01mexcept\u001b[39;00m \u001b[38;5;167;01mBaseException\u001b[39;00m \u001b[38;5;28;01mas\u001b[39;00m e:\n\u001b[0;32m--> 899\u001b[0m \u001b[38;5;28;01mwith\u001b[39;00m util\u001b[38;5;241m.\u001b[39msafe_reraise():\n\u001b[1;32m 900\u001b[0m pool\u001b[38;5;241m.\u001b[39mlogger\u001b[38;5;241m.\u001b[39mdebug(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mError on connect(): \u001b[39m\u001b[38;5;132;01m%s\u001b[39;00m\u001b[38;5;124m\"\u001b[39m, e)\n\u001b[1;32m 901\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m 902\u001b[0m \u001b[38;5;66;03m# in SQLAlchemy 1.4 the first_connect event is not used by\u001b[39;00m\n\u001b[1;32m 903\u001b[0m \u001b[38;5;66;03m# the engine, so this will usually not be set\u001b[39;00m\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/util/langhelpers.py:224\u001b[0m, in \u001b[0;36msafe_reraise.__exit__\u001b[0;34m(self, type_, value, traceback)\u001b[0m\n\u001b[1;32m 222\u001b[0m \u001b[38;5;28;01massert\u001b[39;00m exc_value \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[1;32m 223\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_exc_info \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;66;03m# remove potential circular references\u001b[39;00m\n\u001b[0;32m--> 224\u001b[0m \u001b[38;5;28;01mraise\u001b[39;00m exc_value\u001b[38;5;241m.\u001b[39mwith_traceback(exc_tb)\n\u001b[1;32m 225\u001b[0m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[1;32m 226\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39m_exc_info \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;66;03m# remove potential circular references\u001b[39;00m\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/pool/base.py:895\u001b[0m, in \u001b[0;36m_ConnectionRecord.__connect\u001b[0;34m(self)\u001b[0m\n\u001b[1;32m 893\u001b[0m \u001b[38;5;28;01mtry\u001b[39;00m:\n\u001b[1;32m 894\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mstarttime \u001b[38;5;241m=\u001b[39m time\u001b[38;5;241m.\u001b[39mtime()\n\u001b[0;32m--> 895\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mdbapi_connection \u001b[38;5;241m=\u001b[39m connection \u001b[38;5;241m=\u001b[39m \u001b[43mpool\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43m_invoke_creator\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m)\u001b[49m\n\u001b[1;32m 896\u001b[0m pool\u001b[38;5;241m.\u001b[39mlogger\u001b[38;5;241m.\u001b[39mdebug(\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mCreated new connection \u001b[39m\u001b[38;5;132;01m%r\u001b[39;00m\u001b[38;5;124m\"\u001b[39m, connection)\n\u001b[1;32m 897\u001b[0m \u001b[38;5;28mself\u001b[39m\u001b[38;5;241m.\u001b[39mfresh \u001b[38;5;241m=\u001b[39m \u001b[38;5;28;01mTrue\u001b[39;00m\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/engine/create.py:661\u001b[0m, in \u001b[0;36mcreate_engine..connect\u001b[0;34m(connection_record)\u001b[0m\n\u001b[1;32m 658\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m connection \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m 659\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m connection\n\u001b[0;32m--> 661\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mdialect\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconnect\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mcargs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mcparams\u001b[49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/sqlalchemy/engine/default.py:629\u001b[0m, in \u001b[0;36mDefaultDialect.connect\u001b[0;34m(self, *cargs, **cparams)\u001b[0m\n\u001b[1;32m 627\u001b[0m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[38;5;21mconnect\u001b[39m(\u001b[38;5;28mself\u001b[39m, \u001b[38;5;241m*\u001b[39mcargs: Any, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mcparams: Any) \u001b[38;5;241m-\u001b[39m\u001b[38;5;241m>\u001b[39m DBAPIConnection:\n\u001b[1;32m 628\u001b[0m \u001b[38;5;66;03m# inherits the docstring from interfaces.Dialect.connect\u001b[39;00m\n\u001b[0;32m--> 629\u001b[0m \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mloaded_dbapi\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mconnect\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mcargs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mcparams\u001b[49m\u001b[43m)\u001b[49m\n", - "File \u001b[0;32m~/git/collab-environment/.venv-310/lib/python3.10/site-packages/psycopg2/__init__.py:122\u001b[0m, in \u001b[0;36mconnect\u001b[0;34m(dsn, connection_factory, cursor_factory, **kwargs)\u001b[0m\n\u001b[1;32m 119\u001b[0m kwasync[\u001b[38;5;124m'\u001b[39m\u001b[38;5;124masync_\u001b[39m\u001b[38;5;124m'\u001b[39m] \u001b[38;5;241m=\u001b[39m kwargs\u001b[38;5;241m.\u001b[39mpop(\u001b[38;5;124m'\u001b[39m\u001b[38;5;124masync_\u001b[39m\u001b[38;5;124m'\u001b[39m)\n\u001b[1;32m 121\u001b[0m dsn \u001b[38;5;241m=\u001b[39m _ext\u001b[38;5;241m.\u001b[39mmake_dsn(dsn, \u001b[38;5;241m*\u001b[39m\u001b[38;5;241m*\u001b[39mkwargs)\n\u001b[0;32m--> 122\u001b[0m conn \u001b[38;5;241m=\u001b[39m \u001b[43m_connect\u001b[49m\u001b[43m(\u001b[49m\u001b[43mdsn\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mconnection_factory\u001b[49m\u001b[38;5;241;43m=\u001b[39;49m\u001b[43mconnection_factory\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[38;5;241;43m*\u001b[39;49m\u001b[43mkwasync\u001b[49m\u001b[43m)\u001b[49m\n\u001b[1;32m 123\u001b[0m \u001b[38;5;28;01mif\u001b[39;00m cursor_factory \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n\u001b[1;32m 124\u001b[0m conn\u001b[38;5;241m.\u001b[39mcursor_factory \u001b[38;5;241m=\u001b[39m cursor_factory\n", - "\u001b[0;31mOperationalError\u001b[0m: (psycopg2.OperationalError) connection to server at \"34.67.80.127\", port 5432 failed: Operation timed out\n\tIs the server running on that host and accepting TCP/IP connections?\n\n(Background on this error at: https://sqlalche.me/e/20/e3q8)" - ] - } - ], - "source": [ - "# Initialize QueryBackend with custom Cloud SQL config\n", - "query_backend = QueryBackend(config=db_config)\n", - "print(\"✓ QueryBackend initialized with direct Cloud SQL connection\")\n", - "\n", - "# Test connection by listing sessions\n", - "try:\n", - " sessions = query_backend.get_sessions()\n", - " print(f\"\\n✓ Connection successful! Found {len(sessions)} sessions\")\n", - "except Exception as e:\n", - " print(f\"\\n✗ Connection failed: {e}\")\n", - " print(\"\\nTroubleshooting:\")\n", - " print(\" 1. Verify Cloud SQL public IP is correct\")\n", - " print(\" 2. Ensure your IP is in authorized networks:\")\n", - " print(\n", - " \" gcloud sql instances patch INSTANCE_NAME --authorized-networks=YOUR_IP/32\"\n", - " )\n", - " print(\" 3. Check password is correct\")\n", - " print(\" 4. Verify database exists and user has permissions\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": "**Notes on Direct Public IP Connections:**\n\n⚠️ **Google Cloud Recommendation**: Use the Cloud SQL Auth Proxy instead of direct connections for better security.\n\nIf using direct connections:\n1. **Security**: Requires SSL/TLS configuration and authorized networks management\n2. **Public IP**: Get your instance's public IP with:\n ```bash\n gcloud sql instances describe INSTANCE_NAME --format=\"value(ipAddresses[0].ipAddress)\"\n ```\n3. **Firewall**: Must authorize your IP address:\n ```bash\n gcloud sql instances patch INSTANCE_NAME --authorized-networks=YOUR_IP/32\n ```\n4. **SSL**: Should configure SSL certificates for encrypted connections\n\n**Recommended**: See [docs/dashboard/CLOUD_SETUP.md](../dashboard/CLOUD_SETUP.md) for Cloud SQL Auth Proxy setup.\n\n---\n\n**If cells 3-6 above worked, skip cell 8 below and proceed directly to cell 9.**\n\n---\n\n### Default: Connect Using Environment Variables (Recommended with Proxy)" - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Initialize QueryBackend (reads from environment variables)\n", - "# Set DUCKDB_PATH or DB_BACKEND as needed before running\n", - "# Example: os.environ['DUCKDB_PATH'] = '/path/to/your.duckdb'\n", - "\n", - "query_backend = QueryBackend()\n", - "print(\"✓ QueryBackend initialized\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# List available sessions\n", - "sessions = query_backend.get_sessions(\"boids_2d_rollout\")\n", - "print(f\"Found {len(sessions)} sessions:\")\n", - "sessions[[\"session_id\"]]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Select a session to analyze (modify as needed)\n", - "session_id = (\n", - " \"rollout-boid_food_basic_vpluspplus_a_n0_h1_vr0.1_s0_rollout5_selfloops-unknown\"\n", - ")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Query Extended Properties at Session Scope\n", - "\n", - "This uses the same high-level API as the dashboard's `ExtendedPropertiesViewerWidget`." - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Create session-scope QueryScope\n", - "scope = QueryScope.from_session(\n", - " session_id=session_id, agent_type=\"agent\" # Filter to 'agent' type (not 'target')\n", - ")\n", - "\n", - "# Create AnalysisContext (mimics dashboard widget pattern)\n", - "context = AnalysisContext(query_backend=query_backend, scope=scope)\n", - "\n", - "print(f\"✓ Created context with scope: {scope}\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Get available properties for this session\n", - "params = context.get_query_params()\n", - "available_props = query_backend.get_available_properties(**params)\n", - "\n", - "print(f\"\\nAvailable properties ({len(available_props)} total):\")\n", - "available_props[[\"property_id\", \"property_name\", \"data_type\", \"unit\"]]" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Query property distributions (aggregated across all episodes in session)\n", - "distributions = query_backend.get_property_distributions(**params)\n", - "\n", - "print(f\"\\n✓ Loaded {len(distributions)} property observations\")\n", - "print(f\" Properties: {distributions['property_id'].unique().tolist()}\")" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "distributions.head()" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Visualize Selected Properties with Violin Plots" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Select properties to visualize (modify as needed)\n", - "# Example: attention weight properties from GNN models\n", - "selected_properties = [\"attn_weight_food\", \"attn_weight_boid\", \"attn_weight_self\"]\n", - "\n", - "# Filter distributions to selected properties\n", - "filtered_df = distributions[\n", - " distributions[\"property_id\"].isin(selected_properties)\n", - "].copy()\n", - "\n", - "if len(filtered_df) == 0:\n", - " print(f\"⚠ No data found for properties: {selected_properties}\")\n", - " print(f\" Available properties: {distributions['property_id'].unique().tolist()}\")\n", - "else:\n", - " print(\n", - " f\"✓ Filtered to {len(filtered_df)} observations across {len(selected_properties)} properties\"\n", - " )" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Create violin plot using seaborn\n", - "if len(filtered_df) > 0:\n", - " # Set figure size and style\n", - " plt.figure(figsize=(12, 6))\n", - " sns.set_style(\"whitegrid\")\n", - "\n", - " # Create violin plot\n", - " sns.violinplot(data=filtered_df, x=\"property_id\", y=\"value_float\", palette=\"Set2\")\n", - "\n", - " # Customize plot\n", - " plt.title(\n", - " f\"Extended Property Distributions - Session: {session_id}\",\n", - " fontsize=14,\n", - " fontweight=\"bold\",\n", - " )\n", - " plt.xlabel(\"Property\", fontsize=12)\n", - " plt.ylabel(\"Value\", fontsize=12)\n", - " plt.xticks(rotation=45, ha=\"right\")\n", - " plt.tight_layout()\n", - "\n", - " plt.show()\n", - "\n", - " # Print summary statistics\n", - " print(\"\\nSummary Statistics:\")\n", - " summary = filtered_df.groupby(\"property_id\")[\"value_float\"].agg(\n", - " [\"count\", \"mean\", \"std\", \"min\", \"max\"]\n", - " )\n", - " print(summary)" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": [ - "## Cleanup" - ] - }, - { - "cell_type": "code", - "execution_count": null, - "metadata": {}, - "outputs": [], - "source": [ - "# Close database connection\n", - "query_backend.close()\n", - "print(\"✓ Connection closed\")" - ] - }, - { - "cell_type": "markdown", - "metadata": {}, - "source": "## Summary\n\nThis notebook demonstrates:\n\n1. **High-Level API**: Uses the same `QueryBackend` and `AnalysisContext` pattern as dashboard widgets\n2. **Session Scope**: Queries property distributions aggregated across all episodes in a session\n3. **Property Agnostic**: Works with any extended property (attention weights, distances, speeds, etc.)\n4. **Visualization**: Creates violin plots using seaborn for distribution analysis\n5. **Cloud SQL Integration**: Multiple connection approaches with security considerations\n\n### Key Components\n\n- `QueryBackend`: High-level database query interface\n- `QueryScope.from_session()`: Define session-scope analysis\n- `AnalysisContext`: Merges scope + shared parameters\n- `get_property_distributions()`: Query raw property values for distributions\n- `sns.violinplot()`: Visualize property distributions\n\n### Database Connection Options (Priority Order)\n\n**Option 1: Cloud SQL Auth Proxy + Environment Variables (RECOMMENDED)**\n- Start proxy: `./cloud-sql-proxy PROJECT:REGION:INSTANCE --port 5433`\n- Set environment variables to connect via localhost:5433\n- Initialize `QueryBackend()` with no arguments\n- **Benefits**: SSL/TLS encryption, IAM auth, no IP management\n- **Best for**: Production notebooks, shared environments, secure connections\n- See Cell 3 and [docs/dashboard/CLOUD_SETUP.md](../dashboard/CLOUD_SETUP.md)\n\n**Option 2: Local DuckDB (Cell 8)**\n- Initialize `QueryBackend()` with `DUCKDB_PATH` environment variable\n- **Best for**: Local development, testing, offline work\n\n**Option 3: Direct Public IP (Cells 4-6, NOT RECOMMENDED)**\n- ⚠️ Bypasses proxy security features (no automatic SSL/TLS, IAM auth)\n- Requires firewall configuration and authorized networks management\n- Construct `PostgresConfig` with public IP and credentials\n- Initialize `QueryBackend(config=db_config)`\n- **Only for**: Special testing scenarios, legacy systems\n\n### Customization\n\nModify these cells to customize the analysis:\n- Cell 3: Configure Cloud SQL Auth Proxy connection\n- Cells 4-6: Configure direct connection (if absolutely necessary)\n- Cell 9/10: Select different session\n- Cell 17: Change `selected_properties` to analyze different properties\n- Cell 18: Customize plot style, colors, or add statistical overlays\n\n### References\n\n- **Recommended Setup**: [docs/dashboard/CLOUD_SETUP.md](../dashboard/CLOUD_SETUP.md)\n- **Google Cloud SQL Auth Proxy**: https://cloud.google.com/sql/docs/postgres/connect-auth-proxy\n- **Dashboard Widgets**: [collab_env/dashboard/widgets/](../../collab_env/dashboard/widgets/)" - } - ], - "metadata": { - "kernelspec": { - "display_name": ".venv-310", - "language": "python", - "name": "python3" - }, - "language_info": { - "codemirror_mode": { - "name": "ipython", - "version": 3 - }, - "file_extension": ".py", - "mimetype": "text/x-python", - "name": "python", - "nbconvert_exporter": "python", - "pygments_lexer": "ipython3", - "version": "3.10.16" - } - }, - "nbformat": 4, - "nbformat_minor": 4 -} \ No newline at end of file From 2f8ca13ab0bf4ec4cb103fde6dbbed8c7c773ef4 Mon Sep 17 00:00:00 2001 From: Dmitry Batenkov Date: Mon, 23 Feb 2026 13:39:30 -0500 Subject: [PATCH 18/21] load models btn was unavailable on auto-population of the model list --- collab_env/tracking_studio/app.py | 1 + 1 file changed, 1 insertion(+) diff --git a/collab_env/tracking_studio/app.py b/collab_env/tracking_studio/app.py index ab910e84..1b626759 100644 --- a/collab_env/tracking_studio/app.py +++ b/collab_env/tracking_studio/app.py @@ -297,6 +297,7 @@ async def list_rf_models(): rf_version_select.value = versions[0]["version"] rf_version_select.enable() rf_detail_btn.visible = True + enable_load_model_btn() ui.notify( f"Found {len(versions)} versions", type="positive", From 5d27b3d47820ea5186a6550b6acc3f9f8007da7a Mon Sep 17 00:00:00 2001 From: Dmitry Batenkov Date: Mon, 6 Apr 2026 16:12:34 -0400 Subject: [PATCH 19/21] add roboflow project dropdown --- collab_env/tracking_studio/app.py | 42 ++++++---- collab_env/tracking_studio/model_manager.py | 88 +++++++++++++++++++++ 2 files changed, 114 insertions(+), 16 deletions(-) diff --git a/collab_env/tracking_studio/app.py b/collab_env/tracking_studio/app.py index 1b626759..2bc19f6d 100644 --- a/collab_env/tracking_studio/app.py +++ b/collab_env/tracking_studio/app.py @@ -256,11 +256,24 @@ async def handle_upload(e): prefs.get("model_source", "Roboflow") == "Roboflow" ) with rf_container: - rf_project_input = ui.input( + # Populate project dropdown from Roboflow workspace + try: + _rf_project_options = model_manager.list_roboflow_projects() + except Exception as _err: + logger.warning(f"Could not list Roboflow projects: {_err}") + _rf_project_options = [] + + _saved_rf_project = prefs.get("rf_project_id", "") + if _saved_rf_project and _saved_rf_project not in _rf_project_options: + _rf_project_options = [_saved_rf_project, *_rf_project_options] + + rf_project_input = ui.select( label="Project ID", - placeholder="workspace/project", - value=prefs.get("rf_project_id", ""), - ).classes("w-full") + options=_rf_project_options, + value=_saved_rf_project or None, + ).classes("w-full").tooltip( + "Pick a project from your Roboflow workspace" + ) # Store raw version data for detail dialog _rf_versions_raw = {} @@ -272,7 +285,6 @@ async def list_rf_models(): ui.notify("Please enter project ID", type="warning") return try: - rf_list_btn.disable() rf_version_select.options = {} rf_version_select.value = None rf_version_select.disable() @@ -307,8 +319,6 @@ async def list_rf_models(): except Exception as error: logger.error(f"Failed to list models: {error}") ui.notify(f"Error: {error}", type="negative") - finally: - rf_list_btn.enable() def show_version_detail(): """Show full JSON for the selected version in a dialog""" @@ -334,20 +344,20 @@ def show_version_detail(): ) dlg.open() - with ui.row().classes("w-full gap-2 items-center"): - rf_list_btn = ui.button( - "List Models", on_click=list_rf_models - ).props("size=sm color=primary") + with ui.row().classes("w-full gap-2 items-center no-wrap"): + rf_version_select = ui.select( + label="Version", + options=[], + ).classes("flex-grow") + rf_version_select.disable() rf_detail_btn = ui.button( "Details", on_click=show_version_detail ).props("size=sm flat") rf_detail_btn.visible = False - rf_version_select = ui.select( - label="Version", - options=[], - ).classes("w-full") - rf_version_select.disable() + rf_project_input.on( + "update:model-value", lambda _e: list_rf_models() + ) # Custom model upload custom_container = ui.column().classes("w-full mt-2") diff --git a/collab_env/tracking_studio/model_manager.py b/collab_env/tracking_studio/model_manager.py index 9b8650d7..13ace507 100644 --- a/collab_env/tracking_studio/model_manager.py +++ b/collab_env/tracking_studio/model_manager.py @@ -328,6 +328,94 @@ def _load_roboflow_with_pipeline(self, model_id: str): logger.error(error_msg) raise ValueError(error_msg) from e + def list_roboflow_projects(self) -> List[str]: + """ + Query Roboflow API for all projects in the workspace tied to the API key. + + Returns: + List of project IDs in "workspace/project" format, sorted alphabetically. + """ + import requests + + if not self.roboflow_api_key: + raise ValueError("ROBOFLOW_API_KEY not set") + + try: + # Root endpoint with API key returns workspace info (may include + # workspace name and/or a nested workspace object with projects). + root = requests.get( + "https://api.roboflow.com/", + params={"api_key": self.roboflow_api_key}, + timeout=10, + ) + root.raise_for_status() + root_data = root.json() + logger.debug(f"Roboflow root response keys: {list(root_data.keys())}") + + # Collect candidate workspace names from various possible shapes + workspace_names: List[str] = [] + ws_field = root_data.get("workspace") + if isinstance(ws_field, str): + workspace_names.append(ws_field) + elif isinstance(ws_field, dict): + name = ws_field.get("url") or ws_field.get("name") + if name: + workspace_names.append(name) + for w in root_data.get("workspaces", []) or []: + if isinstance(w, str): + workspace_names.append(w) + elif isinstance(w, dict): + name = w.get("url") or w.get("name") + if name: + workspace_names.append(name) + + if not workspace_names: + raise ValueError( + f"Could not resolve any workspace from API key. " + f"Root response: {root_data}" + ) + + project_ids: List[str] = [] + for workspace in workspace_names: + ws = requests.get( + f"https://api.roboflow.com/{workspace}", + params={"api_key": self.roboflow_api_key}, + timeout=10, + ) + ws.raise_for_status() + data = ws.json() + projects = data.get("workspace", {}).get("projects") or data.get( + "projects" + ) or [] + logger.info( + f"Roboflow workspace '{workspace}': {len(projects)} projects" + ) + for p in projects: + if isinstance(p, str): + pid = p + else: + pid = p.get("id") or p.get("url") or p.get("name") or "" + if not pid: + continue + if "/" not in pid: + pid = f"{workspace}/{pid}" + project_ids.append(pid) + + project_ids = sorted(set(project_ids)) + logger.info( + f"Found {len(project_ids)} total Roboflow projects across " + f"{len(workspace_names)} workspace(s)" + ) + return project_ids + except requests.exceptions.HTTPError as e: + error_msg = f"Failed to list Roboflow projects: HTTP {e.response.status_code}" + logger.error(error_msg) + raise ValueError(error_msg) from e + except Exception as e: + error_msg = f"Failed to list Roboflow projects: {str(e)}" + logger.error(error_msg) + raise ValueError(error_msg) from e + def list_roboflow_project_models(self, project_id: str) -> List[dict]: """ Query Roboflow API for available model versions in a project. From 91002aafbeac3942b47a331ddbc1dc3ae4ca48fc Mon Sep 17 00:00:00 2001 From: Dmitry Batenkov Date: Mon, 6 Apr 2026 16:13:07 -0400 Subject: [PATCH 20/21] lint --- collab_env/tracking_studio/app.py | 26 ++++++++++++++------- collab_env/tracking_studio/model_manager.py | 12 ++++++---- 2 files changed, 25 insertions(+), 13 deletions(-) diff --git a/collab_env/tracking_studio/app.py b/collab_env/tracking_studio/app.py index 2bc19f6d..3a698b00 100644 --- a/collab_env/tracking_studio/app.py +++ b/collab_env/tracking_studio/app.py @@ -264,15 +264,23 @@ async def handle_upload(e): _rf_project_options = [] _saved_rf_project = prefs.get("rf_project_id", "") - if _saved_rf_project and _saved_rf_project not in _rf_project_options: - _rf_project_options = [_saved_rf_project, *_rf_project_options] - - rf_project_input = ui.select( - label="Project ID", - options=_rf_project_options, - value=_saved_rf_project or None, - ).classes("w-full").tooltip( - "Pick a project from your Roboflow workspace" + if ( + _saved_rf_project + and _saved_rf_project not in _rf_project_options + ): + _rf_project_options = [ + _saved_rf_project, + *_rf_project_options, + ] + + rf_project_input = ( + ui.select( + label="Project ID", + options=_rf_project_options, + value=_saved_rf_project or None, + ) + .classes("w-full") + .tooltip("Pick a project from your Roboflow workspace") ) # Store raw version data for detail dialog diff --git a/collab_env/tracking_studio/model_manager.py b/collab_env/tracking_studio/model_manager.py index 13ace507..d3b1383c 100644 --- a/collab_env/tracking_studio/model_manager.py +++ b/collab_env/tracking_studio/model_manager.py @@ -384,9 +384,11 @@ def list_roboflow_projects(self) -> List[str]: ) ws.raise_for_status() data = ws.json() - projects = data.get("workspace", {}).get("projects") or data.get( - "projects" - ) or [] + projects = ( + data.get("workspace", {}).get("projects") + or data.get("projects") + or [] + ) logger.info( f"Roboflow workspace '{workspace}': {len(projects)} projects" ) @@ -408,7 +410,9 @@ def list_roboflow_projects(self) -> List[str]: ) return project_ids except requests.exceptions.HTTPError as e: - error_msg = f"Failed to list Roboflow projects: HTTP {e.response.status_code}" + error_msg = ( + f"Failed to list Roboflow projects: HTTP {e.response.status_code}" + ) logger.error(error_msg) raise ValueError(error_msg) from e except Exception as e: From d2ff6e1532175f9933d19c7bcd7dbb56c4d87507 Mon Sep 17 00:00:00 2001 From: Dmitry Batenkov Date: Thu, 23 Apr 2026 10:42:55 -0400 Subject: [PATCH 21/21] gnn3d notebook takes too long for CI --- scripts/test_notebooks.sh | 1 + 1 file changed, 1 insertion(+) diff --git a/scripts/test_notebooks.sh b/scripts/test_notebooks.sh index dcca0946..4965fac0 100755 --- a/scripts/test_notebooks.sh +++ b/scripts/test_notebooks.sh @@ -11,6 +11,7 @@ EXCLUDED_NOTEBOOKS=( "docs/alignment/align.ipynb" "docs/alignment/reprojection.ipynb" "docs/tracking/full_pipeline.ipynb" + "docs/gnn/gnn3D/sample.ipynb" ) # Notebooks requiring GCS credentials - excluded when SKIP_GCS_TESTS is set