Update

app.py

@@ -2,7 +2,7 @@ from pydrive2.auth import GoogleAuth
 from pydrive2.drive import GoogleDrive
 import os
 import gradio as gr
-from datasets import load_dataset, Dataset
+from datasets import load_dataset, Dataset, concatenate_datasets
 import pandas as pd
 from PIL import Image
 from tqdm import tqdm
@@ -14,34 +14,53 @@ logging.basicConfig(level=logging.INFO)
 logger = logging.getLogger(__name__)
 
 # Load settings
+if not os.path.exists("settings.yaml"):
+    raise FileNotFoundError("settings.yaml file is missing. Please add it with 'client_secrets_file'.")
+
 with open('settings.yaml', 'r') as file:
     settings = yaml.safe_load(file)
 
+# Utility Functions
+def safe_load_dataset(dataset_name):
+    """Load Hugging Face dataset safely."""
+    try:
+        dataset = load_dataset(dataset_name)
+        return dataset, len(dataset['train']) if 'train' in dataset else 0
+    except Exception:
+        logger.info("No existing dataset found. Starting fresh.")
+        return None, 0
+
+def is_valid_image(file_path):
+    """Check if a file is a valid image."""
+    try:
+        with Image.open(file_path) as img:
+            img.verify()
+        return True
+    except:
+        return False
+
+# DatasetManager Class
 class DatasetManager:
     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)
-    
+
     def authenticate_drive(self):
-        """Authenticate with Google Drive"""
+        """Authenticate with Google Drive."""
         try:
             gauth = GoogleAuth()
             gauth.settings['client_config_file'] = settings['client_secrets_file']
             
             # Try to load saved credentials
             gauth.LoadCredentialsFile("credentials.txt")
-            
             if gauth.credentials is None:
                 gauth.LocalWebserverAuth()
             elif gauth.access_token_expired:
                 gauth.Refresh()
             else:
                 gauth.Authorize()
-                
             gauth.SaveCredentialsFile("credentials.txt")
             
             self.drive = GoogleDrive(gauth)
@@ -50,156 +69,97 @@ class DatasetManager:
             return False, f"Authentication failed: {str(e)}"
 
     def download_and_rename_files(self, drive_folder_id, naming_convention):
-        """Download files from Google Drive and rename them"""
+        """Download files from Google Drive and rename them."""
         if not self.drive:
             return False, "Google Drive not authenticated", []
         
         try:
             query = f"'{drive_folder_id}' in parents and trashed=false"
             file_list = self.drive.ListFile({'q': query}).GetList()
-            
             if not file_list:
-                file = self.drive.CreateFile({'id': drive_folder_id})
-                if file:
-                    file_list = [file]
-                else:
-                    return False, "No files found with the specified ID", []
-            
+                return False, "No files found in the specified folder.", []
+
+            existing_dataset, start_index = safe_load_dataset(self.dataset_name)
             renamed_files = []
-            try:
-                existing_dataset = load_dataset(self.dataset_name)
-                logger.info(f"Loaded existing dataset: {self.dataset_name}")
-                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/'):
+
+            for i, file in enumerate(tqdm(file_list, desc="Downloading files", unit="file")):
+                if 'mimeType' in file and 'image' in file['mimeType'].lower():
                     new_filename = f"{naming_convention}_{start_index + i + 1}.jpg"
                     file_path = os.path.join(self.local_images_dir, new_filename)
-                    
                     file.GetContentFile(file_path)
-                    
-                    try:
-                        with Image.open(file_path) as img:
-                            img.verify()
+
+                    if is_valid_image(file_path):
                         renamed_files.append({
                             'file_path': file_path,
                             'original_name': file['title'],
-                            'new_name': new_filename,
-                            'image': file_path
+                            'new_name': new_filename
                         })
-                    except Exception as e:
-                        logger.error(f"Error processing image {file['title']}: {str(e)}")
-                        if os.path.exists(file_path):
-                            os.remove(file_path)
+                        logger.info(f"Downloaded and renamed: {file['title']} -> {new_filename}")
+                    else:
+                        logger.error(f"Invalid image detected, removing {file_path}")
+                        os.remove(file_path)
             
-            return True, f"Successfully processed {len(renamed_files)} images", renamed_files
+            return True, f"Processed {len(renamed_files)} images", renamed_files
         except Exception as e:
-            return False, f"Error downloading files: {str(e)}", []
+            return False, f"Error during download: {str(e)}", []
 
     def update_huggingface_dataset(self, renamed_files):
-        """Update the sports-cards dataset with new images"""
+        """Update Hugging Face dataset with new images."""
         try:
             df = pd.DataFrame(renamed_files)
             new_dataset = Dataset.from_pandas(df)
+
+            existing_dataset, _ = safe_load_dataset(self.dataset_name)
+            if existing_dataset and 'train' in existing_dataset:
+                combined_dataset = concatenate_datasets([existing_dataset['train'], new_dataset])
+            else:
+                combined_dataset = new_dataset
             
-            try:
-                existing_dataset = load_dataset(self.dataset_name)
-                if 'train' in existing_dataset:
-                    new_dataset = concatenate_datasets([existing_dataset['train'], new_dataset])
-            except Exception:
-                logger.info("Creating new dataset")
-            
-            new_dataset.push_to_hub(self.dataset_name, split="train")
-            
-            return True, f"Successfully updated dataset '{self.dataset_name}' with {len(renamed_files)} new images"
+            combined_dataset.push_to_hub(self.dataset_name, split="train")
+            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)}"
 
+# Process Pipeline
 def process_pipeline(folder_id, naming_convention):
-    """Main pipeline to process images and update dataset"""
+    """Main pipeline for processing images and updating dataset."""
     manager = DatasetManager()
-    
+
+    # Step 1: Authenticate Google Drive
     auth_success, auth_message = manager.authenticate_drive()
     if not auth_success:
-        return auth_message
-    
+        return auth_message, []
+
+    # Step 2: Download and rename files
     success, message, renamed_files = manager.download_and_rename_files(folder_id, naming_convention)
     if not success:
-        return message
-    
+        return message, []
+
+    # Step 3: Update Hugging Face dataset
     success, hf_message = manager.update_huggingface_dataset(renamed_files)
-    return f"{message}\n{hf_message}"
-
-# Custom CSS for web-safe fonts and clean styling
-custom_css = """
-.gradio-container {
-    font-family: Arial, sans-serif !important;
-}
-
-h1, h2, h3 {
-    font-family: 'Helvetica Neue', Helvetica, Arial, sans-serif !important;
-    font-weight: 600 !important;
-}
-
-.gr-button {
-    font-family: Arial, sans-serif !important;
-}
-
-.gr-input {
-    font-family: 'Courier New', Courier, monospace !important;
-}
-
-.gr-box {
-    border-radius: 8px !important;
-    border: 1px solid #e5e5e5 !important;
-}
-
-.gr-padded {
-    padding: 16px !important;
-}
-"""
-
-# Gradio interface with custom theme
-with gr.Blocks(css=custom_css) as demo:
+    return f"{message}\n{hf_message}", renamed_files
+
+# Gradio Interface
+with gr.Blocks() as demo:
     gr.Markdown("# Sports Cards Dataset Processor")
-    
+
     with gr.Box():
-        gr.Markdown("""
-        ### Instructions
-        1. Enter the Google Drive folder/file ID
-        2. Choose a naming convention for your cards
-        3. Click Process to start
-        """)
+        gr.Markdown("### Instructions: Upload from Google Drive and Update Hugging Face Dataset")
     
     with gr.Row():
-        with gr.Column():
-            folder_id = gr.Textbox(
-                label="Google Drive File/Folder ID",
-                placeholder="Enter the ID from your Google Drive URL",
-                value="151VOxPO91mg0C3ORiioGUd4hogzP1ujm"
-            )
-            naming = gr.Textbox(
-                label="Naming Convention",
-                placeholder="e.g., sports_card",
-                value="sports_card"
-            )
-            process_btn = gr.Button("Process Images", variant="primary")
-    
-    with gr.Box():
-        output = gr.Textbox(
-            label="Processing Status",
-            show_label=True,
-            lines=5
-        )
-    
-    process_btn.click(
-        fn=process_pipeline,
-        inputs=[folder_id, naming],
-        outputs=output
-    )
+        folder_id = gr.Textbox(label="Google Drive Folder ID", placeholder="Enter the folder ID")
+        naming_convention = gr.Textbox(label="Naming Convention", placeholder="e.g., sports_card")
+        process_btn = gr.Button("Process Images")
+
+    output = gr.Textbox(label="Status")
+    output_table = gr.Dataframe(label="Processed Files", headers=["Original Name", "New Name", "File Path"])
+
+    def process_ui(folder_id, naming_convention):
+        status, renamed_files = process_pipeline(folder_id, naming_convention)
+        table_data = [[file['original_name'], file['new_name'], file['file_path']] for file in renamed_files]
+        return status, table_data
+
+    process_btn.click(process_ui, inputs=[folder_id, naming_convention], outputs=[output, output_table])
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()