napsack 0.1.0__py3-none-any.whl

label/__main__.py

@@ -0,0 +1,210 @@
+from pathlib import Path
+import argparse
+from dotenv import load_dotenv
+
+from label.discovery import discover_sessions, discover_screenshots_sessions, create_single_config
+from label.clients import create_client
+from label.processor import Processor
+from label.visualizer import Visualizer
+
+load_dotenv()
+
+
+def parse_args():
+    p = argparse.ArgumentParser(description="Process session recordings with VLM")
+
+    session_group = p.add_mutually_exclusive_group(required=True)
+    session_group.add_argument("--session", type=Path)
+    session_group.add_argument("--sessions-root", type=Path)
+
+    p.add_argument("--chunk-duration", type=int, default=60, help="Chunk duration in seconds")
+    p.add_argument("--fps", type=int, default=1, help="Frames per second for video processing")
+
+    p.add_argument("--screenshots-only", action="store_true", help="Process screenshots folder only without aggregations or annotations")
+    p.add_argument("--image-extensions", nargs="+", default=[".jpg", ".jpeg", ".png"], help="Image file extensions to consider")
+    p.add_argument("--max-time-gap", type=float, default=300.0, help="Maximum time gap (seconds) between images before forcing a video split (default: 120 = 2 minutes)")
+    p.add_argument("--prompt-file", default=None, help="Path to prompt file (default: prompts/default.txt or prompts/screenshots_only.txt if screenshots only)")
+    p.add_argument("--hash-cache", type=str, default=None, help="Path to hash_cache.json for deduplicating consecutive similar images")
+    p.add_argument("--dedupe-threshold", type=int, default=1, help="Hamming distance threshold for deduplication (drop if <= threshold, default: 1)")
+    p.add_argument("--annotate", action="store_true", help="Annotate videos with cursor positions and clicks (only for standard processing)")
+    p.add_argument("--skip-existing", action="store_true", help="Skip sessions that have already been processed")
+    p.add_argument("--visualize", action="store_true", help="Create annotated video visualizations after processing")
+    p.add_argument("--encode-only", action="store_true", help="Only encode videos (create chunks), skip labeling. Useful for pre-processing before running the full pipeline.")
+
+    p.add_argument("--client", choices=["gemini", "vllm", "bigquery"], default="gemini")
+    p.add_argument("--model", default="")
+    p.add_argument("--encode-workers", type=int, default=8, help="Number of parallel workers for video encoding")
+    p.add_argument("--label-workers", type=int, default=4, help="Number of parallel workers for VLM labeling")
+
+    vllm_group = p.add_argument_group("vLLM Options")
+    vllm_group.add_argument("--vllm-url")
+
+    bq_group = p.add_argument_group("BigQuery Options")
+    bq_group.add_argument("--bq-project", help="GCP project ID for AI Platform endpoint")
+    bq_group.add_argument("--bq-bucket-name", help="GCS bucket name for uploading videos")
+    bq_group.add_argument("--bq-gcs-prefix", default="video_chunks", help="Prefix/folder path in GCS bucket")
+    bq_group.add_argument("--bq-object-table-location", default="us", help="Object table location (e.g., 'us' or 'us.screenomics-gemini')")
+
+    args = p.parse_args()
+
+    if not args.model:
+        if args.client == 'gemini':
+            args.model = 'gemini-3-flash-preview'
+        elif args.client == 'vllm':
+            args.model = 'Qwen/Qwen3-VL-8B-Thinking-FP8'
+        elif args.client == 'bigquery':
+            args.model = 'dataset.model'  # Placeholder - user must provide full model reference
+    if not args.prompt_file:
+        args.prompt_file = "prompts/screenshots_only.txt" if args.screenshots_only else "prompts/default.txt"
+
+    return args
+
+
+def setup_configs(args):
+    if args.session:
+        configs = [create_single_config(
+            args.session,
+            args.chunk_duration,
+            args.screenshots_only,
+            tuple(args.image_extensions),
+        )]
+    else:
+        if args.screenshots_only:
+            configs = discover_screenshots_sessions(
+                args.sessions_root,
+                args.chunk_duration,
+                tuple(args.image_extensions),
+            )
+        else:
+            configs = discover_sessions(
+                args.sessions_root,
+                args.chunk_duration,
+                args.skip_existing,
+            )
+
+        if not configs:
+            print(f"No sessions found in {args.sessions_root}")
+            return []
+
+    return configs
+
+
+def process_with_gemini(args, configs):
+    client = create_client(
+        'gemini',
+        model_name=args.model,
+    )
+
+    processor = Processor(
+        client=client,
+        encode_workers=args.encode_workers,
+        label_workers=args.label_workers,
+        screenshots_only=args.screenshots_only,
+        prompt_file=args.prompt_file,
+        max_time_gap=args.max_time_gap,
+        hash_cache_path=args.hash_cache,
+        dedupe_threshold=args.dedupe_threshold,
+    )
+
+    return processor.process_sessions(
+        configs,
+        fps=args.fps,
+        annotate=args.annotate and not args.screenshots_only,
+        encode_only=args.encode_only,
+    )
+
+
+def process_with_vllm(args, configs):
+    client = create_client(
+        'vllm',
+        base_url=args.vllm_url if args.vllm_url.endswith('/v1') else f"{args.vllm_url}/v1",
+        model_name=args.model
+    )
+
+    processor = Processor(
+        client=client,
+        encode_workers=args.encode_workers,
+        label_workers=args.label_workers,
+        screenshots_only=args.screenshots_only,
+        prompt_file=args.prompt_file,
+        max_time_gap=args.max_time_gap,
+        hash_cache_path=args.hash_cache,
+        dedupe_threshold=args.dedupe_threshold,
+    )
+
+    return processor.process_sessions(
+        configs,
+        fps=args.fps,
+        annotate=args.annotate and not args.screenshots_only,
+        encode_only=args.encode_only,
+    )
+
+
+def process_with_bigquery(args, configs):
+    client = create_client(
+        'bigquery',
+        model_name=args.model,
+        bucket_name=args.bq_bucket_name,
+        gcs_prefix=args.bq_gcs_prefix,
+        object_table_location=args.bq_object_table_location,
+        project_id=args.bq_project,
+    )
+
+    processor = Processor(
+        client=client,
+        encode_workers=args.encode_workers,
+        label_workers=args.label_workers,
+        screenshots_only=args.screenshots_only,
+        prompt_file=args.prompt_file,
+        max_time_gap=args.max_time_gap,
+        hash_cache_path=args.hash_cache,
+        dedupe_threshold=args.dedupe_threshold,
+    )
+
+    return processor.process_sessions(
+        configs,
+        fps=args.fps,
+        annotate=args.annotate and not args.screenshots_only,
+        encode_only=args.encode_only,
+    )
+
+
+def main():
+    args = parse_args()
+
+    configs = setup_configs(args)
+    if not configs:
+        return
+
+    print(f"Processing {len(configs)} sessions")
+
+    if args.client == 'gemini':
+        results = process_with_gemini(args, configs)
+    elif args.client == 'vllm':
+        results = process_with_vllm(args, configs)
+    elif args.client == 'bigquery':
+        results = process_with_bigquery(args, configs)
+    else:
+        raise ValueError(f"Unknown client: {args.client}")
+
+    print(f"✓ Processed {len(results)} sessions")
+
+    if args.visualize:
+        print("\nCreating visualizations...")
+        visualizer = Visualizer(args.annotate)
+
+        for config in configs:
+            if not config.matched_captions_jsonl.exists():
+                print(f"Skipping Visualizing {config.session_id}: no data.jsonl")
+                continue
+
+            try:
+                output = config.session_folder / "annotated.mp4"
+                visualizer.visualize(config.session_folder, output, args.fps)
+                print(f"✓ {config.session_id}: {output}")
+            except Exception as e:
+                print(f"✗ {config.session_id}: {e}")
+
+
+if __name__ == '__main__':
+    main()

label/caption_matching.py

@@ -0,0 +1,150 @@
+from pathlib import Path
+from typing import List, Dict, Any, Optional
+import json
+
+
+def match_captions_with_events(
+    captions_path: Path,
+    aggregations_path: Path,
+    output_path: Path,
+    fps: int = 1
+) -> List[Dict[str, Any]]:
+    """
+    Match captions with aggregated events based on timestamps.
+
+    Args:
+        captions_path: Path to captions.jsonl
+        aggregations_path: Path to aggregations.jsonl
+        output_path: Path to save matched_captions.jsonl
+        fps: Frames per second used in video creation
+
+    Returns:
+        List of matched caption-event objects
+    """
+    # Load captions
+    captions = []
+    with open(captions_path, 'r', encoding='utf-8') as f:
+        for line in f:
+            if line.strip():
+                captions.append(json.loads(line))
+
+    # Load aggregations
+    aggregations = []
+    with open(aggregations_path, 'r', encoding='utf-8') as f:
+        for line in f:
+            if line.strip():
+                agg = json.loads(line)
+                aggregations.append(agg)
+
+    # Sort aggregations by timestamp
+    aggregations.sort(key=lambda x: x.get('timestamp', 0))
+
+    if not aggregations:
+        print("[Matcher] Warning: No aggregations found")
+        return []
+
+    # Get first aggregation timestamp (video start time)
+    first_timestamp = aggregations[0].get('timestamp', 0)
+
+    print(f"[Matcher] Video start time: {first_timestamp}")
+    print(f"[Matcher] Total aggregations: {len(aggregations)}")
+    print(f"[Matcher] FPS: {fps}")
+
+    # Match captions with events
+    matched_data = []
+
+    for caption in captions:
+        # Convert MM:SS to seconds
+        start_seconds = caption['start_seconds']
+        end_seconds = caption['end_seconds']
+
+        # Convert video time to aggregation indices
+        # Each aggregation represents 1 frame, so index = seconds * fps
+        start_index = int(start_seconds * fps)
+        end_index = int(end_seconds * fps)
+
+        # Clamp to valid range
+        start_index = max(0, min(start_index, len(aggregations) - 1))
+        end_index = max(start_index, min(end_index, len(aggregations) - 1))
+
+        print(f"[Matcher] Caption '{caption['caption'][:50]}...' -> indices [{start_index}, {end_index}]")
+
+        # Get aggregations in this range
+        matched_aggs = aggregations[start_index:end_index + 1]
+
+        if not matched_aggs:
+            # No events matched, but still save the caption
+            matched_entry = {
+                "start_time": first_timestamp + start_seconds,
+                "end_time": first_timestamp + end_seconds,
+                "start_index": start_index,
+                "end_index": end_index,
+                "img": None,
+                "caption": caption['caption'],
+                "raw_events": [],
+                "num_aggregations": 0,
+                "start_formatted": caption['start'],
+                "end_formatted": caption['end'],
+            }
+        else:
+            # Get first and last aggregation for time and image
+            first_agg = matched_aggs[0]
+            last_agg = matched_aggs[-1]
+
+            # Concatenate all events from matched aggregations
+            all_events = []
+            for agg in matched_aggs:
+                events = agg.get('events', [])
+                all_events.extend(events)
+
+            matched_entry = {
+                "start_time": first_agg.get('timestamp'),
+                "end_time": last_agg.get('timestamp'),
+                "start_index": start_index,
+                "end_index": end_index,
+                "img": first_agg.get('screenshot_path'),
+                "caption": caption['caption'],
+                "raw_events": all_events,
+                "num_aggregations": len(matched_aggs),
+                "start_formatted": caption['start'],
+                "end_formatted": caption['end'],
+            }
+
+        matched_data.append(matched_entry)
+
+    # Save matched data
+    with open(output_path, 'w', encoding='utf-8') as f:
+        for entry in matched_data:
+            f.write(json.dumps(entry, ensure_ascii=False) + '\n')
+
+    print(f"[Matcher] Saved {len(matched_data)} matched entries to {output_path}")
+
+    return matched_data
+
+
+def create_matched_captions_for_session(session_dir: Path, fps: int = 1) -> Optional[Path]:
+    """
+    Create matched_captions.jsonl for a session directory.
+
+    Args:
+        session_dir: Path to session directory
+        fps: Frames per second used in video creation
+
+    Returns:
+        Path to created matched_captions.jsonl or None if failed
+    """
+    captions_path = session_dir / "captions.jsonl"
+    aggregations_path = session_dir / "aggregations.jsonl"
+    output_path = session_dir / "matched_captions.jsonl"
+
+    if not captions_path.exists():
+        print(f"[Matcher] Warning: {captions_path} not found")
+        return None
+
+    if not aggregations_path.exists():
+        print(f"[Matcher] Warning: {aggregations_path} not found")
+        return None
+
+    match_captions_with_events(captions_path, aggregations_path, output_path, fps)
+
+    return output_path

label/clients/__init__.py

@@ -0,0 +1,28 @@
+from label.clients.client import VLMClient, CAPTION_SCHEMA
+from label.clients.gemini import GeminiClient, GeminiResponse
+from label.clients.vllm import VLLMClient, VLLMResponse
+from label.clients.bigquery import BigQueryClient, BigQueryResponse
+
+
+def create_client(client_type: str, **kwargs) -> VLMClient:
+    if client_type == 'gemini':
+        return GeminiClient(**kwargs)
+    elif client_type == 'vllm':
+        return VLLMClient(**kwargs)
+    elif client_type == 'bigquery':
+        return BigQueryClient(**kwargs)
+    else:
+        raise ValueError(f"Unknown client type: {client_type}")
+
+
+__all__ = [
+    "VLMClient",
+    "GeminiClient",
+    "GeminiResponse",
+    "VLLMClient",
+    "VLLMResponse",
+    "BigQueryClient",
+    "BigQueryResponse",
+    "CAPTION_SCHEMA",
+    "create_client",
+]

label/clients/bigquery.py

@@ -0,0 +1,171 @@
+from typing import Optional, Any, Dict
+import json
+from pathlib import Path
+
+from google.cloud import storage
+from google.cloud import bigquery
+
+from label.clients.client import VLMClient, CAPTION_SCHEMA  # still imported, but we don't rely on it by default
+
+# example call-
+# uv run -m label \
+#   --session /home/jupyter/Omar/downloads/test@gmail.com \
+#   --screenshots-only \
+#   --client bigquery \
+#   --model gemini-3-pro-preview \
+#   --bq-project hs-nero-phi-reeves-haitech \
+#   --bq-bucket-name hs-nero-phi-reeves-haitech-project \
+#   --bq-gcs-prefix Shaikh_Omar \
+#   --bq-object-table-location us.screenomics-gemini
+
+class BigQueryResponse:
+    def __init__(self, result_row):
+        self.result_row = result_row
+        self._json = None
+
+    @property
+    def text(self) -> str:
+        return self.result_row
+
+    @property
+    def json(self):
+        if self._json is None:
+            self._json = json.loads(self.text)
+        return self._json
+
+
+class BigQueryClient(VLMClient):
+    def __init__(
+        self,
+        model_name: str,
+        bucket_name: str,
+        gcs_prefix: str = "video_chunks",
+        object_table_location: str = "us",  # e.g., "us.screenomics-gemini"
+        temperature: float = 0.0,
+        max_output_tokens: int = 65535,
+        project_id: Optional[str] = None,
+    ):
+        """
+        Initialize BigQuery client for ML.GENERATE_TEXT with video analysis.
+
+        Args:
+            model_name: Full BigQuery model reference (e.g., "dataset.model" or "project.dataset.model")
+            bucket_name: GCS bucket name for uploading videos
+            gcs_prefix: Prefix/folder path in GCS bucket
+            object_table_location: Object table location (e.g., "us.screenomics-gemini")
+            temperature: Model temperature parameter
+            max_output_tokens: Maximum number of tokens in the generated response
+            project_id: Optional GCP project ID (if not provided, uses default credentials)
+        """
+        self.model_name = model_name
+        self.bucket_name = bucket_name
+        self.gcs_prefix = gcs_prefix
+        self.object_table_location = object_table_location
+        self.temperature = temperature
+        self.max_output_tokens = max_output_tokens
+        self.project_id = project_id
+
+        self.storage_client = storage.Client(project=project_id)
+        self.bq_client = bigquery.Client(project=project_id)
+
+    @staticmethod
+    def _escape_for_bq_single_quoted_string(s: str) -> str:
+        r"""
+        Escape a Python string for use in a BigQuery single-quoted string literal.
+
+        BigQuery-style escaping:
+        - Backslash:    \  -> \\
+        - Newline:      actual newline -> \n  (two chars)
+        - Carriage ret: actual \r      -> \r
+        - Single quote: '  -> \'
+        """
+        # Escape backslashes first so we don't re-escape ones we add later
+        s = s.replace("\\", "\\\\")
+        # Encode newlines and carriage returns as literal escape sequences
+        s = s.replace("\r", "\\r")
+        s = s.replace("\n", "\\n")
+        # Escape single quotes for BigQuery
+        s = s.replace("'", "\\'")
+        return s
+
+    def upload_file(self, path: str, session_id: str = None) -> str:
+        """
+        Upload file to GCS and return the GCS URI.
+
+        Args:
+            path: Local file path
+            session_id: Optional session identifier for namespacing uploads
+
+        Returns:
+            GCS URI (gs://bucket/path/to/file)
+        """
+        file_path = Path(path)
+        if session_id:
+            destination_blob_name = f"{self.gcs_prefix}/{session_id}/{file_path.name}"
+        else:
+            destination_blob_name = f"{self.gcs_prefix}/{file_path.name}"
+
+        bucket = self.storage_client.bucket(self.bucket_name)
+        blob = bucket.blob(destination_blob_name)
+
+        # Upload the file
+        blob.upload_from_filename(path)
+
+        gcs_uri = f"gs://{self.bucket_name}/{destination_blob_name}"
+        print(f"Uploaded {path} to {gcs_uri}")
+
+        return gcs_uri
+    
+    def generate(
+        self,
+        prompt: str,
+        file_descriptor: Optional[Any] = None,
+        schema: Optional[Dict] = None,
+    ) -> BigQueryResponse:
+        if not file_descriptor:
+            raise ValueError("file_descriptor (GCS URI) is required")
+
+        gcs_uri = file_descriptor
+
+        # Escape everything that will go inside single-quoted SQL string literals
+        escaped_prompt = self._escape_for_bq_single_quoted_string(prompt)
+        escaped_gcs_uri = self._escape_for_bq_single_quoted_string(gcs_uri)
+        escaped_location = self._escape_for_bq_single_quoted_string(
+            self.object_table_location
+        )
+                
+        response_params = {
+            "generation_config": {
+                "media_resolution": "MEDIA_RESOLUTION_HIGH",
+                "response_mime_type": "application/json"                
+            }
+        }
+        
+        response_params_json = json.dumps(response_params)
+        escaped_response_params = self._escape_for_bq_single_quoted_string(response_params_json)
+        
+        query = f"""
+        SELECT
+          AI.GENERATE(
+            (
+              '{escaped_prompt}',
+              OBJ.FETCH_METADATA(
+                OBJ.MAKE_REF('{escaped_gcs_uri}', '{escaped_location}')
+              )
+            ),
+            endpoint => 'https://aiplatform.googleapis.com/v1/projects/{self.project_id}/locations/global/publishers/google/models/{self.model_name}',
+            model_params => JSON '{escaped_response_params}'
+          ) AS gen;
+        """
+        
+        job_config = bigquery.QueryJobConfig()
+
+        query_job = self.bq_client.query(query, job_config=job_config)
+        results = query_job.result()
+        
+        for row in results:
+            print(row[0]["result"])
+            return BigQueryResponse(row[0]["result"])
+
+        raise RuntimeError("No results returned from BigQuery")
+

label/clients/client.py

@@ -0,0 +1,28 @@
+from __future__ import annotations
+from abc import ABC, abstractmethod
+from typing import Any, Dict, Optional, List, Union
+
+CAPTION_SCHEMA = {
+    "type": "array",
+    "items": {
+        "type": "object",
+        "properties": {
+            "start": {"type": "string"},
+            "end": {"type": "string"},
+            "caption": {"type": "string"}
+        },
+        "required": ["start", "end", "caption"]
+    }
+}
+
+
+class VLMClient(ABC):
+    @abstractmethod
+    def upload_file(self, path: str, session_id: str = None) -> Any:
+        pass
+
+    @abstractmethod
+    def generate(self, prompt: Union[str, List[str]],
+                 file_descriptor: Optional[Union[Any, List[Any]]] = None,
+                 schema: Optional[Dict] = None) -> Union[Any, List[Any]]:
+        pass

label/clients/gemini.py

@@ -0,0 +1,84 @@
+from typing import Optional, Any, Dict
+import json
+import os
+import time
+from label.clients.client import VLMClient, CAPTION_SCHEMA
+
+from google import genai
+from google.genai import types
+
+
+class GeminiResponse:
+    def __init__(self, response):
+        self.response = response
+        self._json = None
+
+    @property
+    def text(self) -> str:
+        return self.response.text
+
+    @property
+    def json(self):
+        if self._json is None:
+            self._json = json.loads(self.text)
+        return self._json
+
+
+class GeminiClient(VLMClient):
+    def __init__(self, api_key: Optional[str] = None, model_name: str = "gemini-2.5-flash"):
+        if genai is None:
+            raise RuntimeError("google-genai not installed")
+
+        api_key = api_key or os.environ.get("GEMINI_API_KEY")
+        if not api_key:
+            raise RuntimeError("GEMINI_API_KEY not set")
+
+        self.client = genai.Client(api_key=api_key)
+        self.model_name = model_name
+
+    def upload_file(self, path: str, session_id: str = None) -> Any:
+        video_file = self.client.files.upload(file=path)
+
+        while True:
+            video_file = self.client.files.get(name=video_file.name)
+            state = getattr(getattr(video_file, "state", None), "name", None)
+
+            if state == "PROCESSING":
+                time.sleep(2)
+            elif state == "FAILED":
+                raise RuntimeError("Gemini failed processing file")
+            elif state == "ACTIVE":
+                break
+            else:
+                break
+
+        return video_file
+
+    def generate(self, prompt: str, file_descriptor: Optional[Any] = None,
+                 schema: Optional[Dict] = None) -> GeminiResponse:
+
+        inputs = [prompt]
+        if file_descriptor:
+            inputs.append(file_descriptor)
+
+        if "gemini-3" in self.model_name:
+            config = types.GenerateContentConfig(
+                response_mime_type="application/json",
+                response_schema=schema or CAPTION_SCHEMA,
+                thinking_config=types.ThinkingConfig(thinking_budget=-1),
+                media_resolution=types.MediaResolution.MEDIA_RESOLUTION_HIGH,
+            )
+        else:
+            config = types.GenerateContentConfig(
+                response_mime_type="application/json",
+                temperature=0.0,
+                response_schema=schema or CAPTION_SCHEMA
+            )
+
+        res = self.client.models.generate_content(
+            model=self.model_name,
+            contents=inputs,
+            config=config
+        )
+
+        return GeminiResponse(res)