Upload 2 files

data_utils.py

@@ -0,0 +1,377 @@
+"""
+Data utilities for fire detection classification
+Handles data loading, transformations, and dataset management
+"""
+
+import os
+import torch
+from torch.utils.data import Dataset, DataLoader, WeightedRandomSampler
+from torchvision import transforms, datasets
+from PIL import Image
+import numpy as np
+from typing import Tuple, Dict, List, Optional
+from collections import Counter
+import random
+
+class FireDetectionDataset(Dataset):
+    """
+    Custom dataset for fire detection images
+    Supports both training and validation modes with appropriate transforms
+    """
+    
+    def __init__(self, data_dir: str, split: str = 'train', image_size: int = 224):
+        """
+        Initialize fire detection dataset
+        
+        Args:
+            data_dir: Root directory containing train/val folders
+            split: 'train' or 'val'
+            image_size: Size to resize images to
+        """
+        self.data_dir = data_dir
+        self.split = split
+        self.image_size = image_size
+        
+        # Define class mapping
+        self.classes = ['fire', 'no_fire']  # 0: fire, 1: no_fire
+        self.class_to_idx = {cls: idx for idx, cls in enumerate(self.classes)}
+        
+        # Load image paths and labels
+        self.samples = self._load_samples()
+        
+        # Define transforms
+        self.transform = self._get_transforms()
+        
+        print(f"🔥 {split.upper()} Dataset loaded:")
+        print(f"   Total samples: {len(self.samples)}")
+        print(f"   Classes: {self.classes}")
+        self._print_class_distribution()
+    
+    def _load_samples(self) -> List[Tuple[str, int]]:
+        """Load image paths and corresponding labels"""
+        samples = []
+        split_dir = os.path.join(self.data_dir, self.split)
+        
+        for class_name in self.classes:
+            class_dir = os.path.join(split_dir, class_name)
+            if not os.path.exists(class_dir):
+                print(f"⚠️ Warning: {class_dir} not found")
+                continue
+            
+            class_idx = self.class_to_idx[class_name]
+            
+            # Load all images from class directory and subdirectories
+            for root, dirs, files in os.walk(class_dir):
+                for img_name in files:
+                    if img_name.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tiff')):
+                        img_path = os.path.join(root, img_name)
+                        samples.append((img_path, class_idx))
+        
+        return samples
+    
+    def _print_class_distribution(self):
+        """Print class distribution for the dataset"""
+        class_counts = Counter([label for _, label in self.samples])
+        for class_name, class_idx in self.class_to_idx.items():
+            count = class_counts.get(class_idx, 0)
+            print(f"   {class_name}: {count} samples")
+    
+    def _get_transforms(self) -> transforms.Compose:
+        """Get appropriate transforms for the split"""
+        if self.split == 'train':
+            return transforms.Compose([
+                transforms.Resize((self.image_size + 32, self.image_size + 32)),
+                transforms.RandomResizedCrop(self.image_size, scale=(0.8, 1.0)),
+                transforms.RandomHorizontalFlip(p=0.5),
+                transforms.RandomRotation(degrees=10),
+                transforms.ColorJitter(
+                    brightness=0.2,
+                    contrast=0.2,
+                    saturation=0.2,
+                    hue=0.1
+                ),
+                transforms.ToTensor(),
+                transforms.Normalize(
+                    mean=[0.485, 0.456, 0.406],
+                    std=[0.229, 0.224, 0.225]
+                ),
+                transforms.RandomErasing(p=0.1, scale=(0.02, 0.08))
+            ])
+        else:
+            return transforms.Compose([
+                transforms.Resize((self.image_size, self.image_size)),
+                transforms.ToTensor(),
+                transforms.Normalize(
+                    mean=[0.485, 0.456, 0.406],
+                    std=[0.229, 0.224, 0.225]
+                )
+            ])
+    
+    def __len__(self) -> int:
+        return len(self.samples)
+    
+    def __getitem__(self, idx: int) -> Tuple[torch.Tensor, int]:
+        """Get a sample from the dataset"""
+        img_path, label = self.samples[idx]
+        
+        # Load image
+        try:
+            image = Image.open(img_path).convert('RGB')
+        except Exception as e:
+            print(f"⚠️ Error loading image {img_path}: {e}")
+            # Return a black image as fallback
+            image = Image.new('RGB', (self.image_size, self.image_size), color='black')
+        
+        # Apply transforms
+        if self.transform:
+            image = self.transform(image)
+        
+        return image, label
+
+def create_data_loaders(
+    data_dir: str,
+    batch_size: int = 16,
+    num_workers: int = 4,
+    image_size: int = 224,
+    use_weighted_sampling: bool = True
+) -> Tuple[DataLoader, DataLoader]:
+    """
+    Create train and validation data loaders
+    
+    Args:
+        data_dir: Root directory containing train/val folders
+        batch_size: Batch size for data loaders
+        num_workers: Number of worker processes
+        image_size: Size to resize images to
+        use_weighted_sampling: Whether to use weighted sampling for imbalanced data
+    
+    Returns:
+        Tuple of (train_loader, val_loader)
+    """
+    # Create datasets
+    train_dataset = FireDetectionDataset(data_dir, 'train', image_size)
+    val_dataset = FireDetectionDataset(data_dir, 'val', image_size)
+    
+    # Create samplers
+    train_sampler = None
+    if use_weighted_sampling and len(train_dataset) > 0:
+        train_sampler = create_weighted_sampler(train_dataset)
+    
+    # Create data loaders
+    train_loader = DataLoader(
+        train_dataset,
+        batch_size=batch_size,
+        sampler=train_sampler,
+        shuffle=(train_sampler is None),
+        num_workers=num_workers,
+        pin_memory=torch.cuda.is_available(),
+        drop_last=True
+    )
+    
+    val_loader = DataLoader(
+        val_dataset,
+        batch_size=batch_size,
+        shuffle=False,
+        num_workers=num_workers,
+        pin_memory=torch.cuda.is_available()
+    )
+    
+    print(f"📦 Data loaders created:")
+    print(f"   Batch size: {batch_size}")
+    print(f"   Num workers: {num_workers}")
+    print(f"   Train batches: {len(train_loader)}")
+    print(f"   Val batches: {len(val_loader)}")
+    print(f"   Weighted sampling: {use_weighted_sampling}")
+    
+    return train_loader, val_loader
+
+def create_weighted_sampler(dataset: FireDetectionDataset) -> WeightedRandomSampler:
+    """
+    Create weighted random sampler for imbalanced datasets
+    
+    Args:
+        dataset: The dataset to create sampler for
+    
+    Returns:
+        WeightedRandomSampler for balanced sampling
+    """
+    # Count samples per class
+    class_counts = Counter([label for _, label in dataset.samples])
+    total_samples = len(dataset.samples)
+    
+    # Calculate weights (inverse frequency)
+    class_weights = {}
+    for class_idx, count in class_counts.items():
+        class_weights[class_idx] = total_samples / count
+    
+    # Create sample weights
+    sample_weights = [class_weights[label] for _, label in dataset.samples]
+    
+    # Create sampler
+    sampler = WeightedRandomSampler(
+        weights=sample_weights,
+        num_samples=total_samples,
+        replacement=True
+    )
+    
+    print(f"⚖️ Weighted sampler created:")
+    for class_name, class_idx in dataset.class_to_idx.items():
+        count = class_counts.get(class_idx, 0)
+        weight = class_weights.get(class_idx, 0)
+        print(f"   {class_name}: {count} samples, weight: {weight:.2f}")
+    
+    return sampler
+
+def get_inference_transform(image_size: int = 224) -> transforms.Compose:
+    """
+    Get transforms for inference/prediction
+    
+    Args:
+        image_size: Size to resize images to
+    
+    Returns:
+        Transform pipeline for inference
+    """
+    return transforms.Compose([
+        transforms.Resize((image_size, image_size)),
+        transforms.ToTensor(),
+        transforms.Normalize(
+            mean=[0.485, 0.456, 0.406],
+            std=[0.229, 0.224, 0.225]
+        )
+    ])
+
+def prepare_image_for_inference(image: Image.Image, transform: transforms.Compose) -> torch.Tensor:
+    """
+    Prepare an image for inference
+    
+    Args:
+        image: PIL Image
+        transform: Transform pipeline
+    
+    Returns:
+        Tensor ready for model inference
+    """
+    # Apply transforms
+    image_tensor = transform(image)
+    
+    # Add batch dimension
+    image_tensor = image_tensor.unsqueeze(0)
+    
+    return image_tensor
+
+def visualize_batch(data_loader: DataLoader, num_samples: int = 8) -> None:
+    """
+    Visualize a batch of images from the data loader
+    
+    Args:
+        data_loader: DataLoader to sample from
+        num_samples: Number of samples to visualize
+    """
+    import matplotlib.pyplot as plt
+    
+    # Get a batch
+    images, labels = next(iter(data_loader))
+    
+    # Denormalize images
+    mean = torch.tensor([0.485, 0.456, 0.406]).view(3, 1, 1)
+    std = torch.tensor([0.229, 0.224, 0.225]).view(3, 1, 1)
+    
+    # Create figure
+    fig, axes = plt.subplots(2, 4, figsize=(15, 8))
+    axes = axes.flatten()
+    
+    class_names = ['Fire', 'No Fire']
+    
+    for i in range(min(num_samples, len(images))):
+        # Denormalize
+        img = images[i] * std + mean
+        img = torch.clamp(img, 0, 1)
+        
+        # Convert to numpy
+        img_np = img.permute(1, 2, 0).numpy()
+        
+        # Plot
+        axes[i].imshow(img_np)
+        axes[i].set_title(f'{class_names[labels[i]]}')
+        axes[i].axis('off')
+    
+    plt.tight_layout()
+    plt.show()
+
+def check_data_directory(data_dir: str) -> Dict[str, int]:
+    """
+    Check data directory structure and count samples
+    
+    Args:
+        data_dir: Directory to check
+    
+    Returns:
+        Dictionary with data counts
+    """
+    data_counts = {}
+    
+    if not os.path.exists(data_dir):
+        print(f"❌ Data directory not found: {data_dir}")
+        return data_counts
+    
+    print(f"📊 Data Directory Analysis: {data_dir}")
+    print("=" * 50)
+    
+    total_samples = 0
+    
+    for split in ['train', 'val']:
+        split_dir = os.path.join(data_dir, split)
+        if not os.path.exists(split_dir):
+            continue
+        
+        print(f"\n{split.upper()} SET:")
+        split_total = 0
+        
+        for class_name in ['fire', 'no_fire']:
+            class_dir = os.path.join(split_dir, class_name)
+            if not os.path.exists(class_dir):
+                continue
+            
+            # Count images recursively
+            count = 0
+            for root, dirs, files in os.walk(class_dir):
+                for file in files:
+                    if file.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp', '.tiff')):
+                        count += 1
+            
+            print(f"   {class_name}: {count} images")
+            data_counts[f"{split}_{class_name}"] = count
+            split_total += count
+        
+        print(f"   Total {split}: {split_total}")
+        total_samples += split_total
+        data_counts[f"{split}_total"] = split_total
+    
+    print(f"\nOVERALL TOTAL: {total_samples} images")
+    data_counts['total'] = total_samples
+    print("=" * 50)
+    
+    return data_counts
+
+def create_sample_data_structure():
+    """Create sample data structure for testing"""
+    print("🔥 Creating sample fire detection data structure...")
+    
+    # Create directories
+    directories = [
+        'data/train/fire',
+        'data/train/no_fire',
+        'data/val/fire',
+        'data/val/no_fire'
+    ]
+    
+    for directory in directories:
+        os.makedirs(directory, exist_ok=True)
+    
+    print("✅ Sample data structure created")
+    print("   Please add your fire detection images to the appropriate directories")
+    print("   - data/train/fire/     (training fire images)")
+    print("   - data/train/no_fire/  (training no-fire images)")
+    print("   - data/val/fire/       (validation fire images)")
+    print("   - data/val/no_fire/    (validation no-fire images)") 

model_utils.py

@@ -0,0 +1,239 @@
+"""
+Model utilities for fire detection classification
+Handles ConvNeXt model loading and adaptation for transfer learning
+"""
+
+import torch
+import torch.nn as nn
+import timm
+import os
+from typing import Dict, Any, Optional, Tuple
+
+class FireDetectionClassifier(nn.Module):
+    """
+    ConvNeXt-based fire detection classifier
+    Uses transfer learning from ImageNet pretrained model
+    """
+    
+    def __init__(self, num_classes: int = 2, pretrained: bool = True):
+        super(FireDetectionClassifier, self).__init__()
+        
+        # Load ConvNeXt Large model
+        self.backbone = timm.create_model(
+            'convnext_large.fb_in22k_ft_in1k', 
+            pretrained=pretrained, 
+            num_classes=0  # Remove classification head
+        )
+        
+        # Get feature dimensions
+        self.feature_dim = self.backbone.num_features
+        
+        # Custom classification head for fire detection
+        self.classifier = nn.Sequential(
+            nn.LayerNorm(self.feature_dim),
+            nn.Linear(self.feature_dim, 512),
+            nn.ReLU(inplace=True),
+            nn.Dropout(0.3),
+            nn.Linear(512, 128),
+            nn.ReLU(inplace=True),
+            nn.Dropout(0.2),
+            nn.Linear(128, num_classes)
+        )
+        
+        # Initialize classifier weights
+        self._init_classifier_weights()
+    
+    def _init_classifier_weights(self):
+        """Initialize classifier weights using Xavier initialization"""
+        for module in self.classifier.modules():
+            if isinstance(module, nn.Linear):
+                nn.init.xavier_uniform_(module.weight)
+                nn.init.constant_(module.bias, 0)
+    
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        """Forward pass through the model"""
+        # Extract features using ConvNeXt backbone
+        features = self.backbone(x)
+        
+        # Classify using custom head
+        output = self.classifier(features)
+        
+        return output
+    
+    def freeze_backbone(self):
+        """Freeze backbone parameters for transfer learning"""
+        for param in self.backbone.parameters():
+            param.requires_grad = False
+        print("🔒 Backbone frozen for transfer learning")
+    
+    def unfreeze_backbone(self):
+        """Unfreeze backbone parameters for fine-tuning"""
+        for param in self.backbone.parameters():
+            param.requires_grad = True
+        print("🔓 Backbone unfrozen for fine-tuning")
+    
+    def get_parameter_count(self) -> Dict[str, int]:
+        """Get parameter counts for different parts of the model"""
+        backbone_params = sum(p.numel() for p in self.backbone.parameters())
+        classifier_params = sum(p.numel() for p in self.classifier.parameters())
+        total_params = backbone_params + classifier_params
+        
+        trainable_params = sum(p.numel() for p in self.parameters() if p.requires_grad)
+        
+        return {
+            'backbone': backbone_params,
+            'classifier': classifier_params,
+            'total': total_params,
+            'trainable': trainable_params
+        }
+
+def create_fire_detection_model(
+    num_classes: int = 2,
+    freeze_backbone: bool = True
+) -> FireDetectionClassifier:
+    """
+    Create fire detection classifier model with transfer learning
+    
+    Args:
+        num_classes: Number of output classes (2 for fire/no_fire)
+        freeze_backbone: Whether to freeze backbone for transfer learning
+    
+    Returns:
+        FireDetectionClassifier model ready for training
+    """
+    print("🔥 Creating fire detection classifier...")
+    
+    # Create the model
+    model = FireDetectionClassifier(num_classes=num_classes, pretrained=True)
+    
+    # Freeze backbone if requested
+    if freeze_backbone:
+        model.freeze_backbone()
+    
+    # Print model information
+    param_counts = model.get_parameter_count()
+    print(f"📊 Model Statistics:")
+    print(f"   Backbone parameters: {param_counts['backbone']:,}")
+    print(f"   Classifier parameters: {param_counts['classifier']:,}")
+    print(f"   Total parameters: {param_counts['total']:,}")
+    print(f"   Trainable parameters: {param_counts['trainable']:,}")
+    print(f"   Model size: ~{param_counts['total'] * 4 / 1024**2:.1f} MB")
+    
+    return model
+
+def save_model(
+    model: FireDetectionClassifier,
+    save_path: str,
+    epoch: int,
+    best_acc: float,
+    optimizer_state: Optional[Dict] = None,
+    additional_info: Optional[Dict] = None
+) -> None:
+    """
+    Save model checkpoint with training information
+    
+    Args:
+        model: The model to save
+        save_path: Path to save the model
+        epoch: Current epoch number
+        best_acc: Best accuracy achieved
+        optimizer_state: Optimizer state dict
+        additional_info: Additional information to save
+    """
+    # Create directory if it doesn't exist
+    os.makedirs(os.path.dirname(save_path), exist_ok=True)
+    
+    # Prepare checkpoint
+    checkpoint = {
+        'model_state_dict': model.state_dict(),
+        'epoch': epoch,
+        'best_acc': best_acc,
+        'model_info': {
+            'num_classes': 2,
+            'class_names': ['fire', 'no_fire'],
+            'parameter_count': model.get_parameter_count()
+        }
+    }
+    
+    # Add optional information
+    if optimizer_state:
+        checkpoint['optimizer_state_dict'] = optimizer_state
+    
+    if additional_info:
+        checkpoint.update(additional_info)
+    
+    # Save checkpoint
+    torch.save(checkpoint, save_path)
+    print(f"💾 Model saved to: {save_path}")
+    print(f"📈 Best accuracy: {best_acc:.4f}")
+
+def load_model(
+    model_path: str,
+    num_classes: int = 2,
+    device: str = 'cpu'
+) -> Tuple[FireDetectionClassifier, Dict[str, Any]]:
+    """
+    Load a trained fire detection model
+    
+    Args:
+        model_path: Path to the saved model
+        num_classes: Number of classes (should be 2)
+        device: Device to load model on
+    
+    Returns:
+        Tuple of (model, model_info)
+    """
+    if not os.path.exists(model_path):
+        raise FileNotFoundError(f"Model not found at: {model_path}")
+    
+    # Load checkpoint
+    checkpoint = torch.load(model_path, map_location=device)
+    
+    # Create model
+    model = FireDetectionClassifier(num_classes=num_classes, pretrained=False)
+    
+    # Load state dict
+    model.load_state_dict(checkpoint['model_state_dict'])
+    
+    # Move to device
+    model = model.to(device)
+    
+    # Extract model info
+    model_info = checkpoint.get('model_info', {})
+    model_info['epoch'] = checkpoint.get('epoch', 'Unknown')
+    model_info['best_acc'] = checkpoint.get('best_acc', 'Unknown')
+    
+    print(f"✅ Model loaded from: {model_path}")
+    print(f"📊 Model accuracy: {model_info.get('best_acc', 'Unknown')}")
+    
+    return model, model_info
+
+def get_model_summary(model: FireDetectionClassifier) -> str:
+    """
+    Get a summary of the model architecture
+    
+    Args:
+        model: The model to summarize
+    
+    Returns:
+        String summary of the model
+    """
+    param_counts = model.get_parameter_count()
+    
+    summary = f"""
+🔥 Fire Detection Model Summary
+{'='*50}
+Architecture: ConvNeXt Large + Custom Classifier
+Classes: fire, no_fire
+
+Parameters:
+  Backbone: {param_counts['backbone']:,}
+  Classifier: {param_counts['classifier']:,}
+  Total: {param_counts['total']:,}
+  Trainable: {param_counts['trainable']:,}
+
+Model Size: ~{param_counts['total'] * 4 / 1024**2:.1f} MB
+{'='*50}
+"""
+    
+    return summary