preprocess_datasets.py

#!/usr/bin/env python3
"""
Dataset Preprocessor for SkinMate
Prepares and organizes datasets for model training
"""

import os
import sys
import shutil
import json
from pathlib import Path
from PIL import Image
import numpy as np
from sklearn.model_selection import train_test_split
import matplotlib.pyplot as plt

# Add parent directory to path to import app modules
sys.path.append(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

class DatasetPreprocessor:
    def __init__(self):
        self.base_dir = Path(__file__).parent
        self.datasets_dir = self.base_dir / "datasets"
        self.kaggle_dir = self.datasets_dir / "kaggle-skin-types"
        self.roboflow_dir = self.datasets_dir / "roboflow-skin-types"
        self.combined_dir = self.datasets_dir / "combined-skin-types"
        self.processed_dir = self.datasets_dir / "processed"
        
        # Create processed directory
        self.processed_dir.mkdir(exist_ok=True)
        
        # Standard image size for training
        self.target_size = (224, 224)
        
        # Skin type mapping
        self.skin_types = {
            'normal': 0,
            'oily': 1, 
            'dry': 2
        }
    
    def analyze_image_quality(self, image_path):
        """Analyze image quality and properties"""
        try:
            with Image.open(image_path) as img:
                # Basic image properties
                width, height = img.size
                mode = img.mode
                format_type = img.format
                
                # Calculate aspect ratio
                aspect_ratio = width / height
                
                # Check if image is too small
                min_size = min(width, height)
                
                return {
                    'width': width,
                    'height': height,
                    'mode': mode,
                    'format': format_type,
                    'aspect_ratio': aspect_ratio,
                    'min_size': min_size,
                    'is_valid': min_size >= 64 and mode in ['RGB', 'L']  # Minimum requirements
                }
        except Exception as e:
            return {
                'error': str(e),
                'is_valid': False
            }
    
    def preprocess_image(self, image_path, output_path):
        """Preprocess single image for training"""
        try:
            with Image.open(image_path) as img:
                # Convert to RGB if necessary
                if img.mode != 'RGB':
                    img = img.convert('RGB')
                
                # Resize to target size
                img_resized = img.resize(self.target_size, Image.Resampling.LANCZOS)
                
                # Save preprocessed image
                img_resized.save(output_path, 'JPEG', quality=95)
                
                return True
        except Exception as e:
            print(f"❌ Error processing {image_path}: {e}")
            return False
    
    def process_kaggle_dataset(self):
        """Process Kaggle dataset into standard format"""
        print("🔄 Processing Kaggle dataset...")
        
        kaggle_source = self.kaggle_dir / "Oily-Dry-Skin-Types"
        processed_kaggle = self.processed_dir / "kaggle"
        processed_kaggle.mkdir(exist_ok=True)
        
        dataset_info = {
            'source': 'kaggle',
            'total_processed': 0,
            'splits': {},
            'quality_stats': {
                'valid_images': 0,
                'invalid_images': 0,
                'errors': []
            }
        }
        
        # Process each split (train, valid, test)
        for split_name in ['train', 'valid', 'test']:
            split_dir = kaggle_source / split_name
            if not split_dir.exists():
                continue
            
            output_split_dir = processed_kaggle / split_name
            output_split_dir.mkdir(exist_ok=True)
            
            split_stats = {'total': 0, 'processed': 0, 'skipped': 0}
            
            # Process each skin type
            for skin_type in ['normal', 'oily', 'dry']:
                skin_dir = split_dir / skin_type
                if not skin_dir.exists():
                    continue
                
                output_skin_dir = output_split_dir / skin_type
                output_skin_dir.mkdir(exist_ok=True)
                
                print(f"  📂 Processing {split_name}/{skin_type}...")
                
                # Process each image
                image_files = list(skin_dir.glob("*.jpg")) + list(skin_dir.glob("*.jpeg")) + list(skin_dir.glob("*.png"))
                
                for i, image_path in enumerate(image_files):
                    split_stats['total'] += 1
                    
                    # Analyze image quality
                    quality_info = self.analyze_image_quality(image_path)
                    
                    if not quality_info['is_valid']:
                        split_stats['skipped'] += 1
                        dataset_info['quality_stats']['invalid_images'] += 1
                        if 'error' in quality_info:
                            dataset_info['quality_stats']['errors'].append(quality_info['error'])
                        continue
                    
                    # Generate output filename
                    output_filename = f"{skin_type}_{split_name}_{i:04d}.jpg"
                    output_path = output_skin_dir / output_filename
                    
                    # Preprocess and save
                    if self.preprocess_image(image_path, output_path):
                        split_stats['processed'] += 1
                        dataset_info['quality_stats']['valid_images'] += 1
                    else:
                        split_stats['skipped'] += 1
                        dataset_info['quality_stats']['invalid_images'] += 1
            
            dataset_info['splits'][split_name] = split_stats
            dataset_info['total_processed'] += split_stats['processed']
            
            print(f"  ✅ {split_name}: {split_stats['processed']}/{split_stats['total']} images processed")
        
        # Save dataset info
        info_file = processed_kaggle / "dataset_info.json"
        with open(info_file, 'w') as f:
            json.dump(dataset_info, f, indent=2)
        
        print(f"✅ Kaggle dataset processed: {dataset_info['total_processed']} images")
        return dataset_info
    
    def create_training_dataset(self):
        """Create final training dataset with proper splits"""
        print("🔄 Creating training dataset...")
        
        training_dir = self.processed_dir / "training"
        training_dir.mkdir(exist_ok=True)
        
        # Check if Kaggle dataset is processed
        kaggle_processed = self.processed_dir / "kaggle"
        if not kaggle_processed.exists():
            print("❌ Kaggle dataset not processed yet. Run process_kaggle_dataset() first.")
            return None
        
        # Create training structure
        for split in ['train', 'val', 'test']:
            for skin_type in ['normal', 'oily', 'dry']:
                (training_dir / split / skin_type).mkdir(parents=True, exist_ok=True)
        
        # Copy from Kaggle processed dataset
        kaggle_splits = {
            'train': 'train',
            'val': 'valid', 
            'test': 'test'
        }
        
        training_stats = {
            'train': {'normal': 0, 'oily': 0, 'dry': 0},
            'val': {'normal': 0, 'oily': 0, 'dry': 0},
            'test': {'normal': 0, 'oily': 0, 'dry': 0}
        }
        
        for train_split, kaggle_split in kaggle_splits.items():
            kaggle_split_dir = kaggle_processed / kaggle_split
            if not kaggle_split_dir.exists():
                continue
            
            for skin_type in ['normal', 'oily', 'dry']:
                src_dir = kaggle_split_dir / skin_type
                dst_dir = training_dir / train_split / skin_type
                
                if src_dir.exists():
                    # Copy all images
                    image_files = list(src_dir.glob("*.jpg"))
                    for img_file in image_files:
                        shutil.copy2(img_file, dst_dir)
                        training_stats[train_split][skin_type] += 1
        
        # Save training dataset info
        training_info = {
            'total_images': sum(sum(split.values()) for split in training_stats.values()),
            'splits': training_stats,
            'skin_types': list(self.skin_types.keys()),
            'image_size': self.target_size,
            'created_from': ['kaggle']
        }
        
        info_file = training_dir / "dataset_info.json"
        with open(info_file, 'w') as f:
            json.dump(training_info, f, indent=2)
        
        print("📊 Training Dataset Summary:")
        for split, counts in training_stats.items():
            total = sum(counts.values())
            print(f"  {split}: {total} images")
            for skin_type, count in counts.items():
                print(f"    - {skin_type}: {count}")
        
        print(f"✅ Training dataset created: {training_info['total_images']} total images")
        return training_info
    
    def visualize_dataset(self):
        """Create visualization of dataset distribution"""
        print("📊 Creating dataset visualization...")
        
        training_dir = self.processed_dir / "training"
        info_file = training_dir / "dataset_info.json"
        
        if not info_file.exists():
            print("❌ Training dataset not found. Create it first.")
            return
        
        with open(info_file, 'r') as f:
            dataset_info = json.load(f)
        
        # Create visualization
        fig, axes = plt.subplots(1, 2, figsize=(15, 6))
        
        # Plot 1: Distribution by split
        splits = list(dataset_info['splits'].keys())
        split_totals = [sum(dataset_info['splits'][split].values()) for split in splits]
        
        axes[0].bar(splits, split_totals, color=['#3498db', '#e74c3c', '#f39c12'])
        axes[0].set_title('Dataset Distribution by Split')
        axes[0].set_ylabel('Number of Images')
        for i, total in enumerate(split_totals):
            axes[0].text(i, total + 10, str(total), ha='center', va='bottom')
        
        # Plot 2: Distribution by skin type (stacked)
        skin_types = ['normal', 'oily', 'dry']
        train_counts = [dataset_info['splits']['train'][st] for st in skin_types]
        val_counts = [dataset_info['splits']['val'][st] for st in skin_types]
        test_counts = [dataset_info['splits']['test'][st] for st in skin_types]
        
        x = np.arange(len(skin_types))
        width = 0.25
        
        axes[1].bar(x - width, train_counts, width, label='Train', color='#3498db')
        axes[1].bar(x, val_counts, width, label='Validation', color='#e74c3c')
        axes[1].bar(x + width, test_counts, width, label='Test', color='#f39c12')
        
        axes[1].set_title('Dataset Distribution by Skin Type')
        axes[1].set_ylabel('Number of Images')
        axes[1].set_xlabel('Skin Type')
        axes[1].set_xticks(x)
        axes[1].set_xticklabels(skin_types)
        axes[1].legend()
        
        plt.tight_layout()
        
        # Save visualization
        viz_file = training_dir / "dataset_visualization.png"
        plt.savefig(viz_file, dpi=300, bbox_inches='tight')
        plt.close()
        
        print(f"✅ Dataset visualization saved: {viz_file}")
    
    def run_full_preprocessing(self):
        """Run complete preprocessing pipeline"""
        print("🚀 Starting full dataset preprocessing...")
        print("=" * 50)
        
        # Step 1: Process Kaggle dataset
        kaggle_info = self.process_kaggle_dataset()
        
        # Step 2: Create training dataset
        training_info = self.create_training_dataset()
        
        # Step 3: Create visualization
        self.visualize_dataset()
        
        print("=" * 50)
        print("✅ Full preprocessing completed!")
        
        if training_info:
            print(f"📊 Final Training Dataset: {training_info['total_images']} images")
            print("📁 Location: datasets/processed/training/")
            print("\n📝 Next steps:")
            print("1. Review dataset visualization")
            print("2. Train initial skin classification model")
            print("3. Evaluate model performance")
        
        return training_info

if __name__ == "__main__":
    preprocessor = DatasetPreprocessor()
    preprocessor.run_full_preprocessing()