Improved UI with better functionalities and UX.

app.py

@@ -1,4 +1,5 @@
 import os
+import dataclasses
 import io
 from PIL import Image
 import gradio as gr
@@ -11,6 +12,15 @@ os.environ["AWS_SECRET_ACCESS_KEY"] = os.getenv("AWS_SECRET_ACCESS_KEY")
 os.environ["S3_BUCKET_NAME"] = os.getenv("AWS_BUCKET")
 
 
+@dataclasses.dataclass
+class ComicFrame:
+    description: str
+    narration: str
+    character_dilouge: str
+    character: str
+    compositions: list
+
+
 def list_current_dir(bucket_name: str, folder_path: str = "") -> list:
     response = aws_utils.S3_CLIENT.list_objects_v2(
         Bucket=bucket_name, Prefix=folder_path, Delimiter="/"
@@ -24,85 +34,100 @@ def list_current_dir(bucket_name: str, folder_path: str = "") -> list:
     return folders
 
 
-def load_text_data(
-    episodes: list, current_episode: int, current_scene: int, current_frame: int
-):
-    curr_frame = episodes[current_episode]["scenes"][current_scene]["frames"][
-        current_frame
-    ]
-    return (
-        episodes,
-        current_episode,
-        current_scene,
-        current_frame,
-        str(current_episode + 1),
-        str(current_scene + 1),
-        str(current_frame + 1),
-        curr_frame["description"],
-        curr_frame["narration"],
-        curr_frame["audio_cue_character"],
-        curr_frame["audio_cue_text"],
-    )
-
-
-def load_data(
-    episodes: list, current_episode: int, current_scene: int, current_frame: int
-):
-    if current_frame + 1 < len(
-        episodes[current_episode]["scenes"][current_scene]["frames"]
-    ):
-        current_frame += 1
-    else:
-        if current_scene + 1 < len(episodes[current_episode]["scenes"]):
-            current_scene += 1
-            current_frame = 0
-        else:
-            if current_episode + 1 < len(episodes):
-                current_episode += 1
-                current_scene = 0
-                current_frame = 0
-            else:
-                return [], current_episode, current_scene, current_frame
-
+def load_data_inner(episodes_data: list, current_episode: int, current_frame: int):
     images = []
+    curr_frame = episodes_data[current_episode][current_frame]
     # Loading the 0th frame of 0th scene in 0th episode.
-    for comps in episodes[current_episode]["scenes"][current_scene]["frames"][
-        current_frame
-    ]["compositions"]:
-        data = aws_utils.fetch_from_s3(comps["image"])
+    for comps in curr_frame.compositions:
+        data = aws_utils.fetch_from_s3(comps)
         images.append(Image.open(io.BytesIO(data)))
 
-    return images, *load_text_data(
-        episodes, current_episode, current_scene, current_frame
+    return (
+        images,
+        episodes_data,
+        current_episode,
+        current_frame,
+        curr_frame.description,
+        curr_frame.narration,
+        curr_frame.character,
+        curr_frame.character_dilouge,
     )
 
 
-def load_data_once(
-    comic_id: str, current_episode: int, current_scene: int, current_frame: int
-):
+def load_metadata_fn(comic_id: str):
     # Logic to load and return images based on comic_id and episode
     # You can replace this with actual image paths or generation logic
     print(f"Getting episodes for comic id: {comic_id}")
-    episodes = {}
-    idx = 0
+    episodes_data = {}
+    episode_idx = []
     for folder in list_current_dir(AWS_BUCKET, f"{comic_id}/"):
         if "episode" in folder:
             json_path = f"s3://{AWS_BUCKET}/{folder}episode.json"
-            episodes[idx] = eval(
-                aws_utils.fetch_from_s3(source=json_path).decode("utf-8")
-            )
-            idx += 1
+            idx = int(folder.split("/")[1].split("-")[-1])
+            episode_idx.append(idx)
+            data = eval(aws_utils.fetch_from_s3(source=json_path).decode("utf-8"))
+            comic_frames = []
+            for scene in data["scenes"]:
+                for frame in scene["frames"]:
+                    comic_frames.append(
+                        ComicFrame(
+                            description=frame["description"],
+                            narration=frame["narration"],
+                            character=frame["audio_cue_character"],
+                            character_dilouge=frame["audio_cue_text"],
+                            compositions=[
+                                comp["image"] for comp in frame["compositions"]
+                            ],
+                        )
+                    )
+            episodes_data[idx] = comic_frames
+    current_episode, current_frame = min(episode_idx), 0
+    return (
+        gr.update(choices=episode_idx, value=episode_idx[0]),
+        gr.update(
+            choices=range(len(episodes_data[current_episode])), value=current_frame
+        ),
+        episodes_data
+    )
 
-    images = []
-    # Loading the 0th frame of 0th scene in 0th episode.
-    for comps in episodes[current_episode]["scenes"][current_scene]["frames"][
-        current_frame
-    ]["compositions"]:
-        data = aws_utils.fetch_from_s3(comps["image"])
-        images.append(Image.open(io.BytesIO(data)))
 
-    return images, *load_text_data(
-        episodes, current_episode, current_scene, current_frame
+def load_data_next(episodes_data: list, current_episode: int, current_frame: int):
+    if current_frame + 1 < len(episodes_data[current_episode]):
+        current_frame += 1
+    elif current_episode + 1 < len(episodes_data):
+        current_episode += 1
+        current_frame = 0
+    else:
+        return [], current_episode, current_frame
+    return (
+        gr.update(value=current_episode),
+        gr.update(value=current_frame),
+        *load_data_inner(episodes_data, current_episode, current_frame),
+    )
+
+
+def load_data_prev(episodes_data: list, current_episode: int, current_frame: int):
+    if current_frame - 1 >= 0:
+        current_frame -= 1
+    elif current_episode - 1 > min(list(episodes_data.keys())):
+        current_episode -= 1
+        current_frame = 0
+    else:
+        return [], current_episode, current_frame
+    return (
+        gr.update(value=current_episode),
+        gr.update(value=current_frame),
+        *load_data_inner(episodes_data, current_episode, current_frame),
+    )
+
+
+def load_from_dropdown(
+    episodes_data: dict, selected_episode, selected_frame,
+):
+    return (
+        gr.update(value=selected_episode),
+        gr.update(value=selected_frame),
+        *load_data_inner(episodes_data, selected_episode, selected_frame),
     )
 
 
@@ -110,11 +135,10 @@ def save_image(
     selected_image,
     comic_id: str,
     current_episode: int,
-    current_scene: int,
     current_frame: int,
 ):
     # Implement your AWS S3 save logic here
-    print(f"Saving image: {selected_image}")
+    # print(f"Saving image: {selected_image}")
     with Image.open(selected_image[0]) as img:
         # Convert and save as JPG
         img_bytes = io.BytesIO()
@@ -123,58 +147,68 @@ def save_image(
 
         aws_utils.save_to_s3(
             AWS_BUCKET,
-            f"{comic_id}/episode-{current_episode}/images/scene-{current_scene}",
+            f"{comic_id}/episode-{current_episode}/images",
             img_bytes,
             f"{current_frame}.jpg",
         )
-    print("Image saved successfully!")
     gr.Info("Saved Image successfully!")
 
 
 with gr.Blocks() as demo:
     selected_image = gr.State(None)
-    current_episode = gr.State(0)
-    current_scene = gr.State(0)
-    current_frame = gr.State(0)
+    current_episode = gr.State(-1)
+    current_frame = gr.State(-1)
     episodes_data = gr.State({})
 
     with gr.Row():
         comic_id = gr.Textbox(label="Enter Comic ID:", placeholder="Enter Comic ID")
-        load_button = gr.Button("Load Data")
+        load_metadata = gr.Button("Load Metadata")
+
+    # Display information about current Image
+    with gr.Row():
+        episode_dropdown = gr.Dropdown(choices=[], label="Current Episode", interactive=True)
+        frame_dropdown = gr.Dropdown(choices=[], label="Current Frame", interactive=True)
+        load_images = gr.Button("Load Images")
 
     images = gr.Gallery(
         label="Select an Image", elem_id="image_select", columns=4, height=300
     )
 
-    # Display information about current Image
     with gr.Row():
-        episode = gr.Textbox(label="Current Episode", interactive=False)
-        scene = gr.Textbox(label="Current Scene", interactive=False)
-        frame = gr.Textbox(label="Current Frame", interactive=False)
+        image_description = gr.Textbox(label="Description", interactive=False)
+        narration = gr.Textbox(label="narration", interactive=False)
 
-    image_description = gr.Textbox(label="Description", interactive=False)
-    narration = gr.Textbox(label="narration", interactive=False)
     with gr.Row():
         character = gr.Textbox(label="Character", interactive=False)
         dialouge = gr.Textbox(label="dialouge", interactive=False)
 
     # buttons to interact with the data
     with gr.Row():
+        prev_button = gr.Button("Prev Image")
         save_button = gr.Button("Save Image")
         next_button = gr.Button("Next Image")
 
-    load_button.click(
-        load_data_once,
-        inputs=[comic_id, current_episode, current_scene, current_frame],
+    load_metadata.click(
+        load_metadata_fn,
+        inputs=[comic_id],
         outputs=[
+            episode_dropdown,
+            frame_dropdown,
+            episodes_data
+        ],
+    )
+
+    load_images.click(
+        load_from_dropdown,
+        inputs=[episodes_data, episode_dropdown, frame_dropdown],
+        outputs=[
+            episode_dropdown,
+            frame_dropdown,
+            # Textual data returned from load_text_data()
             images,
             episodes_data,
             current_episode,
-            current_scene,
             current_frame,
-            episode,
-            scene,
-            frame,
             image_description,
             narration,
             character,
@@ -188,30 +222,35 @@ with gr.Blocks() as demo:
 
     images.select(get_select_index, images, selected_image)
 
-    save_button.click(
-        save_image,
-        inputs=[
-            selected_image,
-            comic_id,
+    next_button.click(
+        load_data_next,
+        inputs=[episodes_data, current_episode, current_frame],
+        outputs=[
+            episode_dropdown,
+            frame_dropdown,
+            # Textual data returned from load_text_data()
+            images,
+            episodes_data,
             current_episode,
-            current_scene,
             current_frame,
+            image_description,
+            narration,
+            character,
+            dialouge
         ],
-        outputs=[],
     )
 
-    next_button.click(
-        load_data,
-        inputs=[episodes_data, current_episode, current_scene, current_frame],
+    prev_button.click(
+        load_data_prev,
+        inputs=[episodes_data, current_episode, current_frame],
         outputs=[
+            episode_dropdown,
+            frame_dropdown,
+            # Textual data returned from load_text_data()
             images,
             episodes_data,
             current_episode,
-            current_scene,
             current_frame,
-            episode,
-            scene,
-            frame,
             image_description,
             narration,
             character,
@@ -219,4 +258,15 @@ with gr.Blocks() as demo:
         ],
     )
 
+    save_button.click(
+        save_image,
+        inputs=[
+            selected_image,
+            comic_id,
+            current_episode,
+            current_frame,
+        ],
+        outputs=[],
+    )
+
 demo.launch(auth=("admin", "Qrt@12*34#immersfy"), share=True, ssr_mode=False)