Update app.py

app.py

@@ -16,64 +16,97 @@ def save_dataset_to_zip(dataset_name, dataset):
     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
+
+    # Create a zip file with the dataset_name as the top-level folder
     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)
+                rel_file = os.path.relpath(abs_file, temp_dir)
                 zipf.write(abs_file, rel_file)
+
     zip_buffer.seek(0)
     return zip_buffer
 
 def load_dataset_from_zip(zip_file):
     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
+    try:
+        with zipfile.ZipFile(zip_file.name, 'r') as zip_ref:
+            zip_ref.extractall(temp_dir)
+
+        # Get dataset name from zip file name
+        dataset_name_guess = os.path.splitext(os.path.basename(zip_file.name))[0]
+        dataset_path = os.path.join(temp_dir, dataset_name_guess)
+
+        if os.path.exists(dataset_path):
+            dataset_name = dataset_name_guess
+        else:
+            # If the dataset_name directory doesn't exist, try to find the top-level directory
+            entries = [entry for entry in os.listdir(temp_dir) if os.path.isdir(os.path.join(temp_dir, entry))]
+            if entries:
+                dataset_name = entries[0]
+                dataset_path = os.path.join(temp_dir, dataset_name)
+            else:
+                # Files are directly in temp_dir
+                dataset_name = dataset_name_guess
+                dataset_path = temp_dir
+
+        images_dir = os.path.join(dataset_path, 'images')
+        annotations_path = os.path.join(dataset_path, 'annotations.jsonl')
+        dataset = []
+
+        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
+                    })
+        else:
+            # If annotations file not found
+            return None, []
+
+        return dataset_name, dataset
+    except Exception as e:
+        print(f"Error loading dataset: {e}")
+        return None, []
 
 def display_dataset_html(dataset):
     if dataset:
@@ -97,7 +130,7 @@ with gr.Blocks() as demo:
     datasets = gr.State({})
     current_dataset_name = gr.State("")
     dataset_selector = gr.Dropdown(label="Select Dataset", interactive=True)
-    entry_selector = gr.Dropdown(label="Select Entry to Edit/Delete")  # Moved outside
+    entry_selector = gr.Dropdown(label="Select Entry to Edit/Delete")
     dataset_html = gr.HTML()
     message_box = gr.Textbox(interactive=False, label="Message")
 
@@ -109,7 +142,7 @@ with gr.Blocks() as demo:
                 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_input = gr.File(label="Upload Dataset Zip", type="file")
                 upload_button = gr.Button("Upload Dataset")
 
     def create_dataset(name, datasets):
@@ -130,6 +163,8 @@ with gr.Blocks() as demo:
         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 is None:
+            return gr.update(), "Failed to load dataset from zip file."
         if dataset_name in datasets:
             return gr.update(), f"Dataset '{dataset_name}' already exists."
         datasets[dataset_name] = dataset