Update src/streamlit_app.py

src/streamlit_app.py

@@ -2,39 +2,444 @@ import altair as alt
 import numpy as np
 import pandas as pd
 import streamlit as st
+import os
+import io
+import json
+import numpy as np
+from PIL import Image
+import requests
+from segments import SegmentsClient
 
 """
-# Welcome to Streamlit!
+# Copy images from one frame to other frames in the same sample
+
+This HF-application first updates selected frames for a given sample UUID by
+replacing each annotation’s id with its track_id and updating the segmentation bitmap
+accordingly (to avoid potential conflicts). Only the frames specified by the source or target
+frame numbers are processed.
+    
+Then it copies annotations from one source frame to one or more target frames. When copying:
+   - The annotations from the source frame are merged with the target's annotations,
+     adding only those from the source that are not already present (based on id).
+   - The segmentation bitmap is merged: for each pixel, if the target's r-value is 0, the corresponding
+     r-value from the source is used.
+   - After merging, the bitmap is scanned for unique r-values and any annotation in the target
+     frame whose id is not present in these unique values is deleted.
+Finally, the updated datalabel is uploaded. 
 
-Edit `/streamlit_app.py` to customize this app to your heart's desire :heart:.
-If you have any questions, checkout our [documentation](https://docs.streamlit.io) and [community
-forums](https://discuss.streamlit.io).
+**Important** 
+If a loop is detected in any frame’s id mappings, or if any track_id in the
+selected frames is greater than 255, no changes are made and nothing is uploaded to Segments.ai.
+For each track_id higher than 255, the warning will include:
+   - Which track_id is too high and the first frame (1-indexed) where it appears.
+   - The lowest available id values (searched across all frames) – one for each offending track_id.
+The user must resolve these issues before updating.
 
-In the meantime, below is an example of what you can do with just a few lines of code:
 """
 
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
-
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
-
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
-
-df = pd.DataFrame({
-    "x": x,
-    "y": y,
-    "idx": indices,
-    "rand": np.random.randn(num_points),
-})
-
-st.altair_chart(alt.Chart(df, height=700, width=700)
-    .mark_point(filled=True)
-    .encode(
-        x=alt.X("x", axis=None),
-        y=alt.Y("y", axis=None),
-        color=alt.Color("idx", legend=None, scale=alt.Scale()),
-        size=alt.Size("rand", legend=None, scale=alt.Scale(range=[1, 150])),
-    ))
+# ---------------- Utility Functions ----------------
+
+def download_image(url: str) -> Image.Image:
+    """Download an image from the given URL and return a PIL Image in RGB mode."""
+    resp = requests.get(url)
+    resp.raise_for_status()
+    return Image.open(io.BytesIO(resp.content)).convert("RGB")
+
+def topological_sort(mapping: dict) -> list:
+    """
+    Given a mapping (original id -> new id) for non-trivial changes,
+    compute a processing order so that if a new id exists among the original ids,
+    it is processed first.
+    Returns the order as a list, or None if a cycle is detected.
+    """
+    nodes = set(mapping.keys())
+    graph = {}
+    for x in mapping:
+        new_id = mapping[x]
+        if new_id in nodes:
+            graph.setdefault(new_id, set()).add(x)
+    visited = {}
+    result = []
+    cycle_found = False
+
+    def dfs(node):
+        nonlocal cycle_found
+        if cycle_found:
+            return
+        if node in visited:
+            if visited[node] == "visiting":
+                cycle_found = True
+            return
+        visited[node] = "visiting"
+        for neighbor in graph.get(node, set()):
+            dfs(neighbor)
+        visited[node] = "visited"
+        result.append(node)
+
+    for node in nodes:
+        if node not in visited:
+            dfs(node)
+    if cycle_found:
+        return None
+    return result[::-1]
+
+def detect_cycle(mapping: dict):
+    """
+    For each chain in the mapping (original id -> new id), follow it.
+    If a cycle is detected, return the first conflicting mapping as a tuple (original, new).
+    Otherwise, return None.
+    """
+    for orig in mapping:
+        visited = set()
+        current = orig
+        while current in mapping:
+            if current in visited:
+                return (orig, mapping[orig])
+            visited.add(current)
+            current = mapping[current]
+    return None
+
+def parse_frame_numbers(frame_str: str):
+    """
+    Parse a string representing frame numbers.
+    Accepts comma-separated values and ranges (e.g., "2,4-6,8").
+    Returns a list of integers (1-indexed).
+    """
+    result = []
+    for part in frame_str.split(","):
+        part = part.strip()
+        if not part:
+            continue
+        if "-" in part:
+            tokens = part.split("-")
+            start = int(tokens[0].strip())
+            end = int(tokens[1].strip())
+            result.extend(range(start, end+1))
+        else:
+            result.append(int(part))
+    return result
+
+# ---------------- Update Ids and Bitmaps ----------------
+
+def update_frame_annotations_and_bitmap(client, frame: dict) -> (int, bool, tuple):
+    """
+    For a given frame, update annotations (set id = track_id) and update the segmentation bitmap.
+    Only non-trivial mappings (where original id != track_id) are processed.
+    The bitmap's R channel is updated accordingly.
+    
+    IMPORTANT: If a loop is detected in the mapping, the function returns immediately without
+    modifying any values in the frame.
+    
+    Returns a tuple:
+       (collision_count, cycle_detected, conflict_pair)
+    """
+    annotations = frame.get("annotations", [])
+    mapping = {}
+    original_ids = set()
+    for ann in annotations:
+        try:
+            orig_id = int(ann.get("id"))
+            new_id = int(ann.get("track_id"))
+            if orig_id != new_id:
+                mapping[orig_id] = new_id
+                original_ids.add(orig_id)
+        except (ValueError, TypeError):
+            continue
+
+    # Check for a cycle first. If a cycle is detected, do not modify the frame.
+    conflict = detect_cycle(mapping)
+    if conflict is not None:
+        return 0, True, conflict
+
+    collision_count = sum(1 for orig, new in mapping.items() if new in original_ids)
+    if mapping:
+        order = topological_sort(mapping)
+        if order is not None:
+            seg_info = frame.get("segmentation_bitmap", {})
+            seg_url = seg_info.get("url")
+            if seg_url:
+                try:
+                    image = download_image(seg_url)
+                    arr = np.array(image)
+                    r_channel = arr[:, :, 0]
+                    for orig in order:
+                        new = mapping[orig]
+                        r_channel[r_channel == orig] = new
+                    arr[:, :, 0] = r_channel
+                    updated_image = Image.fromarray(arr.astype(np.uint8))
+                    buf = io.BytesIO()
+                    updated_image.save(buf, format="PNG")
+                    buf.seek(0)
+                    base = os.path.basename(seg_url)
+                    name, _ = os.path.splitext(base)
+                    new_filename = f"{name}_updated.png"
+                    asset = client.upload_asset(buf, filename=new_filename)
+                    new_url = asset.url
+                    frame.setdefault("segmentation_bitmap", {})["url"] = new_url
+                except Exception:
+                    pass
+    # Update all annotations: set id = track_id.
+    for ann in annotations:
+        ann["id"] = ann.get("track_id")
+    return collision_count, False, None
+
+def update_datalabel(sample_uuid: str, api_key: str, frames_to_update: set, labelset: str = "ground-truth") -> str:
+    """
+    Retrieves the label for the given sample UUID, updates only the specified frames by modifying annotations
+    and updating the segmentation bitmap (via update_frame_annotations_and_bitmap), then uploads
+    the updated datalabel.
+
+    IMPORTANT: If a loop is detected in any processed frame, no changes are applied and nothing is uploaded.
+    The user must resolve the loop before updating Segments.ai.
+
+    frames_to_update: a set of 0-indexed frame indices that should be processed.
+
+    Returns a single summary line describing the operation with extra newlines for readability.
+    """
+    client = SegmentsClient(api_key)
+    try:
+        label = client.get_label(sample_uuid)
+    except Exception as e:
+        return f"Error retrieving label for sample {sample_uuid}: {e}"
+    
+    attributes = label.attributes.model_dump()
+    frames = attributes.get("frames", [])
+    
+    total_collisions = 0
+    conflict_found = None
+    for i, frame in enumerate(frames):
+        if i in frames_to_update:
+            collisions, cycle, conflict_pair = update_frame_annotations_and_bitmap(client, frame)
+            if cycle:
+                conflict_found = conflict_pair
+                break
+            total_collisions += collisions
+
+    if conflict_found is not None:
+        return (f"Error: Cycle detected in annotation id mappings: original id {conflict_found[0]} -> new id {conflict_found[1]}.\n"
+                "Please resolve the loop before updating. No changes have been uploaded to Segments.ai.")
+    
+    try:
+        client.update_label(sample_uuid, labelset=labelset, attributes=attributes)
+    except Exception as e:
+        return f"Error updating label on Segments.ai: {e}"
+    
+    if total_collisions > 0:
+        summary = ("Updated annotation ids with track_ids and updated bitmap for the specified frames.\n"
+                   "Collisions in ids were detected and resolved using logical processing.")
+    else:
+        summary = "Updated annotation ids with track_ids and updated bitmap for the specified frames."
+    
+    return summary
+
+# ---------------- Copy Annotations ----------------
+
+def copy_annotations_to_frames(client, attributes: dict, sample_uuid: str, source_index: int, target_indexes: list) -> str:
+    """
+    Copies annotations from the source frame to each target frame and merges the segmentation bitmap.
+    For each target frame:
+      - The annotations from the source frame are merged with the target's annotations.
+        Only those annotations from the source that are not already present (based on id) are added.
+      - The segmentation bitmap is merged: for each pixel, if the target's r-value is 0, the corresponding
+        r-value from the source is used.
+      - After merging, only annotations whose id appears in the set of unique r-values in the merged bitmap are kept.
+    After processing all target frames, the updated label is uploaded.
+    Returns a single summary line with extra newlines.
+    """
+    frames = attributes.get("frames", [])
+    if source_index < 0 or source_index >= len(frames):
+        return f"Source frame index {source_index+1} is out of range."
+    source_frame = frames[source_index]
+    
+    for tgt in target_indexes:
+        if tgt < 0 or tgt >= len(frames):
+            return f"Target frame index {tgt+1} is out of range."
+        target_frame = frames[tgt]
+        # Merge annotations: keep existing target annotations and add source annotations not already present.
+        target_annotations = target_frame.get("annotations", [])
+        source_annotations = source_frame.get("annotations", [])
+        existing_ids = {ann.get("id") for ann in target_annotations}
+        for ann in source_annotations:
+            if ann.get("id") not in existing_ids:
+                target_annotations.append(ann)
+        target_frame["annotations"] = target_annotations
+        
+        # Merge segmentation bitmaps if both exist.
+        source_seg_url = source_frame.get("segmentation_bitmap", {}).get("url")
+        target_seg_url = target_frame.get("segmentation_bitmap", {}).get("url")
+        if source_seg_url and target_seg_url:
+            try:
+                source_img = download_image(source_seg_url)
+                target_img = download_image(target_seg_url)
+                arr_source = np.array(source_img)
+                arr_target = np.array(target_img)
+                if arr_source.shape != arr_target.shape:
+                    # If shapes differ, simply use source bitmap.
+                    target_frame["segmentation_bitmap"] = source_frame.get("segmentation_bitmap", {})
+                else:
+                    # For pixels where the target's r-value is 0, use the source's pixel.
+                    r_target = arr_target[:, :, 0]
+                    mask = (r_target == 0)
+                    merged_arr = arr_target.copy()
+                    merged_arr[mask] = arr_source[mask]
+                    # Upload the merged image.
+                    merged_img = Image.fromarray(merged_arr.astype(np.uint8))
+                    buf = io.BytesIO()
+                    merged_img.save(buf, format="PNG")
+                    buf.seek(0)
+                    base = os.path.basename(target_seg_url)
+                    name, _ = os.path.splitext(base)
+                    new_filename = f"{name}_merged.png"
+                    asset = client.upload_asset(buf, filename=new_filename)
+                    merged_url = asset.url
+                    target_frame.setdefault("segmentation_bitmap", {})["url"] = merged_url
+                    # Determine unique r-values from the merged bitmap.
+                    unique_vals = set(np.unique(merged_arr[:, :, 0])) - {0}
+                    # Filter annotations: keep only those whose id (as int) is in unique_vals.
+                    filtered_annotations = []
+                    for ann in target_frame.get("annotations", []):
+                        try:
+                            ann_id = int(ann.get("id"))
+                            if ann_id in unique_vals:
+                                filtered_annotations.append(ann)
+                        except (ValueError, TypeError):
+                            pass
+                    target_frame["annotations"] = filtered_annotations
+            except Exception:
+                # If any error occurs during merge, fall back to using the source segmentation bitmap.
+                target_frame["segmentation_bitmap"] = source_frame.get("segmentation_bitmap", {})
+        else:
+            # If one of the bitmaps is missing, use the source's.
+            target_frame["segmentation_bitmap"] = source_frame.get("segmentation_bitmap", {})
+    
+    try:
+        client.update_label(sample_uuid, labelset="ground-truth", attributes=attributes)
+    except Exception as e:
+        return f"Error updating label on Segments.ai during annotation copy: {e}"
+    
+    target_frames_str = ", ".join(str(tgt+1) for tgt in target_indexes)
+    return (f"Annotations merged from frame {source_index+1} into frames {target_frames_str}.\n"
+            f"Bitmap updated with merged r-values and annotations filtered accordingly.")
+
+# ---------------------- Main UI ----------------------
+
+st.title("Copy/Merge Annotations to Target Frames")
+
+# Prompt user for API key as the first input (use type="password" for security if desired)
+api_key = st.text_input("API Key", type="password")
+sample_uuid = st.text_input("Sample UUID", value="")
+source_frame_num = st.number_input("Source Frame Number (1-indexed)", min_value=1, step=1)
+target_frames_str = st.text_input("Target Frame Numbers (comma-separated or range, e.g., '2,4-6')", value="2")
+
+if "result" not in st.session_state:
+    st.session_state["result"] = ""
+if "original_label" not in st.session_state:
+    st.session_state["original_label"] = ""
+if "new_label" not in st.session_state:
+    st.session_state["new_label"] = ""
+
+if st.button("Update and Copy Annotations"):
+    if not api_key:
+        st.error("Please enter your API Key.")
+    else:
+        client = SegmentsClient(api_key)
+        try:
+            orig_label_obj = client.get_label(sample_uuid)
+        except Exception as e:
+            st.error("Error retrieving original label: " + str(e))
+            orig_label_obj = None
+        if orig_label_obj is not None:
+            original_label_json = json.dumps(orig_label_obj.attributes.model_dump(), indent=4)
+            try:
+                target_frames_nums = parse_frame_numbers(target_frames_str)
+            except Exception as e:
+                st.error(f"Error parsing target frame numbers: {e}")
+                target_frames_nums = []
+            source_index = int(source_frame_num) - 1
+            target_indexes = [n - 1 for n in target_frames_nums]
+            # Only update the frames that are in the source or target list.
+            frames_to_update = set([source_index] + target_indexes)
+            
+            # --- Warning Check: Verify no track_id > 255 in selected frames ---
+            attributes = orig_label_obj.attributes.model_dump()
+            frames = attributes.get("frames", [])
+            # Dictionary to record offending track_id -> first (1-indexed) frame number where it appears
+            track_id_warnings = {}
+            # Also accumulate id and track_id values in the selected frames (for reference)
+            selected_existing_values = set()
+            for i in sorted(frames_to_update):
+                if i < 0 or i >= len(frames):
+                    continue
+                frame = frames[i]
+                for ann in frame.get("annotations", []):
+                    try:
+                        t_id = int(ann.get("track_id"))
+                        selected_existing_values.add(t_id)
+                        selected_existing_values.add(int(ann.get("id")))
+                        if t_id > 255 and t_id not in track_id_warnings:
+                            track_id_warnings[t_id] = i + 1  # record first appearance (1-indexed)
+                    except Exception:
+                        continue
+            
+            # Compute available lowest values across ALL frames.
+            all_existing_values = set()
+            for frame in frames:
+                for ann in frame.get("annotations", []):
+                    try:
+                        t_id = int(ann.get("track_id"))
+                        all_existing_values.add(t_id)
+                        all_existing_values.add(int(ann.get("id")))
+                    except Exception:
+                        continue
+            num_offending = len(track_id_warnings)
+            lowest_available_list = []
+            candidate = 1
+            while len(lowest_available_list) < num_offending:
+                if candidate not in all_existing_values:
+                    lowest_available_list.append(candidate)
+                candidate += 1
+
+            if track_id_warnings:
+                warning_message = "Warning: The following track_id values exceed 255:\n"
+                for t_id, frame_no in sorted(track_id_warnings.items()):
+                    warning_message += f" - Track_id {t_id} appears first in frame {frame_no}.\n"
+                warning_message += "\nPlease change these values on Segments.ai before proceeding.\n"
+                warning_message += f"The lowest available id values (across all frames) are: {', '.join(map(str, lowest_available_list))}."
+                st.error(warning_message)
+            else:
+                update_summary = update_datalabel(sample_uuid, api_key, frames_to_update)
+                if update_summary.startswith("Error"):
+                    st.error(update_summary)
+                else:
+                    try:
+                        # Retrieve the label after the update.
+                        label = client.get_label(sample_uuid)
+                    except Exception as e:
+                        st.error("Error retrieving updated label: " + str(e))
+                        label = None
+                    if label is not None:
+                        attributes = label.attributes.model_dump()
+                        copy_summary = copy_annotations_to_frames(client, attributes, sample_uuid, source_index, target_indexes)
+                        final_summary = update_summary + "\n\n" + copy_summary
+                        st.session_state["result"] = final_summary
+                        # Retrieve the final updated label.
+                        try:
+                            label_after = client.get_label(sample_uuid)
+                        except Exception as e:
+                            st.error("Error retrieving final updated label: " + str(e))
+                            label_after = None
+                        if label_after is not None:
+                            new_label_json = json.dumps(label_after.attributes.model_dump(), indent=4)
+                            st.session_state["original_label"] = original_label_json
+                            st.session_state["new_label"] = new_label_json
+
+if st.session_state["result"]:
+    st.text_area("Output", value=st.session_state["result"], height=150)
+    st.download_button("Download Original Label",
+                       data=st.session_state["original_label"],
+                       file_name=f"{sample_uuid}_original.json",
+                       mime="application/json")
+    st.download_button("Download Updated Label",
+                       data=st.session_state["new_label"],
+                       file_name=f"{sample_uuid}_updated.json",
+                       mime="application/json")