Update

app.py

@@ -10,7 +10,10 @@ import logging
 import yaml
 
 # Set up logging
-logging.basicConfig(level=logging.INFO)
+logging.basicConfig(
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s'
+)
 logger = logging.getLogger(__name__)
 
 # Load settings
@@ -26,8 +29,8 @@ def safe_load_dataset(dataset_name):
     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.")
+    except Exception as e:
+        logger.info(f"No existing dataset found. Starting fresh. Error: {str(e)}")
         return None, 0
 
 def is_valid_image(file_path):
@@ -35,10 +38,21 @@ def is_valid_image(file_path):
     try:
         with Image.open(file_path) as img:
             img.verify()
-        return True
-    except:
+            return True
+    except Exception as e:
+        logger.error(f"Invalid image: {file_path}. Error: {str(e)}")
         return False
 
+def validate_input(folder_id, naming_convention):
+    """Validate user input."""
+    if not folder_id or not folder_id.strip():
+        return False, "Folder ID cannot be empty"
+    if not naming_convention or not naming_convention.strip():
+        return False, "Naming convention cannot be empty"
+    if not naming_convention.replace('_', '').isalnum():
+        return False, "Naming convention should only contain letters, numbers, and underscores"
+    return True, ""
+
 # DatasetManager Class
 class DatasetManager:
     def __init__(self, local_images_dir="downloaded_cards"):
@@ -66,6 +80,7 @@ class DatasetManager:
             self.drive = GoogleDrive(gauth)
             return True, "Successfully authenticated with Google Drive"
         except Exception as e:
+            logger.error(f"Authentication failed: {str(e)}")
             return False, f"Authentication failed: {str(e)}"
 
     def download_and_rename_files(self, drive_folder_id, naming_convention):
@@ -76,35 +91,55 @@ class DatasetManager:
         try:
             query = f"'{drive_folder_id}' in parents and trashed=false"
             file_list = self.drive.ListFile({'q': query}).GetList()
+            
             if not file_list:
+                logger.warning(f"No files found in folder: {drive_folder_id}")
                 return False, "No files found in the specified folder.", []
 
             existing_dataset, start_index = safe_load_dataset(self.dataset_name)
             renamed_files = []
+            processed_count = 0
+            error_count = 0
 
             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"
+                    new_filename = f"{naming_convention}_{start_index + processed_count + 1}.jpg"
                     file_path = os.path.join(self.local_images_dir, new_filename)
-                    file.GetContentFile(file_path)
-
-                    if is_valid_image(file_path):
-                        renamed_files.append({
-                            'file_path': file_path,
-                            'original_name': file['title'],
-                            'new_name': new_filename
-                        })
-                        logger.info(f"Downloaded and renamed: {file['title']} -> {new_filename}")
-                    else:
-                        logger.error(f"Invalid image detected, removing {file_path}")
-                        os.remove(file_path)
+                    
+                    try:
+                        file.GetContentFile(file_path)
+                        if is_valid_image(file_path):
+                            renamed_files.append({
+                                'file_path': file_path,
+                                'original_name': file['title'],
+                                'new_name': new_filename
+                            })
+                            processed_count += 1
+                            logger.info(f"Successfully processed: {file['title']} -> {new_filename}")
+                        else:
+                            error_count += 1
+                            if os.path.exists(file_path):
+                                os.remove(file_path)
+                    except Exception as e:
+                        error_count += 1
+                        logger.error(f"Error processing file {file['title']}: {str(e)}")
+                        if os.path.exists(file_path):
+                            os.remove(file_path)
+
+            status_message = f"Processed {processed_count} images successfully"
+            if error_count > 0:
+                status_message += f" ({error_count} files failed)"
             
-            return True, f"Processed {len(renamed_files)} images", renamed_files
+            return True, status_message, renamed_files
         except Exception as e:
+            logger.error(f"Download error: {str(e)}")
             return False, f"Error during download: {str(e)}", []
 
     def update_huggingface_dataset(self, renamed_files):
         """Update Hugging Face dataset with new images."""
+        if not renamed_files:
+            return False, "No files to update"
+            
         try:
             df = pd.DataFrame(renamed_files)
             new_dataset = Dataset.from_pandas(df)
@@ -118,11 +153,17 @@ class DatasetManager:
             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:
+            logger.error(f"Dataset update error: {str(e)}")
             return False, f"Error updating Hugging Face dataset: {str(e)}"
 
 # Process Pipeline
 def process_pipeline(folder_id, naming_convention):
     """Main pipeline for processing images and updating dataset."""
+    # Validate input
+    is_valid, error_message = validate_input(folder_id, naming_convention)
+    if not is_valid:
+        return error_message, []
+
     manager = DatasetManager()
 
     # Step 1: Authenticate Google Drive
@@ -140,26 +181,60 @@ def process_pipeline(folder_id, naming_convention):
     return f"{message}\n{hf_message}", renamed_files
 
 # Gradio Interface
-with gr.Blocks() as demo:
+with gr.Blocks(title="Sports Cards Dataset Processor") as demo:
     gr.Markdown("# Sports Cards Dataset Processor")
 
     with gr.Box():
-        gr.Markdown("### Instructions: Upload from Google Drive and Update Hugging Face Dataset")
+        gr.Markdown("""
+        ### Instructions
+        1. Enter the Google Drive folder ID (found in the folder's URL)
+        2. Specify a naming convention for the files (e.g., 'sports_card')
+        3. Click 'Process Images' to start
+        
+        Note: Only image files will be processed. Invalid images will be skipped.
+        """)
     
     with gr.Row():
-        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"])
+        folder_id = gr.Textbox(
+            label="Google Drive Folder ID",
+            placeholder="Enter the folder ID from the URL",
+            info="Found in your Google Drive folder's URL"
+        )
+        naming_convention = gr.Textbox(
+            label="Naming Convention",
+            placeholder="e.g., sports_card",
+            value="sports_card",
+            info="Use only letters, numbers, and underscores"
+        )
+    
+    process_btn = gr.Button("Process Images", variant="primary")
+    
+    with gr.Row():
+        with gr.Column():
+            output = gr.Textbox(
+                label="Processing Status",
+                show_label=True,
+                lines=3
+            )
+        
+        with gr.Column():
+            output_table = gr.Dataframe(
+                label="Processed Files",
+                headers=["Original Name", "New Name", "File Path"],
+                wrap=True
+            )
 
     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]
+        table_data = [[file['original_name'], file['new_name'], file['file_path']] 
+                     for file in renamed_files] if renamed_files else []
         return status, table_data
 
-    process_btn.click(process_ui, inputs=[folder_id, naming_convention], outputs=[output, output_table])
+    process_btn.click(
+        fn=process_ui,
+        inputs=[folder_id, naming_convention],
+        outputs=[output, output_table]
+    )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch()