Update src/streamlit_app.py

src/streamlit_app.py

@@ -1,40 +1,68 @@
-import altair as alt
-import numpy as np
-import pandas as pd
 import streamlit as st
+from segments import SegmentsClient
+import sys
+import copy
+import os
+sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../../0_label_scripts")))
+from get_labels_from_samples import export_all_sensor_frames_and_annotations
 
-"""
-# Welcome to Streamlit!
+# ---------------- Streamlit UI ----------------
+st.title("Multi-Sensor: Overwrite Target Frame Annotations with Source Frame (Cuboids)")
 
-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).
+api_key = st.text_input("API Key", type="password")
+source_uuid = st.text_input("Source UUID", value="")
+target_uuid = st.text_input("Target UUID", value="")
+source_frame_num = st.number_input("Source Frame Number (1-indexed)", min_value=1, value=50, step=1)
+target_frame_num = st.number_input("Target Frame Number (1-indexed)", min_value=1, value=1, step=1)
 
-In the meantime, below is an example of what you can do with just a few lines of code:
-"""
+if st.button("Overwrite Target Frame Annotations"):
+    if not api_key or not source_uuid or not target_uuid:
+        st.error("Please fill in all fields.")
+    else:
+        try:
+            client = SegmentsClient(api_key)
+            # Fetch source and target labels
+            source_label = client.get_label(source_uuid)
+            target_label = client.get_label(target_uuid)
 
-num_points = st.slider("Number of points in spiral", 1, 10000, 1100)
-num_turns = st.slider("Number of turns in spiral", 1, 300, 31)
+            # Extract all sensor frames/annotations
+            source_sensors = export_all_sensor_frames_and_annotations(source_label)
+            target_sensors = export_all_sensor_frames_and_annotations(target_label)
 
-indices = np.linspace(0, 1, num_points)
-theta = 2 * np.pi * num_turns * indices
-radius = indices
+            # Defensive copy
+            new_target_label = copy.deepcopy(target_label)
+            sensors_attr = getattr(new_target_label.attributes, "sensors", None)
+            if sensors_attr is None:
+                st.error("Target label has no sensors.")
+                st.stop()
 
-x = radius * np.cos(theta)
-y = radius * np.sin(theta)
+            # For each sensor in both source and target
+            for sensor_idx, sensor in enumerate(sensors_attr):
+                sensor_name = getattr(sensor, "name", None)
+                if not sensor_name or sensor_name not in source_sensors:
+                    continue  # skip sensors not found in source
+                source_frames = source_sensors[sensor_name]
+                target_frames = target_sensors.get(sensor_name, [])
+                # Frame indices are 1-indexed for user, 0-indexed in list
+                src_idx = source_frame_num - 1
+                tgt_idx = target_frame_num - 1
+                if src_idx >= len(source_frames) or tgt_idx >= len(target_frames):
+                    st.warning(f"Sensor '{sensor_name}': Frame index out of range. Skipped.")
+                    continue
+                # Overwrite target frame's annotations with source frame's annotations
+                src_anns = source_frames[src_idx]["annotations"]
+                tgt_frame = getattr(sensor.attributes.frames[tgt_idx], "annotations", None)
+                if tgt_frame is not None:
+                    sensor.attributes.frames[tgt_idx].annotations = copy.deepcopy(src_anns)
+                else:
+                    st.warning(f"Sensor '{sensor_name}': Could not access target frame annotations.")
 
-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])),
-    ))
+            # Upload the updated label to Segments.ai (only the specified frames/annotations are changed)
+            client.update_label(
+                target_uuid,
+                labelset="ground-truth",  # Change this if you use a different labelset name
+                attributes=new_target_label.attributes.model_dump()
+            )
+            st.success("Target label updated and uploaded to Segments.ai!")
+        except Exception as e:
+            st.error(f"Error: {e}")