Update app.py

app.py

@@ -14,10 +14,10 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 class DatasetManager:
-    def __init__(self, dataset_name=None, local_images_dir="downloaded_cards"):
-        self.dataset_name = dataset_name
+    def __init__(self, local_images_dir="downloaded_cards"):
         self.local_images_dir = local_images_dir
         self.drive = None
+        self.dataset_name = "GotThatData/sports-cards"
         
         # Create local directory if it doesn't exist
         os.makedirs(local_images_dir, exist_ok=True)
@@ -43,12 +43,27 @@ class DatasetManager:
             file_list = self.drive.ListFile({'q': query}).GetList()
             
             if not file_list:
-                return False, "No files found in the specified folder", []
+                # Try to get single file if folder is empty
+                file = self.drive.CreateFile({'id': drive_folder_id})
+                if file:
+                    file_list = [file]
+                else:
+                    return False, "No files found with the specified ID", []
             
             renamed_files = []
+            existing_dataset = None
+            try:
+                existing_dataset = load_dataset(self.dataset_name)
+                logger.info(f"Loaded existing dataset: {self.dataset_name}")
+                # Get the current count of images to continue numbering
+                start_index = len(existing_dataset['train']) if 'train' in existing_dataset else 0
+            except Exception as e:
+                logger.info(f"No existing dataset found, starting fresh: {str(e)}")
+                start_index = 0
+            
             for i, file in enumerate(tqdm(file_list, desc="Downloading files")):
                 if file['mimeType'].startswith('image/'):
-                    new_filename = f"{naming_convention}_{i+1}.jpg"
+                    new_filename = f"{naming_convention}_{start_index + i + 1}.jpg"
                     file_path = os.path.join(self.local_images_dir, new_filename)
                     
                     # Download file
@@ -61,7 +76,8 @@ class DatasetManager:
                         renamed_files.append({
                             'file_path': file_path,
                             'original_name': file['title'],
-                            'new_name': new_filename
+                            'new_name': new_filename,
+                            'image': file_path  # Adding image column for dataset
                         })
                     except Exception as e:
                         logger.error(f"Error processing image {file['title']}: {str(e)}")
@@ -72,23 +88,32 @@ class DatasetManager:
         except Exception as e:
             return False, f"Error downloading files: {str(e)}", []
 
-    def update_huggingface_dataset(self, dataset_name, renamed_files):
-        """Update or create Hugging Face dataset with new images"""
+    def update_huggingface_dataset(self, renamed_files):
+        """Update the sports-cards dataset with new images"""
         try:
             # Create a DataFrame with the file information
             df = pd.DataFrame(renamed_files)
             
-            # Create a Hugging Face Dataset
-            dataset = Dataset.from_pandas(df)
+            # Create a Hugging Face Dataset from the new files
+            new_dataset = Dataset.from_pandas(df)
+            
+            try:
+                # Try to load existing dataset
+                existing_dataset = load_dataset(self.dataset_name)
+                # Concatenate with existing dataset if it exists
+                if 'train' in existing_dataset:
+                    new_dataset = concatenate_datasets([existing_dataset['train'], new_dataset])
+            except Exception:
+                logger.info("Creating new dataset")
             
             # Push to Hugging Face Hub
-            dataset.push_to_hub(dataset_name)
+            new_dataset.push_to_hub(self.dataset_name, split="train")
             
-            return True, f"Successfully updated dataset '{dataset_name}' with {len(renamed_files)} images"
+            return True, f"Successfully updated dataset '{self.dataset_name}' with {len(renamed_files)} new images"
         except Exception as e:
             return False, f"Error updating Hugging Face dataset: {str(e)}"
 
-def process_pipeline(folder_id, naming_convention, dataset_name):
+def process_pipeline(folder_id, naming_convention):
     """Main pipeline to process images and update dataset"""
     manager = DatasetManager()
     
@@ -103,33 +128,27 @@ def process_pipeline(folder_id, naming_convention, dataset_name):
         return message
     
     # Step 3: Update Hugging Face dataset
-    if dataset_name:
-        success, hf_message = manager.update_huggingface_dataset(dataset_name, renamed_files)
-        return f"{message}\n{hf_message}"
-    
-    return message
+    success, hf_message = manager.update_huggingface_dataset(renamed_files)
+    return f"{message}\n{hf_message}"
 
 # Gradio interface
 demo = gr.Interface(
     fn=process_pipeline,
     inputs=[
         gr.Textbox(
-            label="Google Drive Folder ID",
-            placeholder="Enter the folder ID from your Google Drive URL"
+            label="Google Drive File/Folder ID",
+            placeholder="Enter the ID from your Google Drive URL",
+            value="151VOxPO91mg0C3ORiioGUd4hogzP1ujm"  # Pre-filled with provided ID
         ),
         gr.Textbox(
             label="Naming Convention",
             placeholder="e.g., card",
-            value="card"
-        ),
-        gr.Textbox(
-            label="Hugging Face Dataset Name (Optional)",
-            placeholder="username/dataset-name"
+            value="sports_card"
         )
     ],
     outputs=gr.Textbox(label="Status"),
-    title="Card Image Processor",
-    description="Download card images from Google Drive and add them to your Hugging Face dataset"
+    title="Sports Cards Dataset Processor",
+    description="Download card images from Google Drive and add them to the sports-cards dataset"
 )
 
 if __name__ == "__main__":