Add functions to extract first frame labels from cliplabels

poseinterface/clips.py

@@ -8,6 +8,36 @@
 
 import sleap_io as sio
 
+try:
+    import boto3
+    from botocore.exceptions import ClientError
+
+    BOTO3_AVAILABLE = True
+except ImportError:
+    BOTO3_AVAILABLE = False
+
+
+def _cliplabels_suffix_error(name: str) -> str:
+    return f"Input file must end with '_cliplabels.json', got {name}"
+
+
+def _extract_start_labels(clip_labels: dict) -> dict:
+    start_images = [img for img in clip_labels["images"] if img["id"] == 0]
+    if len(start_images) != 1:
+        raise ValueError(
+            "Clip labels must contain exactly one first-frame image with id 0"
+        )
+
+    return {
+        "images": start_images,
+        "annotations": [
+            annot
+            for annot in clip_labels["annotations"]
+            if annot["image_id"] == 0
+        ],
+        "categories": clip_labels["categories"],
+    }
+
 
 def extract_clip(
     video_path: str | Path,
@@ -178,6 +208,170 @@ def _extract_cliplabels(
     return clip_json
 
 
+def extract_start_frames_label(cliplabels_path: str | Path) -> Path:
+    """Extract only the first frame's labels from a cliplabels.json file.
+
+    Reads a ``*_cliplabels.json`` file and creates a corresponding
+    ``*_startlabels.json`` file containing only the labels for the first
+    frame (frame with id=0).
+
+    Parameters
+    ----------
+    cliplabels_path
+        Path to the input ``*_cliplabels.json`` file.
+
+    Returns
+    -------
+    Path
+        Path to the output ``*_startlabels.json`` file.
+
+    Raises
+    ------
+    FileNotFoundError
+        If the input file does not exist.
+    ValueError
+        If the input filename does not end with ``_cliplabels.json``.
+    """
+    cliplabels_path = Path(cliplabels_path)
+
+    # Validate input file exists
+    if not cliplabels_path.exists():
+        raise FileNotFoundError(f"Input file not found: {cliplabels_path}")
+
+    # Validate input filename
+    if not cliplabels_path.name.endswith("_cliplabels.json"):
+        raise ValueError(_cliplabels_suffix_error(cliplabels_path.name))
+
+    # Read the cliplabels file
+    with open(cliplabels_path) as f:
+        clip_labels = json.load(f)
+
+    start_labels = _extract_start_labels(clip_labels)
+
+    # Generate output path by replacing _cliplabels.json with _startlabels.json
+    output_path = cliplabels_path.parent / cliplabels_path.name.replace(
+        "_cliplabels.json", "_startlabels.json"
+    )
+
+    # Save the start labels
+    with open(output_path, "w") as f:
+        json.dump(start_labels, f)
+
+    logging.info(
+        f"Extracted start frame labels from {cliplabels_path.name} "
+        f"to {output_path.name}"
+    )
+
+    return output_path
+
+
+def extract_start_frames_label_s3(
+    s3_cliplabels_uri: str,
+    aws_profile: str | None = None,
+) -> str:
+    """Extract first frame's labels from a cliplabels.json file on S3.
+
+    Reads a ``*_cliplabels.json`` file from S3 and creates a corresponding
+    ``*_startlabels.json`` file on S3 containing only the labels for the first
+    frame (frame with id=0).
+
+    Parameters
+    ----------
+    s3_cliplabels_uri
+        S3 URI of the input ``*_cliplabels.json`` file in the format
+        ``s3://bucket-name/path/to/file_cliplabels.json``.
+    aws_profile
+        Optional AWS profile name to use for authentication. If None, uses
+        the default AWS credentials chain.
+
+    Returns
+    -------
+    str
+        S3 URI of the output ``*_startlabels.json`` file.
+
+    Raises
+    ------
+    ImportError
+        If boto3 is not installed.
+    ValueError
+        If the S3 URI format is invalid or filename does not end with
+        ``_cliplabels.json``.
+    ClientError
+        If there are issues accessing S3 (permissions, file not found, etc.).
+    """
+    if not BOTO3_AVAILABLE:
+        raise ImportError(
+            "boto3 is required for S3 operations. "
+            "Install it with: pip install boto3"
+        )
+
+    # Parse S3 URI
+    if not s3_cliplabels_uri.startswith("s3://"):
+        raise ValueError(
+            f"Invalid S3 URI format. Expected 's3://bucket/key', "
+            f"got {s3_cliplabels_uri}"
+        )
+
+    # Extract bucket and key from URI
+    uri_parts = s3_cliplabels_uri[5:].split("/", 1)
+    if len(uri_parts) != 2:
+        raise ValueError(
+            f"Invalid S3 URI format. Expected 's3://bucket/key', "
+            f"got {s3_cliplabels_uri}"
+        )
+    bucket_name, key = uri_parts
+
+    # Validate filename
+    if not key.endswith("_cliplabels.json"):
+        raise ValueError(_cliplabels_suffix_error(key))
+
+    # Initialize S3 client
+    session = boto3.Session(profile_name=aws_profile)
+    s3_client = session.client("s3")
+
+    try:
+        # Download and read the cliplabels file from S3
+        logging.info(f"Downloading {s3_cliplabels_uri}")
+        response = s3_client.get_object(Bucket=bucket_name, Key=key)
+        clip_labels = json.loads(response["Body"].read().decode("utf-8"))
+
+        start_labels = _extract_start_labels(clip_labels)
+
+        # Generate output key by replacing _cliplabels.json with
+        # _startlabels.json
+        output_key = key.replace("_cliplabels.json", "_startlabels.json")
+        output_uri = f"s3://{bucket_name}/{output_key}"
+
+        # Upload the start labels to S3
+        logging.info(f"Uploading to {output_uri}")
+        s3_client.put_object(
+            Bucket=bucket_name,
+            Key=output_key,
+            Body=json.dumps(start_labels),
+            ContentType="application/json",
+        )
+
+        logging.info(
+            f"Extracted start frame labels from {key} to {output_key}"
+        )
+
+        return output_uri
+
+    except ClientError as e:
+        error_code = e.response.get("Error", {}).get("Code", "Unknown")
+        if error_code == "NoSuchKey":
+            raise FileNotFoundError(
+                f"File not found on S3: {s3_cliplabels_uri}"
+            ) from e
+        # Preserve other S3 client errors, but log them for easier diagnosis.
+        logging.error(
+            "Failed to extract start labels from %s (S3 error: %s)",
+            s3_cliplabels_uri,
+            error_code,
+        )
+        raise
+
+
 def main(args: argparse.Namespace) -> None:
     """Run clip extraction from parsed command-line arguments.
 

tests/test_unit/test_clips.py

@@ -7,6 +7,7 @@
 
 from poseinterface.clips import (
     _extract_cliplabels,
+    _extract_start_labels,
     extract_clip,
     main,
     parse_args,
@@ -157,6 +158,37 @@ def test_extract_clip_invalid(
         extract_clip(video_path, start_frame, duration)
 
 
+def test_extract_start_labels(video_labels):
+    """Test start labels contain only the first frame and its annotations."""
+    start_labels = _extract_start_labels(video_labels)
+    first_frame_annotations = [
+        annotation
+        for annotation in video_labels["annotations"]
+        if annotation["image_id"] == 0
+    ]
+
+    assert start_labels["images"] == [{"id": 0}]
+    assert start_labels["annotations"] == first_frame_annotations
+    assert start_labels["categories"] == video_labels["categories"]
+
+
+def test_extract_start_labels_requires_first_frame(video_labels):
+    """Test start-label extraction validates the presence of frame 0."""
+    clip_labels = {
+        **video_labels,
+        "images": [{"id": 1}],
+        "annotations": [{"image_id": 1, "id": 1}],
+    }
+
+    with pytest.raises(
+        ValueError,
+        match=(
+            "Clip labels must contain exactly one first-frame image with id 0"
+        ),
+    ):
+        _extract_start_labels(clip_labels)
+
+
 def test_parse_args():
     """Check arguments are parsed with correct types"""
     video_path = "foo.mp4"