gemini_output_visualisation.py

import gradio as gr
import json
import os
import re

# List of JSON files you want to expose in the dropdown.
JSON_FILES = [
    "prompts_outputs/prompt2_output_with_uid.json",
    "prompts_outputs/prompt3_output_with_uid.json",
    "prompts_outputs/prompt4_output_with_uid.json",
    "prompts_outputs/prompt5_output_with_uid.json",
    "prompts_outputs/prompt5_whole_episode_with_uid.json"
]

# You can define a default JSON file and a default UID that you know exists in that file:
DEFAULT_JSON = "prompts_outputs/prompt2_output_with_uid.json"
DEFAULT_UID = "03aee888-17e3-4e76-bb2d-ec12f59190be"  # <-- Make sure this is a valid key in the default file

# Global variable to temporarily store the data from the selected file.
# This will map uid -> { "timestamp": "...", "type": "...", "description": "...", ... }
data_dict = {}


def get_prompt_text_filename(json_file: str) -> str:
    """
    Given a JSON filename like 'prompt2_output_with_uid.json',
    return the corresponding text filename, e.g. 'prompt_2.txt'.

    You can adapt the logic to your naming convention if needed.
    """
    match = re.search(r"(prompt)(\d+)", json_file)
    if match:
        # e.g. 'prompt' + '2'
        number_str = match.group(2)  # '2'
        return f"prompt{number_str}.txt"
    else:
        return ""  # If no match, return empty or handle differently.


def load_json_and_text(json_file):
    """
    1) Load the JSON file (dict of uid->segment_info).
    2) Return a gr.update(...) object for the UID dropdown (choices, default value).
    3) Load the corresponding prompt text file and return its contents.
    """
    global data_dict

    # -------------- Load the JSON --------------
    if not os.path.exists(json_file):
        data_dict = {}
        uid_dropdown_update = gr.update(choices=[], value=None)
        prompt_text = "No JSON file found."
        return uid_dropdown_update, prompt_text

    with open(json_file, "r", encoding="utf-8") as f:
        data = json.load(f)

    data_dict = data
    new_uids = list(data_dict.keys())
    default_uid = new_uids[0] if new_uids else None

    uid_dropdown_update = gr.update(choices=new_uids, value=default_uid)

    # -------------- Load the Prompt Text --------------
    text_filename = get_prompt_text_filename(json_file)
    if text_filename and os.path.exists(text_filename):
        with open(text_filename, "r", encoding="utf-8") as tf:
            prompt_text = tf.read()
    else:
        prompt_text = f"No corresponding text file found for {json_file}."

    return uid_dropdown_update, prompt_text


def show_segment(uid, json_file):
    """
    Given the uid, display the corresponding video, description, and comments.
    Assumes the video file is named segment_{uid}.mp4, or you could store
    the actual file path in your dictionary if you prefer.
    """
    if uid not in data_dict:
        return gr.update(value=None), gr.update(value=""), gr.update(value="")

    segment = data_dict[uid]
    video_path = f"./gemini_shorts/segment_{uid}.mp4"

    description = segment.get("description", "")
    comments = segment.get("comments", "")

    # Return: (video_component, description_text, comments_text)
    return video_path, description, comments


def update_comments(uid, json_file, new_comments):
    """
    Updates the 'comments' field in the chosen JSON for the selected uid.
    Saves it back to the same JSON file.
    """
    if not os.path.exists(json_file):
        return "Error: File not found."

    if uid not in data_dict:
        return "Error: UID not in data."

    # Update local data_dict
    data_dict[uid]["comments"] = new_comments

    # Update on disk
    with open(json_file, "r", encoding="utf-8") as f:
        original_data = json.load(f)

    if uid in original_data:
        original_data[uid]["comments"] = new_comments
    else:
        # If somehow missing, create or handle error
        original_data[uid] = {"comments": new_comments}

    with open(json_file, "w", encoding="utf-8") as f:
        json.dump(original_data, f, indent=2)

    return "Comments saved successfully!"


########################################
# Pre-load the default JSON to set up default UIDs & text
########################################
if os.path.exists(DEFAULT_JSON):
    with open(DEFAULT_JSON, "r", encoding="utf-8") as f:
        data_dict = json.load(f)
    initial_uids = list(data_dict.keys())
    default_uid_value = initial_uids[0] if initial_uids else None
else:
    data_dict = {}
    initial_uids = []
    default_uid_value = None

def_text_filename = get_prompt_text_filename(DEFAULT_JSON)
if def_text_filename and os.path.exists(def_text_filename):
    with open(def_text_filename, "r", encoding="utf-8") as tf:
        default_prompt_text = tf.read()
else:
    default_prompt_text = f"No corresponding text file found for {DEFAULT_JSON}."


# Build the Gradio interface
with gr.Blocks() as demo:
    with gr.Row():
        # Left Column
        with gr.Column():
            description_out = gr.Textbox(
                label="Description",
                interactive=False
            )
            comments_box = gr.Textbox(
                label="Comments"
            )
            save_button = gr.Button("Save Comments")
            status_text = gr.Textbox(
                label="Status",
                interactive=False
            )
            prompt_textbox = gr.Textbox(
                label="Prompt Text File",
                value=default_prompt_text,
                lines=8,
                interactive=False
            )

        # Right Column
        with gr.Column():
            json_file_dropdown = gr.Dropdown(
                choices=JSON_FILES,
                value=DEFAULT_JSON,
                label="Select JSON File"
            )
            uid_dropdown = gr.Dropdown(
                choices=initial_uids,
                value=default_uid_value,
                label="Select UID"
            )
            video_out = gr.Video(label="Video Segment")

    # Define interaction flows

    # When a JSON file is selected, load_json_and_text returns:
    #  (a) new UID dropdown update
    #  (b) text file contents
    json_file_dropdown.change(
        fn=load_json_and_text,
        inputs=[json_file_dropdown],
        outputs=[uid_dropdown, prompt_textbox]
    )

    # When a UID is selected, show_segment updates the video, description, and comments boxes
    uid_dropdown.change(
        fn=show_segment,
        inputs=[uid_dropdown, json_file_dropdown],
        outputs=[video_out, description_out, comments_box]
    )

    # When "Save Comments" is pressed, update_comments writes to JSON and returns a status
    save_button.click(
        fn=update_comments,
        inputs=[uid_dropdown, json_file_dropdown, comments_box],
        outputs=[status_text]
    )

if __name__ == "__main__":
    demo.launch(share=True)