Upload model.py

model.py

@@ -0,0 +1,373 @@
+import os
+import numpy as np
+import matplotlib.pyplot as plt
+from datasets import load_dataset
+from sklearn.model_selection import train_test_split
+from sklearn.metrics import classification_report, confusion_matrix
+from sklearn.utils.class_weight import compute_class_weight
+import seaborn as sns
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import Dataset, DataLoader
+from torchvision import transforms, models
+from PIL import Image
+from tqdm import tqdm
+import warnings
+warnings.filterwarnings('ignore')
+
+# Set random seeds for reproducibility
+torch.manual_seed(42)
+np.random.seed(42)
+
+# Configuration
+CONFIG = {
+    'img_size': 224,
+    'batch_size': 16,  # Reduced batch size
+    'num_epochs': 30,
+    'learning_rate': 0.0001,
+    'patience': 7,
+    'device': 'cuda' if torch.cuda.is_available() else 'cpu',
+    'num_workers': 0,  # Set to 0 to avoid multiprocessing issues
+    'model_save_path': 'best_trash_classifier.pth',
+}
+
+print(f"Using device: {CONFIG['device']}")
+
+# Memory-Efficient Dataset Class
+class TrashDatasetLazy(Dataset):
+    def __init__(self, dataset, indices, transform=None):
+        self.dataset = dataset
+        self.indices = indices
+        self.transform = transform
+    
+    def __len__(self):
+        return len(self.indices)
+    
+    def __getitem__(self, idx):
+        actual_idx = self.indices[idx]
+        item = self.dataset[actual_idx]
+        
+        image = item['image']
+        label = item['label']
+        
+        # Convert to PIL Image if needed
+        if not isinstance(image, Image.Image):
+            image = Image.fromarray(np.array(image))
+        
+        # Convert to RGB if grayscale
+        if image.mode != 'RGB':
+            image = image.convert('RGB')
+        
+        if self.transform:
+            image = self.transform(image)
+        
+        return image, label
+
+# Data Augmentation and Normalization
+train_transform = transforms.Compose([
+    transforms.Resize((CONFIG['img_size'], CONFIG['img_size'])),
+    transforms.RandomHorizontalFlip(p=0.5),
+    transforms.RandomRotation(15),
+    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2),
+    transforms.RandomAffine(degrees=0, translate=(0.1, 0.1)),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+])
+
+val_transform = transforms.Compose([
+    transforms.Resize((CONFIG['img_size'], CONFIG['img_size'])),
+    transforms.ToTensor(),
+    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225])
+])
+
+# Load Dataset (streaming mode)
+print("\n" + "="*60)
+print("LOADING DATASET")
+print("="*60)
+
+ds = load_dataset("rootstrap-org/waste-classifier", split="train")
+print(f"Dataset loaded successfully!")
+print(f"Total samples: {len(ds)}")
+
+# Get class names
+class_names = ds.features['label'].names
+num_classes = len(class_names)
+print(f"\nNumber of classes: {num_classes}")
+print(f"Classes: {class_names}")
+
+# Extract only labels for splitting (not images!)
+labels = [item['label'] for item in ds]
+
+# Check class distribution
+unique, counts = np.unique(labels, return_counts=True)
+print("\nClass Distribution:")
+for cls_idx, count in zip(unique, counts):
+    print(f"  {class_names[cls_idx]}: {count} samples ({count/len(labels)*100:.2f}%)")
+
+# Split dataset: 70% train, 15% val, 15% test
+print("\n" + "="*60)
+print("SPLITTING DATASET")
+print("="*60)
+
+indices = np.arange(len(ds))
+train_idx, temp_idx, y_train, y_temp = train_test_split(
+    indices, labels, test_size=0.3, random_state=42, stratify=labels
+)
+val_idx, test_idx, y_val, y_test = train_test_split(
+    temp_idx, y_temp, test_size=0.5, random_state=42, stratify=y_temp
+)
+
+print(f"Train set: {len(train_idx)} samples")
+print(f"Validation set: {len(val_idx)} samples")
+print(f"Test set: {len(test_idx)} samples")
+
+# Calculate class weights for handling imbalance
+class_weights = compute_class_weight(
+    class_weight='balanced',
+    classes=np.unique(y_train),
+    y=y_train
+)
+class_weights = torch.FloatTensor(class_weights).to(CONFIG['device'])
+print(f"\nClass weights (for imbalance): {class_weights.cpu().numpy()}")
+
+# Create datasets and dataloaders
+train_dataset = TrashDatasetLazy(ds, train_idx, transform=train_transform)
+val_dataset = TrashDatasetLazy(ds, val_idx, transform=val_transform)
+test_dataset = TrashDatasetLazy(ds, test_idx, transform=val_transform)
+
+train_loader = DataLoader(
+    train_dataset, batch_size=CONFIG['batch_size'], 
+    shuffle=True, num_workers=CONFIG['num_workers'], pin_memory=True
+)
+val_loader = DataLoader(
+    val_dataset, batch_size=CONFIG['batch_size'], 
+    shuffle=False, num_workers=CONFIG['num_workers'], pin_memory=True
+)
+test_loader = DataLoader(
+    test_dataset, batch_size=CONFIG['batch_size'], 
+    shuffle=False, num_workers=CONFIG['num_workers'], pin_memory=True
+)
+
+# Build Model using EfficientNetV2 (pretrained)
+print("\n" + "="*60)
+print("BUILDING MODEL")
+print("="*60)
+
+model = models.efficientnet_v2_s(weights='IMAGENET1K_V1')
+
+# Freeze early layers
+for param in list(model.parameters())[:-20]:
+    param.requires_grad = False
+
+# Modify classifier for our number of classes
+num_features = model.classifier[1].in_features
+model.classifier = nn.Sequential(
+    nn.Dropout(p=0.3, inplace=True),
+    nn.Linear(num_features, 512),
+    nn.ReLU(),
+    nn.Dropout(p=0.3),
+    nn.Linear(512, num_classes)
+)
+
+model = model.to(CONFIG['device'])
+print(f"Model: EfficientNetV2-S (pretrained)")
+print(f"Total parameters: {sum(p.numel() for p in model.parameters()):,}")
+print(f"Trainable parameters: {sum(p.numel() for p in model.parameters() if p.requires_grad):,}")
+
+# Loss function with class weights and optimizer
+criterion = nn.CrossEntropyLoss(weight=class_weights)
+optimizer = optim.AdamW(model.parameters(), lr=CONFIG['learning_rate'], weight_decay=0.01)
+scheduler = optim.lr_scheduler.ReduceLROnPlateau(
+    optimizer, mode='min', factor=0.5, patience=3, verbose=True
+)
+
+# Training and Validation Functions
+def train_epoch(model, loader, criterion, optimizer, device):
+    model.train()
+    running_loss = 0.0
+    correct = 0
+    total = 0
+    
+    pbar = tqdm(loader, desc='Training')
+    for images, labels in pbar:
+        images, labels = images.to(device), labels.to(device)
+        
+        optimizer.zero_grad()
+        outputs = model(images)
+        loss = criterion(outputs, labels)
+        loss.backward()
+        optimizer.step()
+        
+        running_loss += loss.item() * images.size(0)
+        _, predicted = outputs.max(1)
+        total += labels.size(0)
+        correct += predicted.eq(labels).sum().item()
+        
+        pbar.set_postfix({'loss': f'{loss.item():.4f}', 'acc': f'{100.*correct/total:.2f}%'})
+    
+    epoch_loss = running_loss / total
+    epoch_acc = 100. * correct / total
+    return epoch_loss, epoch_acc
+
+def validate_epoch(model, loader, criterion, device):
+    model.eval()
+    running_loss = 0.0
+    correct = 0
+    total = 0
+    
+    with torch.no_grad():
+        pbar = tqdm(loader, desc='Validation')
+        for images, labels in pbar:
+            images, labels = images.to(device), labels.to(device)
+            
+            outputs = model(images)
+            loss = criterion(outputs, labels)
+            
+            running_loss += loss.item() * images.size(0)
+            _, predicted = outputs.max(1)
+            total += labels.size(0)
+            correct += predicted.eq(labels).sum().item()
+            
+            pbar.set_postfix({'loss': f'{loss.item():.4f}', 'acc': f'{100.*correct/total:.2f}%'})
+    
+    epoch_loss = running_loss / total
+    epoch_acc = 100. * correct / total
+    return epoch_loss, epoch_acc
+
+# Training Loop
+print("\n" + "="*60)
+print("TRAINING MODEL")
+print("="*60)
+
+history = {
+    'train_loss': [], 'train_acc': [],
+    'val_loss': [], 'val_acc': []
+}
+
+best_val_acc = 0.0
+patience_counter = 0
+
+for epoch in range(CONFIG['num_epochs']):
+    print(f"\nEpoch {epoch+1}/{CONFIG['num_epochs']}")
+    print("-" * 60)
+    
+    train_loss, train_acc = train_epoch(model, train_loader, criterion, optimizer, CONFIG['device'])
+    val_loss, val_acc = validate_epoch(model, val_loader, criterion, CONFIG['device'])
+    
+    history['train_loss'].append(train_loss)
+    history['train_acc'].append(train_acc)
+    history['val_loss'].append(val_loss)
+    history['val_acc'].append(val_acc)
+    
+    print(f"\nTrain Loss: {train_loss:.4f} | Train Acc: {train_acc:.2f}%")
+    print(f"Val Loss: {val_loss:.4f} | Val Acc: {val_acc:.2f}%")
+    
+    scheduler.step(val_loss)
+    
+    # Save best model
+    if val_acc > best_val_acc:
+        best_val_acc = val_acc
+        torch.save({
+            'epoch': epoch,
+            'model_state_dict': model.state_dict(),
+            'optimizer_state_dict': optimizer.state_dict(),
+            'val_acc': val_acc,
+            'class_names': class_names
+        }, CONFIG['model_save_path'])
+        print(f"✓ Model saved! (Val Acc: {val_acc:.2f}%)")
+        patience_counter = 0
+    else:
+        patience_counter += 1
+        print(f"No improvement ({patience_counter}/{CONFIG['patience']})")
+    
+    if patience_counter >= CONFIG['patience']:
+        print("\nEarly stopping triggered!")
+        break
+
+# Plot Training History
+print("\n" + "="*60)
+print("SAVING TRAINING GRAPHS")
+print("="*60)
+
+fig, axes = plt.subplots(1, 2, figsize=(15, 5))
+
+# Loss plot
+axes[0].plot(history['train_loss'], label='Train Loss', marker='o')
+axes[0].plot(history['val_loss'], label='Val Loss', marker='s')
+axes[0].set_xlabel('Epoch')
+axes[0].set_ylabel('Loss')
+axes[0].set_title('Training and Validation Loss')
+axes[0].legend()
+axes[0].grid(True, alpha=0.3)
+
+# Accuracy plot
+axes[1].plot(history['train_acc'], label='Train Acc', marker='o')
+axes[1].plot(history['val_acc'], label='Val Acc', marker='s')
+axes[1].set_xlabel('Epoch')
+axes[1].set_ylabel('Accuracy (%)')
+axes[1].set_title('Training and Validation Accuracy')
+axes[1].legend()
+axes[1].grid(True, alpha=0.3)
+
+plt.tight_layout()
+plt.savefig('training_history.png', dpi=300, bbox_inches='tight')
+print("✓ Training graphs saved as 'training_history.png'")
+
+# Load Best Model and Test
+print("\n" + "="*60)
+print("LOADING BEST MODEL AND TESTING")
+print("="*60)
+
+checkpoint = torch.load(CONFIG['model_save_path'])
+model.load_state_dict(checkpoint['model_state_dict'])
+print(f"✓ Loaded best model from epoch {checkpoint['epoch']+1}")
+print(f"  Best validation accuracy: {checkpoint['val_acc']:.2f}%")
+
+# Test the model
+model.eval()
+all_preds = []
+all_labels = []
+
+with torch.no_grad():
+    for images, labels in tqdm(test_loader, desc='Testing'):
+        images = images.to(CONFIG['device'])
+        outputs = model(images)
+        _, predicted = outputs.max(1)
+        
+        all_preds.extend(predicted.cpu().numpy())
+        all_labels.extend(labels.numpy())
+
+# Calculate test accuracy
+test_acc = 100. * np.sum(np.array(all_preds) == np.array(all_labels)) / len(all_labels)
+print(f"\n{'='*60}")
+print(f"TEST SET ACCURACY: {test_acc:.2f}%")
+print(f"{'='*60}")
+
+# Classification Report
+print("\n" + "="*60)
+print("CLASSIFICATION REPORT")
+print("="*60)
+print(classification_report(all_labels, all_preds, target_names=class_names, digits=4))
+
+# Confusion Matrix
+cm = confusion_matrix(all_labels, all_preds)
+plt.figure(figsize=(10, 8))
+sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', 
+            xticklabels=class_names, yticklabels=class_names)
+plt.title('Confusion Matrix - Test Set')
+plt.ylabel('True Label')
+plt.xlabel('Predicted Label')
+plt.xticks(rotation=45, ha='right')
+plt.yticks(rotation=0)
+plt.tight_layout()
+plt.savefig('confusion_matrix.png', dpi=300, bbox_inches='tight')
+print("\n✓ Confusion matrix saved as 'confusion_matrix.png'")
+
+print("\n" + "="*60)
+print("TRAINING COMPLETE!")
+print("="*60)
+print(f"✓ Best model saved: {CONFIG['model_save_path']}")
+print(f"✓ Training history: training_history.png")
+print(f"✓ Confusion matrix: confusion_matrix.png")
+print(f"✓ Final test accuracy: {test_acc:.2f}%")