Update app.py

app.py

@@ -10,7 +10,6 @@ import io
 import uuid
 import time
 import random
-from typing import List, Dict
 
 DATA_DIR = "/data"
 IMAGES_DIR = os.path.join(DATA_DIR, "images")
@@ -78,10 +77,10 @@ class DatasetBuilder:
         os.makedirs(IMAGES_DIR, exist_ok=True)
         self.hf_token = os.getenv("HF_Token")  # Access the token from the environment variable
 
-    def get_dataset_file(self) -> str:
+    def get_dataset_file(self):
         return os.path.join(DATA_DIR, f"{self.dataset_name}.json")
 
-    def load_dataset(self) -> List[Dict]:
+    def load_dataset(self):
         dataset_file = self.get_dataset_file()
         if os.path.exists(dataset_file):
             with open(dataset_file, 'r') as f:
@@ -137,14 +136,6 @@ class DatasetBuilder:
     def get_dataset_info(self):
         return f"Current dataset size ({self.dataset_name}): {len(self.dataset)} images"
 
-    def get_storage_usage(self) -> float:
-        total_size = 0
-        for dirpath, dirnames, filenames in os.walk(DATA_DIR):
-            for f in filenames:
-                fp = os.path.join(dirpath, f)
-                total_size += os.path.getsize(fp)
-        return total_size / (1024 * 1024 * 1024)  # Convert to GB
-
     def get_dataset_preview(self, num_images=5):
         preview = []
         for item in self.dataset[-num_images:]:
@@ -152,35 +143,6 @@ class DatasetBuilder:
             preview.append((image_path, item['tags']))
         return preview
 
-    def get_dataset_with_paths(self) -> List[Dict]:
-        return [
-            {
-                'image': os.path.join(IMAGES_DIR, item['image']),
-                'tags': item['tags'],
-                'filename': item['image']
-            }
-            for item in self.dataset
-        ]      
-
-    def remove_image(self, filename):
-        for item in self.dataset:
-            if item['image'] == filename:
-                self.dataset.remove(item)
-                image_path = os.path.join(IMAGES_DIR, filename)
-                if os.path.exists(image_path):
-                    os.remove(image_path)
-                self.save_dataset()
-                return f"Removed image: {filename}"
-        return f"Image not found: {filename}"
-
-    def edit_tags(self, filename, new_tags):
-        for item in self.dataset:
-            if item['image'] == filename:
-                item['tags'] = new_tags
-                self.save_dataset()
-                return f"Updated tags for image: {filename}"
-        return f"Image not found: {filename}"
-    
     def upload_to_huggingface(self, private=True):
         if not self.dataset:
             return "Dataset is empty. Add some images first."
@@ -235,90 +197,44 @@ def create_huggingface_dataset(dataset_name):
 
 def view_dataset(dataset_name):
     builder = DatasetBuilder(dataset_name)
-    return builder.get_dataset_preview(num_images=20)
+    return builder.get_dataset_preview(num_images=60)
 
 def upload_huggingface_dataset(dataset_name, privacy):
     builder = DatasetBuilder(dataset_name)
-    return builder.upload_to_huggingface(private=privacy) 
-
-def view_dataset_for_editing(dataset_name: str) -> List[Dict]:
-    builder = DatasetBuilder(dataset_name)
-    return builder.get_dataset_with_paths()
-
-def remove_image_from_dataset(dataset_name, filename):
-    builder = DatasetBuilder(dataset_name)
-    result = builder.remove_image(filename)
-    return result, builder.get_dataset_with_paths()
-
-def edit_image_tags(dataset_name, filename, new_tags):
-    builder = DatasetBuilder(dataset_name)
-    result = builder.edit_tags(filename, new_tags)
-    return result, builder.get_dataset_with_paths()  
-
-def get_storage_usage(dataset_name: str) -> str:
-    builder = DatasetBuilder(dataset_name)
-    usage = builder.get_storage_usage()
-    return f"Current storage usage: {usage:.2f} GB / 20 GB"
+    return builder.upload_to_huggingface(private=privacy)    
 
 # Create Gradio interface
 with gr.Blocks(theme="huggingface") as iface:
     gr.Markdown("# Image Dataset Builder")
-    
-    dataset_name_input = gr.Textbox(lines=1, label="Dataset Name", placeholder="Enter dataset name...", value="default_dataset")
-    
-    with gr.Tab("Add Images"):
-        gr.Markdown("Enter a URL to add an image and its tags to the dataset. Progress is saved automatically.")
-        with gr.Row():
-            url_input = gr.Textbox(lines=2, label="URL", placeholder="Enter image URL here...")
-            cookies_input = gr.Textbox(lines=2, label="Cookies (optional)", placeholder="Enter cookies")
+    gr.Markdown("Enter a URL to add an image and its tags to the dataset. Progress is saved automatically.")
+
+    with gr.Row():
+        dataset_name_input = gr.Textbox(lines=1, label="Dataset Name", placeholder="Enter dataset name...", value="default_dataset")
+        url_input = gr.Textbox(lines=2, label="URL", placeholder="Enter image URL here...")
+        cookies_input = gr.Textbox(lines=2, label="Cookies (optional)", placeholder="Enter cookies")
         add_button = gr.Button("Add Image")
-        result_output = gr.Textbox(label="Result")
-        dataset_info = gr.Textbox(label="Dataset Info")
-        preview_gallery = gr.Gallery(label="Recent Additions", show_label=False, elem_id="preview_gallery", columns=5, rows=1, height="auto")
-        
-        add_button.click(add_image_to_dataset, inputs=[url_input, cookies_input, dataset_name_input], outputs=[result_output, dataset_info, preview_gallery])
-
-    with gr.Tab("Edit Dataset"):
-        gr.Markdown("## View and Edit Dataset")
-        view_dataset_button = gr.Button("View Dataset")
-        
-        with gr.Column():
-            edit_gallery = gr.Gallery(label="Dataset Images", show_label=True, elem_id="edit_gallery", columns=4, rows=3, height="auto", object_fit="contain")
-            selected_image = gr.State(None)
-            
-            with gr.Row():
-                remove_button = gr.Button("Remove Selected Image")
-                edit_tags_input = gr.Textbox(lines=1, label="Edit Tags", placeholder="Enter new tags, comma-separated...")
-                edit_button = gr.Button("Update Tags")
-            
-            edit_result_output = gr.Textbox(label="Edit Result")
-        
-        def select_image(evt: gr.SelectData, images):
-            selected = images[evt.index]
-            return selected['filename'], selected['tags']
-        
-        edit_gallery.select(select_image, inputs=[edit_gallery], outputs=[selected_image, edit_tags_input])
-        view_dataset_button.click(view_dataset_for_editing, inputs=[dataset_name_input], outputs=[edit_gallery])
-        remove_button.click(remove_image_from_dataset, inputs=[dataset_name_input, selected_image], outputs=[edit_result_output, edit_gallery])
-        edit_button.click(edit_image_tags, inputs=[dataset_name_input, selected_image, edit_tags_input], outputs=[edit_result_output, edit_gallery])
-
-    with gr.Tab("HuggingFace Integration"):
-        gr.Markdown("## Create and Upload HuggingFace Dataset")
-        create_hf_button = gr.Button("Create HuggingFace Dataset")
-        hf_result = gr.Textbox(label="Dataset Creation Result")
-        create_hf_button.click(create_huggingface_dataset, inputs=[dataset_name_input], outputs=hf_result)
-
-        gr.Markdown("## Upload Dataset to Hugging Face")
-        privacy_radio = gr.Radio(choices=["private", "public"], value="private", label="Repository Privacy")
-        upload_hf_button = gr.Button("Upload to Hugging Face")
-        hf_upload_result = gr.Textbox(label="Upload Result")
-        upload_hf_button.click(upload_huggingface_dataset, inputs=[dataset_name_input, privacy_radio], outputs=hf_upload_result)
-
-    with gr.Tab("Storage"):
-        gr.Markdown("## Storage Usage")
-        storage_usage_button = gr.Button("Check Storage Usage")
-        storage_usage_output = gr.Textbox(label="Storage Usage")
-        storage_usage_button.click(get_storage_usage, inputs=[dataset_name_input], outputs=storage_usage_output)
+
+    result_output = gr.Textbox(label="Result")
+    dataset_info = gr.Textbox(label="Dataset Info")
+
+    gr.Markdown("## Dataset Preview")
+    preview_gallery = gr.Gallery(label="Recent Additions", show_label=False, elem_id="preview_gallery", columns=5, rows=1, height="auto")
+
+    add_button.click(add_image_to_dataset, inputs=[url_input, cookies_input, dataset_name_input], outputs=[result_output, dataset_info, preview_gallery])
+
+    create_hf_button = gr.Button("Create HuggingFace Dataset")
+    hf_result = gr.Textbox(label="Dataset Creation Result")
+    create_hf_button.click(create_huggingface_dataset, inputs=[dataset_name_input], outputs=hf_result)
+
+    view_dataset_button = gr.Button("View Dataset")
+    dataset_gallery = gr.Gallery(label="Dataset Contents", show_label=False, elem_id="dataset_gallery", columns=5, rows=4, height="auto")
+    view_dataset_button.click(view_dataset, inputs=[dataset_name_input], outputs=dataset_gallery)
+
+    gr.Markdown("## Upload Dataset to Hugging Face")
+    privacy_radio = gr.Radio(choices=["private", "public"], value="private", label="Repository Privacy")
+    upload_hf_button = gr.Button("Upload to Hugging Face")
+    hf_upload_result = gr.Textbox(label="Upload Result")
+    upload_hf_button.click(upload_huggingface_dataset, inputs=[dataset_name_input, privacy_radio], outputs=hf_upload_result)
 
 # Launch the interface
 iface.launch()