Update app.py

app.py

@@ -93,6 +93,25 @@ class DatasetBuilder:
         with open(dataset_file, 'w') as f:
             json.dump(self.dataset, f)
 
+    def resize_images(self, min_size=512, max_size=768):
+        for item in self.dataset:
+            image_path = os.path.join(IMAGES_DIR, item['image'])
+            image = Image.open(image_path)
+            
+            # Resize the image while maintaining the aspect ratio
+            image.thumbnail((max_size, max_size), resample=Image.BICUBIC)
+            
+            # Save the resized image
+            image.save(image_path)
+
+    def resize_dataset(self):
+        resized_dataset_name = f"{self.dataset_name} (resized)"
+        resized_dataset_builder = DatasetBuilder(resized_dataset_name)
+        resized_dataset_builder.dataset = self.dataset
+        resized_dataset_builder.resize_images()
+        resized_dataset_builder.save_dataset()
+        return f"Resized dataset '{self.dataset_name}' to '{resized_dataset_name}'."
+
     def create_downloadable_dataset(self):
         if not self.dataset:
             return None, "Dataset is empty. Add some images first."
@@ -232,6 +251,15 @@ def download_dataset(dataset_name):
     zip_path, message = builder.create_downloadable_dataset()
     return zip_path, message
 
+def resize_dataset(dataset_name):
+    builder = DatasetBuilder(dataset_name)
+    return builder.resize_dataset()
+
+def download_resized_dataset(dataset_name):
+    builder = DatasetBuilder(f"{dataset_name} (resized)")
+    zip_path, message = builder.create_downloadable_dataset()
+    return zip_path, message
+
 # Create Gradio interface
 with gr.Blocks(theme="huggingface") as iface:
     gr.Markdown("# Image Dataset Builder")
@@ -276,5 +304,24 @@ with gr.Blocks(theme="huggingface") as iface:
         outputs=[download_output, download_message]
     )
 
+    gr.Markdown("## Resize Dataset")
+    resize_button = gr.Button("Resize Dataset")
+    resize_result = gr.Textbox(label="Resize Result")
+    resize_button.click(
+        resize_dataset,
+        inputs=[dataset_name_input],
+        outputs=resize_result
+    )
+
+    gr.Markdown("## Download Resized Dataset")
+    download_resized_button = gr.Button("Download Resized Dataset")
+    download_resized_output = gr.File(label="Download Resized")
+    download_resized_message = gr.Textbox(label="Resized Download Status")
+    download_resized_button.click(
+        download_resized_dataset,
+        inputs=[dataset_name_input],
+        outputs=[download_resized_output, download_resized_message]
+    )
+
 # Launch the interface
 iface.launch()