Upload data/dowload_images_data.py with huggingface_hub

data/dowload_images_data.py

@@ -0,0 +1,217 @@
+#!/usr/bin/env python3
+"""
+Script to download all images from the dataset locally.
+This file downloads all images from URLs in the dataset CSV and saves them locally
+to speed up training by avoiding repeated downloads. It uses parallel processing
+to download multiple images simultaneously and updates the CSV with local paths
+of downloaded images.
+"""
+
+import pandas as pd
+import requests
+from PIL import Image
+from io import BytesIO
+from tqdm import tqdm
+import hashlib
+from pathlib import Path
+import time
+import concurrent.futures
+from threading import Lock
+import config
+
+class ImageDownloader:
+    def __init__(self, df, images_dir=config.images_dir, max_workers=8, timeout=10):
+        """
+        Initialize the image downloader.
+        
+        Args:
+            csv_path: Path to the CSV file containing the URLs
+            images_dir: Directory to save the images
+            max_workers: Number of threads for parallel download
+            timeout: Timeout for HTTP requests (seconds)
+        """
+        self.df = df
+        self.images_dir = Path(images_dir)
+        self.max_workers = max_workers
+        self.timeout = timeout
+        
+        # Create the images directory if it doesn't exist
+        self.images_dir.mkdir(parents=True, exist_ok=True)
+        
+        # Statistics
+        self.stats = {
+            'downloaded': 0,
+            'skipped': 0,
+            'failed': 0,
+            'total': 0
+        }
+        self.stats_lock = Lock()
+        
+    def url_to_filename(self, url):
+        """Convert a URL to a secure filename."""
+        # Use MD5 hash of the URL to avoid character issues
+        url_hash = hashlib.md5(url.encode()).hexdigest()
+        return f"{url_hash}.jpg"
+    
+    def download_single_image(self, row):
+        """
+        Download a single image.
+        
+        Args:
+            row: Tuple (index, pandas.Series) containing the row data
+            
+        Returns:
+            tuple: (success, index, message)
+        """
+        idx, data = row
+        url = data[config.column_url_image]
+        
+        # Filename based on the URL
+        filename = self.url_to_filename(url)
+        filepath = self.images_dir / filename
+        
+        # Check if the image already exists
+        if filepath.exists():
+            with self.stats_lock:
+                self.stats['skipped'] += 1
+            return True, idx, f"Skipped (already exists): {filename}"
+        
+        try:
+            # Download the image
+            response = requests.get(url, timeout=self.timeout, stream=True)
+            response.raise_for_status()
+            
+            # Check the content type
+            content_type = response.headers.get('content-type', '')
+            if not content_type.startswith('image/'):
+                with self.stats_lock:
+                    self.stats['failed'] += 1
+                return False, idx, f"Not an image: {content_type}"
+            
+                # Save the image
+            try:
+                image = Image.open(BytesIO(response.content)).convert("RGB")
+                image.save(filepath, "JPEG", quality=85, optimize=True)
+                
+                with self.stats_lock:
+                    self.stats['downloaded'] += 1
+                return True, idx, f"Downloaded: {filename}"
+                
+            except Exception as img_error:
+                with self.stats_lock:
+                    self.stats['failed'] += 1
+                return False, idx, f"Image processing error: {str(img_error)}"
+                
+        except requests.exceptions.RequestException as e:
+            with self.stats_lock:
+                self.stats['failed'] += 1
+            return False, idx, f"Download error: {str(e)}"
+        except Exception as e:
+            with self.stats_lock:
+                self.stats['failed'] += 1
+            return False, idx, f"Unexpected error: {str(e)}"
+    
+    def download_all_images(self):
+        """Download all images from the dataset."""
+        print(f"📊 Loading dataset from {self.df}")
+        self.stats['total'] = len(self.df)
+        
+        print(f"🔍 Found {len(self.df)} images to download")
+        print(f"📁 Saving in: {self.images_dir}")
+        print(f"🔧 Using {self.max_workers} threads")
+        
+        # Create a new DataFrame with local paths
+        df_local = self.df.copy()
+        df_local[config.column_local_image_path] = ""
+        df_local['download_success'] = False
+        
+        start_time = time.time()
+        
+        # Parallel download
+        with concurrent.futures.ThreadPoolExecutor(max_workers=self.max_workers) as executor:
+            # Submit all tasks
+            future_to_row = {
+                executor.submit(self.download_single_image, row): row 
+                for row in self.df.iterrows()
+            }
+            
+            # Process the results with a progress bar
+            with tqdm(total=len(self.df), desc="📥 Downloading", unit="img") as pbar:
+                for future in concurrent.futures.as_completed(future_to_row):
+                    row = future_to_row[future]
+                    idx = row[0]
+                    
+                    try:
+                        success, _, message = future.result()
+                        
+                        if success:
+                            # Add the local path to the DataFrame
+                            filename = self.url_to_filename(row[1][config.column_url_image])
+                            df_local.loc[idx, config.column_local_image_path] = str(self.images_dir / filename)
+                            df_local.loc[idx, 'download_success'] = True
+                        
+                        # Update the progress bar
+                        pbar.set_postfix({
+                            'OK': self.stats['downloaded'],
+                            'Skip': self.stats['skipped'], 
+                            'Fail': self.stats['failed']
+                        })
+                        pbar.update(1)
+                        
+                    except Exception as e:
+                        print(f"❌ Unexpected error for index {idx}: {e}")
+                        with self.stats_lock:
+                            self.stats['failed'] += 1
+                        pbar.update(1)
+        
+        elapsed_time = time.time() - start_time
+        
+        # Final statistics
+        print("\n" + "="*60)
+        print("📊 DOWNLOAD STATISTICS")
+        print("="*60)
+        print(f"✅ Downloaded: {self.stats['downloaded']}")
+        print(f"⏭️  Skipped (already present): {self.stats['skipped']}")
+        print(f"❌ Failed: {self.stats['failed']}")
+        print(f"📊 Total: {self.stats['total']}")
+        print(f"⏱️  Time elapsed: {elapsed_time:.1f}s")
+        
+        success_rate = (self.stats['downloaded'] + self.stats['skipped']) / self.stats['total'] * 100
+        print(f"🎯 Success rate: {success_rate:.1f}%")
+        
+        if self.stats['downloaded'] > 0:
+            avg_time = elapsed_time / self.stats['downloaded']
+            print(f"⚡ Average time per image: {avg_time:.2f}s")
+        
+        # Save the updated DataFrame
+        output_path = config.local_dataset_path
+        df_local.to_csv(output_path, index=False)
+        print(f"💾 Updated dataset saved: {output_path}")
+        
+        return df_local
+
+def main():
+    """Main function."""
+    print("🚀 STARTING IMAGE DOWNLOADER")
+    print("="*60)
+    
+    # Configuration
+    df = pd.read_csv(config.local_dataset_path)
+    df = df[df['color'] != 'unknown']
+    
+    # Create the downloader
+    downloader = ImageDownloader(
+        df=df,
+        images_dir=config.images_dir,
+        max_workers=8,
+        timeout=10
+    )
+    
+    # Download all images
+    df_with_paths = downloader.download_all_images()
+    
+    print("\n🎉 DOWNLOAD COMPLETED!")
+    print("💡 You can now use the local images for training.")
+
+if __name__ == "__main__":
+    main()