Update app.py

app.py

@@ -10,16 +10,87 @@ import tempfile
 import numpy as np
 
 def save_dataset_to_zip(dataset_name, dataset):
-    # Function implementation remains the same
-    # ...
+    # Create a temporary directory
+    temp_dir = tempfile.mkdtemp()
+    dataset_path = os.path.join(temp_dir, dataset_name)
+    os.makedirs(dataset_path, exist_ok=True)
+    images_dir = os.path.join(dataset_path, 'images')
+    os.makedirs(images_dir, exist_ok=True)
+    annotations = []
+    for idx, entry in enumerate(dataset):
+        image_data = entry['image']
+        prompt = entry['prompt']
+        # Save image to images directory
+        image_filename = f"{uuid.uuid4().hex}.png"
+        image_path = os.path.join(images_dir, image_filename)
+        # Decode the base64 image data
+        image = Image.open(BytesIO(base64.b64decode(image_data.split(",")[1])))
+        image.save(image_path)
+        # Add annotation
+        annotations.append({
+            'file_name': os.path.join('images', image_filename),
+            'text': prompt
+        })
+    # Save annotations to JSONL file
+    annotations_path = os.path.join(dataset_path, 'annotations.jsonl')
+    with open(annotations_path, 'w') as f:
+        for ann in annotations:
+            f.write(json.dumps(ann) + '\n')
+    # Create a zip file
+    zip_buffer = BytesIO()
+    with zipfile.ZipFile(zip_buffer, 'w', zipfile.ZIP_DEFLATED) as zipf:
+        for root, dirs, files in os.walk(dataset_path):
+            for file in files:
+                abs_file = os.path.join(root, file)
+                rel_file = os.path.relpath(abs_file, dataset_path)
+                zipf.write(abs_file, rel_file)
+    zip_buffer.seek(0)
+    return zip_buffer
 
 def load_dataset_from_zip(zip_file):
-    # Function implementation remains the same
-    # ...
+    temp_dir = tempfile.mkdtemp()
+    with zipfile.ZipFile(zip_file.name, 'r') as zip_ref:
+        zip_ref.extractall(temp_dir)
+    dataset_name = os.listdir(temp_dir)[0]
+    dataset_path = os.path.join(temp_dir, dataset_name)
+    dataset = []
+    images_dir = os.path.join(dataset_path, 'images')
+    annotations_path = os.path.join(dataset_path, 'annotations.jsonl')
+    if os.path.exists(annotations_path):
+        with open(annotations_path, 'r') as f:
+            for line in f:
+                ann = json.loads(line)
+                file_name = ann['file_name']
+                prompt = ann['text']
+                image_path = os.path.join(dataset_path, file_name)
+                # Read image and convert to base64
+                with open(image_path, 'rb') as img_f:
+                    image_bytes = img_f.read()
+                    encoded = base64.b64encode(image_bytes).decode()
+                    mime_type = "image/png"
+                    image_data = f"data:{mime_type};base64,{encoded}"
+                dataset.append({
+                    'image': image_data,
+                    'prompt': prompt
+                })
+    return dataset_name, dataset
 
 def display_dataset_html(dataset):
-    # Function implementation remains the same
-    # ...
+    if dataset:
+        html_content = ""
+        for idx, entry in enumerate(dataset):
+            image_data = entry['image']
+            prompt = entry['prompt']
+            html_content += f"""
+            <div style="display: flex; align-items: center; margin-bottom: 10px;">
+                <div style="width: 50px;">{idx}</div>
+                <img src="{image_data}" alt="Image {idx}" style="max-height: 100px; margin-right: 10px;"/>
+                <div>{prompt}</div>
+            </div>
+            """
+        return html_content
+    else:
+        return "<div>No entries in dataset.</div>"
 
 with gr.Blocks() as demo:
     gr.Markdown("<h1 style='text-align: center; margin-bottom: 20px;'>Dataset Builder</h1>")
@@ -31,12 +102,23 @@ with gr.Blocks() as demo:
     message_box = gr.Textbox(interactive=False, label="Message")
 
     with gr.Tab("Create / Upload Dataset"):
-        # Create / Upload Dataset components and functions
-        # ...
+        with gr.Row():
+            with gr.Column():
+                gr.Markdown("### Create a New Dataset")
+                dataset_name_input = gr.Textbox(label="New Dataset Name")
+                create_button = gr.Button("Create Dataset")
+            with gr.Column():
+                gr.Markdown("### Upload Existing Dataset")
+                upload_input = gr.File(label="Upload Dataset Zip", file_types=['.zip'])
+                upload_button = gr.Button("Upload Dataset")
 
     def create_dataset(name, datasets):
-        # Function implementation remains the same
-        # ...
+        if not name:
+            return gr.update(), "Please enter a dataset name."
+        if name in datasets:
+            return gr.update(), f"Dataset '{name}' already exists."
+        datasets[name] = []
+        return gr.update(choices=list(datasets.keys()), value=name), f"Dataset '{name}' created."
 
     create_button.click(
         create_dataset,
@@ -45,8 +127,13 @@ with gr.Blocks() as demo:
     )
 
     def upload_dataset(zip_file, datasets):
-        # Function implementation remains the same
-        # ...
+        if zip_file is None:
+            return gr.update(), "Please upload a zip file."
+        dataset_name, dataset = load_dataset_from_zip(zip_file)
+        if dataset_name in datasets:
+            return gr.update(), f"Dataset '{dataset_name}' already exists."
+        datasets[dataset_name] = dataset
+        return gr.update(choices=list(datasets.keys()), value=dataset_name), f"Dataset '{dataset_name}' uploaded."
 
     upload_button.click(
         upload_dataset,
@@ -71,8 +158,10 @@ with gr.Blocks() as demo:
     )
 
     with gr.Tab("Add Entry"):
-        # Add Entry components and functions
-        # ...
+        with gr.Row():
+            image_input = gr.Image(label="Upload Image")
+            prompt_input = gr.Textbox(label="Prompt")
+        add_button = gr.Button("Add Entry")
 
     def add_entry(image_data, prompt, current_dataset_name, datasets):
         if not current_dataset_name: