Update app with explorer

app.py

@@ -1,3 +1,318 @@
+from datetime import datetime
+import json
+import os
+import time
+import random
+import string
+import shutil
+import traceback
+import sys
+from pathlib import Path
+
 import streamlit as st
 
-st.title('Hello World')
+from utils import (
+    load_json,
+    load_json_no_cache,
+    parse_arguments,
+    format_chat_message,
+    find_screenshot,
+    gather_chat_history,
+    get_screenshot,
+    load_page,
+)
+
+
+def show_selectbox(demonstration_dir):
+    # find all the subdirectories in the current directory
+    dirs = [
+        d
+        for d in os.listdir(demonstration_dir)
+        if os.path.isdir(f"{demonstration_dir}/{d}")
+    ]
+
+    if not dirs:
+        st.title("No recordings found.")
+        return None
+
+    # sort by date
+    dirs.sort(key=lambda x: os.path.getmtime(f"{demonstration_dir}/{x}"), reverse=True)
+
+    # offer the user a dropdown to select which recording to visualize, set a default
+    recording_name = st.sidebar.selectbox("Recording", dirs, index=0)
+
+    return recording_name
+
+
+def show_overview(data, recording_name, basedir):
+    st.title(f"Recording {recording_name}")
+
+    screenshot_size = st.session_state.get("screenshot_size_view_mode", "regular")
+    show_advanced_info = st.session_state.get("show_advanced_information", False)
+
+    if screenshot_size == "regular":
+        col_layout = [1.5, 1.5, 7, 3.5]
+    elif screenshot_size == "small":
+        col_layout = [1.5, 1.5, 7, 2]
+    else:  # screenshot_size == 'large'
+        col_layout = [1.5, 1.5, 11]
+
+    # col_i, col_time, col_act, col_actvis = st.columns(col_layout)
+    # screenshots = load_screenshots(data, basedir)
+
+    for i, d in enumerate(data):
+        if i > 0 and show_advanced_info:
+            # Use html to add a horizontal line with minimal gap
+            st.markdown(
+                "<hr style='margin-top: 0.1rem; margin-bottom: 0.1rem;'/>",
+                unsafe_allow_html=True,
+            )
+        if screenshot_size == "large":
+            col_time, col_i, col_act = st.columns(col_layout)
+            col_actvis = col_act
+        else:
+            col_time, col_i, col_act, col_actvis = st.columns(col_layout)
+        secs_from_start = d["timestamp"] - data[0]["timestamp"]
+        # `secs_from_start` is a float including ms, display in MM:SS.mm format
+        col_time.markdown(
+            f"**{datetime.utcfromtimestamp(secs_from_start).strftime('%M:%S')}**"
+        )
+        
+        if not st.session_state.get("enable_html_download", True):
+            col_i.markdown(f"**#{i}**")
+        
+        elif d["type"] == "browser" and (page_filename := d["state"]["page"]):
+            page_path = f"{basedir}/pages/{page_filename}"
+
+            col_i.download_button(
+                label="#" + str(i),
+                data=load_page(page_path),
+                file_name=recording_name + "-" + page_filename,
+                mime="multipart/related",
+                key=f"page{i}",
+            )
+        else:
+            col_i.button(f"#{i}", type='secondary')
+
+        if d["type"] == "chat":
+            col_act.markdown(format_chat_message(d), unsafe_allow_html=True)
+            continue
+
+        # screenshot_filename = d["state"]["screenshot"]
+        img = get_screenshot(d, basedir)
+        arguments = parse_arguments(d["action"])
+        event_type = d["action"]["intent"]
+
+        action_str = f"**{event_type}**({arguments})"
+
+        if img:
+            col_actvis.image(img)
+        
+        col_act.markdown(action_str)
+
+        if show_advanced_info:
+            colors = {
+                "good": "green",
+                "broken": "red",
+                "delayed": "orange"
+            }
+            status = d["state"].get("screenshot_status", "unknown")
+            color = colors.get(status, "blue")
+
+            text = ""
+            text += f'Status: **:{color}[{status.upper()}]**\n\n'
+            text += f'Screenshot: `{d["state"]["screenshot"]}`\\\n'
+            text += f'Page: `{d["state"]["page"]}`\n'
+
+            col_act.markdown(text)
+
+
+def show_slide(data, example_index, recording_name, basedir):
+    columns = st.columns([1, 5], gap="large")
+
+    chat = gather_chat_history(data, example_index)
+    example = data[example_index]
+
+    # display current message
+    if example["type"] == "chat":
+        columns[0].markdown(format_chat_message(example), unsafe_allow_html=True)
+
+    # display chat history
+    with columns[0]:
+        for d in chat:
+            columns[0].markdown(format_chat_message(d), unsafe_allow_html=True)
+
+    img = find_screenshot(data, example_index, basedir)
+
+    if example["type"] == "browser":
+        event_type = example["action"]["intent"]
+        arguments = parse_arguments(example["action"])
+        event_type = example["action"]["intent"]
+
+        action = f"**{event_type}**({arguments})"
+
+        with columns[1]:
+            st.write(action)
+
+    with columns[1]:
+        st.image(img)
+
+
+def show_presentation(data, recording_name, basedir):
+    example_index = st.slider(
+        "example",
+        min_value=0,
+        max_value=len(data) - 1,
+        step=1,
+        label_visibility="hidden",
+    )
+
+    show_slide(data, example_index, recording_name=recording_name, basedir=basedir)
+
+
+def show_video(basedir):
+    # find the mp4 file in the basedir
+    video_filename = None
+    for filename in os.listdir(f"{basedir}"):
+        if filename.endswith(".mp4"):
+            video_filename = filename
+            video_path = f"{basedir}/{video_filename}"
+            st.video(video_path)
+
+    if not video_filename:
+        st.error("No video found")
+
+
+def load_recording(basedir):
+    # Before loading replay, we need a dropdown that allows us to select replay.json or replay_orig.json
+    # Find all files in basedir starting with "replay" and ending with ".json"
+    replay_files = sorted(
+        [
+            f
+            for f in os.listdir(basedir)
+            if f.startswith("replay") and f.endswith(".json")
+        ]
+    )
+    replay_file = st.sidebar.selectbox("Select replay", replay_files, index=0)
+    st.sidebar.checkbox(
+        "Advanced Screenshot Info", False, key="show_advanced_information"
+    )
+    st.sidebar.checkbox(
+        "Enable HTML download", True, key="enable_html_download"
+    )
+    replay_file = replay_file.replace(".json", "")
+
+    metadata = load_json(basedir, "metadata")
+
+    extension_version = metadata["version"]
+    st.sidebar.markdown(f"**extension version**: {extension_version}")
+
+    # convert timestamp to readable date string
+    recording_start_timestamp = metadata["recordingStart"]
+    recording_start_date = datetime.fromtimestamp(
+        int(recording_start_timestamp) / 1000
+    ).strftime("%Y-%m-%d %H:%M:%S")
+    st.sidebar.markdown(f"**started**: {recording_start_date}")
+
+    # recording_end_timestamp = k["recordingEnd"]
+    # calculate duration
+    # duration = int(recording_end_timestamp) - int(recording_start_timestamp)
+    # duration = time.strftime("%M:%S", time.gmtime(duration / 1000))
+
+    # Read in the JSON data
+    replay_dict = load_json_no_cache(basedir, replay_file)
+    form = load_json(basedir, "form")
+
+    duration = replay_dict["data"][-1]["timestamp"] - replay_dict["data"][0]["timestamp"]
+    duration = time.strftime("%M:%S", time.gmtime(duration))
+    st.sidebar.markdown(f"**duration**: {duration}")
+
+    if not replay_dict:
+        return None
+
+    for key in [
+        "annotator",
+        "description",
+        "tasks",
+        "upload_date",
+        "instructor_sees_screen",
+        "uses_ai_generated_output",
+    ]:
+        if form and key in form:
+            if type(form[key]) == list:
+                st.sidebar.markdown(f"**{key}**: {', '.join(form[key])}")
+            else:
+                st.sidebar.markdown(f"**{key}**: {form[key]}")
+
+    st.sidebar.markdown("---")
+    if replay_dict and "status" in replay_dict:
+        st.sidebar.markdown(f"**Validation status**: {replay_dict['status']}")
+    
+    processed_meta_path = Path(basedir).joinpath('processed_metadata.json')
+    start_frame = 'file not found'
+    
+    if processed_meta_path.exists():
+        with open(processed_meta_path) as f:
+            processed_meta = json.load(f)
+        start_frame = processed_meta.get('start_frame', 'info not in file')
+    
+    st.sidebar.markdown(f"**Recording start frame**: {start_frame}")
+    
+    
+    # st.sidebar.button("Delete recording", type="primary", on_click=delete_recording, args=[basedir])
+
+    data = replay_dict["data"]
+    return data
+
+
+def run():
+    mode = st.sidebar.radio("Mode", ["Overview"])
+    demonstration_dir = "./demonstrations"
+
+    params = st.experimental_get_query_params()
+    # list demonstrations/
+    demo_names = os.listdir(demonstration_dir)
+
+    if params.get("recording"):
+        recording_name = params["recording"][0]
+
+    else:
+        recording_name = demo_names[0]
+    
+    recording_name = st.sidebar.selectbox(
+        "Recordings",
+        demo_names,
+        index=demo_names.index(recording_name),
+    )
+
+    if recording_name != params.get("recording", [None])[0]:
+        st.experimental_set_query_params(recording=recording_name)
+
+    if mode == "Overview":
+        with st.sidebar:
+            # Want a dropdown
+            st.selectbox(
+                "Screenshot size",
+                ["small", "regular", "large"],
+                index=1,
+                key="screenshot_size_view_mode",
+            )
+
+    if recording_name is not None:
+        basedir = f"{demonstration_dir}/{recording_name}"
+        data = load_recording(basedir=basedir)
+
+        if not data:
+            st.stop()
+
+        if mode == "Overview":
+            show_overview(data, recording_name=recording_name, basedir=basedir)
+        elif mode == "Presentation":
+            show_presentation(data, recording_name=recording_name, basedir=basedir)
+        elif mode == "Video":
+            show_video(basedir=basedir)
+
+
+if __name__ == "__main__":
+    st.set_page_config(layout="wide")
+    run()

utils.py

@@ -0,0 +1,356 @@
+#!/usr/bin/env python3
+
+from datetime import datetime
+import os
+import json
+from pathlib import Path
+import sys
+import shutil
+import time
+import traceback
+
+import pandas as pd
+import streamlit as st
+
+import json
+from PIL import Image, ImageDraw
+
+
+CACHE_TTL = 60 * 60 * 24 * 14
+
+"""
+Streamlit app utilities
+"""
+
+
+@st.cache_data(ttl=CACHE_TTL)
+def load_json(basedir, name):
+    if not os.path.exists(f"{basedir}/{name}.json"):
+        return None
+
+    with open(f"{basedir}/{name}.json", "r") as f:
+        j = json.load(f)
+
+    return j
+
+
+def load_json_no_cache(basedir, name):
+    if not os.path.exists(f"{basedir}/{name}.json"):
+        return None
+
+    with open(f"{basedir}/{name}.json", "r") as f:
+        j = json.load(f)
+
+    return j
+
+
+def save_json(basedir, name, data):
+    with open(f"{basedir}/{name}.json", "w") as f:
+        json.dump(data, f, indent=4)
+
+
+@st.cache_data
+def load_image(image_file):
+    img = Image.open(image_file)
+    return img
+
+
+@st.cache_resource
+def load_page(page_path):
+    return open(page_path, "rb")
+
+
+def shorten(s):
+    # shorten to 100 characters
+    if len(s) > 100:
+        s = s[:100] + "..."
+
+    return s
+
+
+@st.cache_data
+def parse_arguments(action):
+    s = []
+    event_type = action["intent"]
+    args = action["arguments"]
+
+    if event_type == "textInput":
+        txt = args["text"]
+
+        txt = txt.strip()
+
+        # escape markdown characters
+        txt = txt.replace("_", "\\_")
+        txt = txt.replace("*", "\\*")
+        txt = txt.replace("`", "\\`")
+        txt = txt.replace("$", "\\$")
+
+        txt = shorten(txt)
+
+        s.append(f'"{txt}"')
+    elif event_type == "change":
+        s.append(f'{args["value"]}')
+    elif event_type == "load":
+        url = args["properties"].get("url") or args.get("url")
+        short_url = shorten(url)
+        s.append(f'"[{short_url}]({url})"')
+
+        if args["properties"].get("transitionType"):
+            s.append(f'*{args["properties"]["transitionType"]}*')
+            s.append(f'*{" ".join(args["properties"]["transitionQualifiers"])}*')
+    elif event_type == "scroll":
+        s.append(f'{args["scrollX"]}, {args["scrollY"]}')
+    elif event_type == "say":
+        s.append(f'"{args["text"]}"')
+    elif event_type == "copy":
+        selected = shorten(args["selected"])
+        s.append(f'"{selected}"')
+    elif event_type == "paste":
+        pasted = shorten(args["pasted"])
+        s.append(f'"{pasted}"')
+    elif event_type == "tabcreate":
+        s.append(f'{args["properties"]["tabId"]}')
+    elif event_type == "tabremove":
+        s.append(f'{args["properties"]["tabId"]}')
+    elif event_type == "tabswitch":
+        s.append(
+            f'{args["properties"]["tabIdOrigin"]} -> {args["properties"]["tabId"]}'
+        )
+
+    if args.get("element"):
+        
+        if event_type == 'click':
+            x = round(args['metadata']['mouseX'], 1)
+            y = round(args['metadata']['mouseY'], 1)
+            uid = args.get('element', {}).get('attributes', {}).get("data-webtasks-id")
+            s.append(f"*x =* {x}, *y =* {y}, *uid =* {uid}")
+        else:
+            top = round(args["element"]["bbox"]["top"], 1)
+            left = round(args["element"]["bbox"]["left"], 1)
+            right = round(args["element"]["bbox"]["right"], 1)
+            bottom = round(args["element"]["bbox"]["bottom"], 1)
+            
+            s.append(f"*top =* {top}, *left =* {left}, *right =* {right}, *bottom =* {bottom}")
+
+    return ", ".join(s)
+
+
+@st.cache_resource(max_entries=50_000, ttl=CACHE_TTL)
+def create_visualization(_img, event_type, bbox, x, y, screenshot_path):
+    # screenshot_path is not used, but we need it for caching since we can't cache
+    # PIL images (hence the leading underscore in the variable name to indicate
+    # that it's not hashed)
+    _img = _img.convert("RGBA")
+    draw = ImageDraw.Draw(_img)
+
+    # draw a bounding box around the element
+    color = {
+        "click": "red",
+        "hover": "orange",
+        "textInput": "blue",
+        "change": "green",
+    }[event_type]
+
+    left = bbox["left"]
+    top = bbox["top"]
+    w = bbox["width"]
+    h = bbox["height"]
+    draw.rectangle((left, top, left + w, top + h), outline=color, width=2)
+
+    if event_type in ["click", "hover"]:
+        r = 15
+        for i in range(1, 5):
+            rx = r * i
+            draw.ellipse((x - rx, y - rx, x + rx, y + rx), outline=color, width=3)
+        draw.ellipse((x - r, y - r, x + r, y + r), fill=color)
+
+    return _img
+
+
+@st.cache_data(max_entries=50_000, ttl=CACHE_TTL)
+def get_screenshot_minimal(screenshot_path, event_type, bbox, x, y, new_width=None):
+    img = load_image(screenshot_path)
+    # vis = None
+
+    if event_type in ["click", "textInput", "change", "hover"]:
+        img = create_visualization(img, event_type, bbox, x, y, screenshot_path)
+
+    if new_width is not None:
+        # Resize to 800px wide
+        w, h = img.size
+        new_w = new_width
+        new_h = int(new_w * h / w)
+        img = img.resize((new_w, new_h))
+        print(f"Resized '{screenshot_path}' to", new_w, new_h)
+
+    return img
+
+
+def get_event_info(d):
+    event_type = d["action"]["intent"]
+
+    try:
+        bbox = d["action"]["arguments"]["element"]["bbox"]
+    except KeyError:
+        bbox = None
+
+    try:
+        x = d["action"]["arguments"]["properties"]["x"]
+        y = d["action"]["arguments"]["properties"]["y"]
+    except KeyError:
+        x = None
+        y = None
+
+    return event_type, bbox, x, y
+
+
+def get_screenshot(d, basedir, new_width=None):
+    screenshot_filename = d["state"]["screenshot"]
+
+    if not screenshot_filename:
+        return None
+
+    event_type, bbox, x, y = get_event_info(d)
+    screenshot_path = f"{basedir}/screenshots/{screenshot_filename}"
+
+    return get_screenshot_minimal(
+        screenshot_path, event_type, bbox, x, y, new_width=new_width
+    )
+
+
+def text_bubble(text, color):
+    text = text.replace("\n", "<br>").replace("\t", "&nbsp;" * 8)
+    return f'<div style="background-color:{color}; padding: 8px; margin: 6px; border-radius:10px; display:inline-block;">{text}</div>'
+
+
+def gather_chat_history(data, example_index):
+    chat = []
+    for i, d in enumerate(data):
+        if d["type"] == "chat":
+            if i >= example_index:
+                break
+            chat.append(d)
+
+    # # leave out just 5 last messages
+    # if len(chat) > 5:
+    #     chat = chat[-5:]
+
+    return reversed(chat)
+
+
+def format_chat_message(d):
+    if d["speaker"] == "instructor":
+        return text_bubble("🧑 " + d["utterance"], "rgba(63, 111, 255, 0.35)")
+    else:
+        return text_bubble("🤖 " + d["utterance"], "rgba(185,185,185,0.35)")
+
+
+def find_screenshot(data, example_index, basedir):
+    # keep looking at previous screenshots until we find one
+    # if there is none, return None
+
+    for i in range(example_index, -1, -1):
+        d = data[i]
+        if d["type"] == "chat":
+            continue
+
+        screenshot = get_screenshot(d, basedir)
+        if screenshot:
+            return screenshot
+
+    return None
+
+
+def create_visualization_2(_img, bbox, color, width, x, y):
+    _img = _img.convert("RGBA")
+    draw = ImageDraw.Draw(_img)
+
+    if bbox:
+        left = bbox["left"]
+        top = bbox["top"]
+        w = bbox["width"]
+        h = bbox["height"]
+        draw.rectangle((left, top, left + w, top + h), outline=color, width=width)
+
+    if x and y:
+        r = 8
+        for i in range(1, 4):
+            rx = r * i
+            draw.ellipse((x - rx, y - rx, x + rx, y + rx), outline=color, width=2)
+        draw.ellipse((x - r, y - r, x + r, y + r), fill=color)
+
+    return _img
+
+
+def rescale_bbox(bbox, scaling_factor):
+    return {
+        k: bbox[k] * scaling_factor
+        for k in ["top", "left", "width", "height", "right", "bottom"]
+        if k in bbox
+    }
+
+
+def show_overlay(
+    _img,
+    pred,
+    ref,
+    turn_args,
+    turn_metadata,
+    scale_pred=True,
+    show=("pred_coords", "ref", "pred_elem"),
+):
+    scaling_factor = turn_metadata.get("zoomLevel", 1.0)
+
+    if "pred_elem" in show:
+        # First, draw red box around predicted element
+        if pred.get("element") and pred["element"].get("bbox"):
+            # rescale the bbox by scaling_factor
+            bbox = rescale_bbox(pred["element"]["bbox"], scaling_factor)
+            _img = create_visualization_2(
+                _img, bbox, color="red", width=9, x=None, y=None
+            )
+
+    if "ref" in show:
+        # Finally, draw a blue box around the reference element (if it exists)
+        if ref.get("element") and ref["element"].get("bbox"):
+            # rescale the bbox
+            bbox = rescale_bbox(ref["element"]["bbox"], scaling_factor)
+            x = turn_args.get("properties", {}).get("x")
+            y = turn_args.get("properties", {}).get("y")
+            _img = create_visualization_2(_img, bbox, color="blue", width=6, x=x, y=y)
+
+    if "pred_coords" in show:
+        # Second draw a green box and x/y coordinate based on predicted coordinates
+        # The predicted coordinates are the raw output of the model,
+        # Whereas the predicted element is the inferred element from the predicted coordinates
+        if pred["args"].get("x") and pred["args"].get("y"):
+            x = pred["args"]["x"]
+            y = pred["args"]["y"]
+
+            if scale_pred:
+                x = x * scaling_factor
+                y = y * scaling_factor
+        else:
+            x = None
+            y = None
+
+        # If the predicted element is a bounding box, draw a green box around it
+        if all(c in pred["args"] for c in ["top", "left", "right", "bottom"]):
+            bbox = {
+                "top": pred["args"]["top"],
+                "left": pred["args"]["left"],
+                "width": (pred["args"]["right"] - pred["args"]["left"]),
+                "height": (pred["args"]["bottom"] - pred["args"]["top"]),
+                "right": pred["args"]["right"],
+                "bottom": pred["args"]["bottom"],
+            }
+
+            if scale_pred:
+                bbox = rescale_bbox(bbox, scaling_factor)
+        else:
+            # Otherwise, do nothing
+            bbox = None
+
+        _img = create_visualization_2(_img, bbox=bbox, color="green", width=3, x=x, y=y)
+
+    return _img