Merge remote-tracking branch 'origin/master' into pso-and-experiments

Author: seofernando25

framework/data_utils.py

@@ -6,11 +6,72 @@
 from datasets import load_from_disk
 from sklearn.model_selection import train_test_split
 from torch.utils.data import DataLoader
+from PIL import Image
 
 from framework import utils
 from framework.datasets import CIFAR10Dataset
 
 
+def convert_to_grayscale(image: np.ndarray) -> np.ndarray:
+    """Convert RGB/RGBA image to grayscale.
+    
+    Args:
+        image: Image array with shape (H, W, C) where C is 3 (RGB) or 4 (RGBA)
+        
+    Returns:
+        Grayscale image with shape (H, W) 
+    """
+    if len(image.shape) == 2:
+        # Already grayscale
+        return image
+    elif len(image.shape) == 3:
+        if image.shape[2] == 1:
+            # Single channel, just squeeze
+            return image.squeeze(axis=2)
+        elif image.shape[2] == 3:
+            # RGB -> Grayscale using luminance weights
+            # Y = 0.2125R + 0.7154G + 0.0721B
+            return np.dot(image[...,:3], [0.2125, 0.7154, 0.0721])
+        elif image.shape[2] == 4:
+            # RGBA -> Grayscale (ignore alpha)
+            return np.dot(image[...,:3], [0.2125, 0.7154, 0.0721])
+    
+    raise ValueError(f"Unsupported image shape: {image.shape}")
+
+
+def preprocess_images_to_grayscale(images: List[np.ndarray]) -> List[np.ndarray]:
+    """Convert a list of images to grayscale.
+    
+    Args:
+        images: List of image arrays
+        
+    Returns:
+        List of grayscale image arrays
+    """
+    return [convert_to_grayscale(img) for img in images]
+
+
+def convert_dataset_to_grayscale(dataset):
+    """Convert HuggingFace dataset images to grayscale in-place preprocessing.
+    
+    Args:
+        dataset: HuggingFace dataset with 'image' column
+        
+    Returns:
+        List of grayscale images and labels
+    """
+    images = []
+    labels = []
+    
+    for item in dataset:
+        img = np.array(item['image'])
+        gray_img = convert_to_grayscale(img)
+        images.append(gray_img)
+        labels.append(item['label'])
+    
+    return images, np.array(labels)
+
+
 def load_cifar10_data():
     """Load CIFAR-10 dataset (grayscale from processed datasets)."""
     repo_root = Path(__file__).resolve().parents[1]

models/cnn.py

@@ -2,7 +2,7 @@
 
 from dataclasses import dataclass
 from pathlib import Path
-from typing import Dict, List, Optional
+from typing import Dict, List, Optional, List
 
 import numpy as np
 import torch
@@ -24,6 +24,9 @@
 from .ParamSpace import ParamSpace
 from .base import BaseModel
 
+import numpy as np
+from torch import tensor
+
 MODEL_PATH = Path(".cache/models/cnn_cifar.pth")
 
 
@@ -89,6 +92,7 @@ def __init__(
             nn.Linear(64, num_classes),
         )
 
+    # Will be run by PyTorch under the hood
     def forward(self, x: torch.Tensor) -> torch.Tensor:
         x = self.features(x)
         return self.classifier(x)
@@ -104,13 +108,25 @@ def __init__(self, num_classes: int = 10) -> None:
         self._input_channels = 1  # grayscale CIFAR-10
 
     def create_model(self, **params) -> None:
+        """Create the CNN model with given parameters."""        
+        # Extract architecture-specific parameters for Backbone creation      
+        kernel_size = params.get("kernel_size", 3)
+        stride = params.get("stride", 1)
+        learning_rate = params.get("learning_rate", 3e-4)
+        batch_size = params.get("batch_size", 64)
+        weight_decay = params.get("weight_decay", 1e-3)
+        optimizer = params.get("optimizer", "AdamW")
+        
+        # Also ensure default values are set if not provided
+        self.params.setdefault("kernel_size", kernel_size)
+        self.params.setdefault("stride", stride)
+        self.params.setdefault("learning_rate", learning_rate)
+        self.params.setdefault("batch_size", batch_size)
+        self.params.setdefault("weight_decay", weight_decay)
+        self.params.setdefault("optimizer", optimizer)
+
         # Store all parameters passed in
         self.params.update(params)
-        
-        # Extract architecture-specific parameters for Backbone creation
-        kernel_size = self.params.get("kernel_size", 3)
-        stride = self.params.get("stride", 1)
-        
         self.network = Backbone(
             in_channels=self._input_channels,
             num_classes=self.num_classes,
@@ -120,48 +136,39 @@ def create_model(self, **params) -> None:
 
     def train(
         self,
-        X_train: List[np.ndarray],
-        y_train: np.ndarray,
-        X_val: List[np.ndarray],
-        y_val: np.ndarray,
+        train_data,  # Can be DataLoader or raw data
+        val_data,    # Can be DataLoader or raw data  
+        config: Optional[TrainingConfig] = None,
         device: Optional[torch.device] = None,
-        epochs: Optional[int] = None,
-        patience: Optional[int] = None,
-        min_delta: Optional[float] = None,
-        checkpoint_path: Optional[Path] = None,
-        grad_clip_norm: Optional[float] = None,
-        writer: Optional[SummaryWriter] = None,
-        num_workers: int = 2,
         verbose: bool = True,
     ) -> Dict[str, float]:
         if self.network is None:
             raise RuntimeError("Train called before model is initialized")
         device = device or get_device()
         self.network = self.network.to(device)
 
-        default_config = TrainingConfig()
-        config = TrainingConfig(
-            learning_rate=float(self.params.get("learning_rate", default_config.learning_rate)),
-            weight_decay=float(self.params.get("weight_decay", default_config.weight_decay)),
-            optimizer=self.params.get("optimizer", default_config.optimizer),
-            batch_size=int(self.params.get("batch_size", default_config.batch_size)),
-            # Infrastructure params: use provided values or defaults
-            epochs=epochs if epochs is not None else default_config.epochs,
-            patience=patience if patience is not None else default_config.patience,
-            min_delta=min_delta if min_delta is not None else default_config.min_delta,
-            checkpoint_path=checkpoint_path if checkpoint_path is not None else default_config.checkpoint_path,
-            grad_clip_norm=grad_clip_norm if grad_clip_norm is not None else default_config.grad_clip_norm,
-            writer=writer if writer is not None else default_config.writer,
-        )
-
-        train_loader, val_loader = create_dataloaders(
-            X_train,
-            y_train,
-            X_val,
-            y_val,
-            batch_size=config.batch_size,
-            num_workers=num_workers,
-        )
+        config = config or TrainingConfig()
+        
+        # Handle different input types - convert to DataLoaders if needed
+        if isinstance(train_data, DataLoader) and isinstance(val_data, DataLoader):
+            train_loader = train_data
+            val_loader = val_data
+        else:
+            # Raw data provided - create DataLoaders
+            from framework.data_utils import create_dataloaders
+            if hasattr(train_data, '__len__') and hasattr(val_data, '__len__'):
+                # Assume train_data is X_train, val_data is y_train for backwards compatibility
+                if len(train_data) > 0 and not isinstance(train_data[0], (int, float)):
+                    # This looks like X_train, y_train, X_test, y_test pattern
+                    # Need to extract from the calling pattern
+                    batch_size = getattr(config, 'batch_size', 64)
+                    train_loader, val_loader = create_dataloaders(
+                        train_data, val_data, train_data[:len(val_data)], val_data, batch_size=batch_size
+                    )
+                else:
+                    raise ValueError("Invalid data format provided to CNN.train()")
+            else:
+                raise ValueError("Invalid data format provided to CNN.train()")
 
         optimizer = self._build_optimizer(self.network, config)
         scheduler = optim.lr_scheduler.OneCycleLR(
@@ -256,7 +263,8 @@ def train(
 
     def predict(
         self,
-        data_loader: DataLoader,
+        data: DataLoader | List | np.ndarray,
+        labels: List | np.ndarray = None,
         device: Optional[torch.device] = None,
         return_probabilities: bool = False,
     ) -> torch.Tensor:
@@ -266,6 +274,19 @@ def predict(
         network = self.network.to(device)
         network.eval()
 
+        # Handle data input - create DataLoader if raw data is provided
+        if isinstance(data, DataLoader):
+            data_loader = data
+        else:
+            # Raw data provided - create DataLoader
+            from framework.data_utils import create_dataloaders
+            if labels is None:
+                labels = [0] * len(data)  # dummy labels for prediction
+            batch_size = getattr(self, 'params', {}).get('batch_size', 64)
+            _, data_loader = create_dataloaders(
+                data[:1], [labels[0]], data, labels, batch_size=batch_size
+            )
+
         outputs = []
         with torch.no_grad():
             for images, _ in data_loader:
@@ -279,8 +300,8 @@ def predict(
 
     def evaluate(
         self,
-        X: List[np.ndarray],
-        y: np.ndarray,
+        data,  # Can be DataLoader or raw data
+        labels=None,  # Required if data is not DataLoader
         device: Optional[torch.device] = None,
         criterion: Optional[nn.Module] = None,
         num_workers: int = 0,
@@ -291,16 +312,18 @@ def evaluate(
         network = self.network.to(device)
         network.eval()
 
-        default_config = TrainingConfig()
-        batch_size = int(self.params.get("batch_size", default_config.batch_size))
-        dataset = CIFAR10Dataset(X, y)
-        data_loader = DataLoader(
-            dataset,
-            batch_size=batch_size,
-            shuffle=False,
-            num_workers=num_workers,
-            pin_memory=torch.cuda.is_available(),
-        )
+        # Handle data input - create DataLoader if raw data is provided
+        if isinstance(data, DataLoader):
+            data_loader = data
+        else:
+            # Raw data provided - create DataLoader
+            from framework.data_utils import create_dataloaders
+            if labels is None:
+                raise ValueError("labels must be provided when data is not a DataLoader")
+            batch_size = getattr(self, 'params', {}).get('batch_size', 64)
+            _, data_loader = create_dataloaders(
+                data[:1], [labels[0]], data, labels, batch_size=batch_size
+            )
 
         criterion = criterion or nn.CrossEntropyLoss()
 
@@ -338,19 +361,26 @@ def evaluate(
         f1_macro = report["macro avg"]["f1-score"]
         f1_micro = report.get("micro avg", {}).get("f1-score", f1_score(y_true, y_pred, average="micro", zero_division=0))
 
-        roc_auc = roc_auc_score(y_true, y_proba, average="macro", multi_class="ovr")
-
-        avg_loss = total_loss / len(data_loader)
-
-        return {
-            "loss": avg_loss,
+        # Initialize metrics without ROC AUC first
+        metrics = {
+            "loss": total_loss / len(data_loader), # average loss
             "accuracy": accuracy,
             "precision_macro": precision_macro,
             "recall_macro": recall_macro,
             "f1_macro": f1_macro,
             "f1_micro": f1_micro,
-            "roc_auc": roc_auc,
         }
+        
+        # Try to calculate ROC AUC, but handle potential errors gracefully
+        try:
+            roc_auc = roc_auc_score(y_true, y_proba, average="macro", multi_class="ovr")
+            metrics["roc_auc"] = roc_auc
+        except ValueError as e:
+            # ROC AUC calculation failed (likely due to insufficient samples per class)
+            # Continue without ROC AUC metric
+            print(f"Warning: Could not calculate ROC AUC: {e}")
+
+        return metrics
 
     def get_param_space(self) -> Dict[str, ParamSpace]:
         return {

models/decision_tree.py

@@ -57,22 +57,37 @@ def evaluate(self, X_test, y_test) -> Dict[str, float]:
         report = classification_report(
             y_test, predictions, output_dict=True, zero_division=0
         )
-
-        proba = self.estimator.predict_proba(X_test)
 
+        # Initialize metrics without ROC AUC first
         metrics: Dict[str, float] = {
             "accuracy": report["accuracy"],
             "precision_macro": report["macro avg"]["precision"],
             "recall_macro": report["macro avg"]["recall"],
             "f1_macro": report["macro avg"]["f1-score"],
             "f1_micro": report.get("micro avg", {}).get("f1-score", f1_score(y_test, predictions, average="micro", zero_division=0)),
-            "roc_auc": roc_auc_score(y_test, proba, average="macro", multi_class="ovr"),
             "precision_weighted": report["weighted avg"]["precision"],
             "recall_weighted": report["weighted avg"]["recall"],
             "f1_weighted": report["weighted avg"]["f1-score"],
-            "roc_auc_weighted": roc_auc_score(y_test, proba, average="weighted", multi_class="ovr"),
         }
 
+        # Add ROC AUC if possible (with proper error handling)
+        if hasattr(self.estimator, "predict_proba"):
+            try:
+                proba = self.estimator.predict_proba(X_test)
+                if proba.ndim == 2 and proba.shape[1] > 1:
+                    # Check if we have enough classes for ROC AUC calculation
+                    unique_classes = len(set(y_test))
+                    if unique_classes >= 2 and proba.shape[1] == len(set(y_test)):
+                        metrics["roc_auc"] = roc_auc_score(
+                            y_test, proba, average="macro", multi_class="ovr"
+                        )
+                        metrics["roc_auc_weighted"] = roc_auc_score(
+                            y_test, proba, average="weighted", multi_class="ovr"
+                        )
+            except (ValueError, Exception) as e:
+                # ROC AUC calculation failed, skip it
+                pass
+
         return metrics
 
     def get_param_space(self) -> Dict[str, ParamSpace]:

models/knn.py

@@ -56,20 +56,37 @@ def evaluate(self, X_test, y_test) -> Dict[str, float]:
         report = classification_report(
             y_test, predictions, output_dict=True, zero_division=0
         )
-        proba = self.estimator.predict_proba(X_test)
 
+        # Initialize metrics without ROC AUC first
         metrics: Dict[str, float] = {
             "accuracy": report["accuracy"],
             "precision_macro": report["macro avg"]["precision"],
             "recall_macro": report["macro avg"]["recall"],
             "f1_macro": report["macro avg"]["f1-score"],
             "f1_micro": report.get("micro avg", {}).get("f1-score", f1_score(y_test, predictions, average="micro", zero_division=0)),
-            "roc_auc": roc_auc_score(y_test, proba, average="macro", multi_class="ovr"),
             "precision_weighted": report["weighted avg"]["precision"],
             "recall_weighted": report["weighted avg"]["recall"],
             "f1_weighted": report["weighted avg"]["f1-score"],
         }
-        
+
+        # Add ROC AUC if possible (with proper error handling)
+        if hasattr(self.estimator, "predict_proba"):
+            try:
+                proba = self.estimator.predict_proba(X_test)
+                if proba.ndim == 2 and proba.shape[1] > 1:
+                    # Check if we have enough classes for ROC AUC calculation
+                    unique_classes = len(set(y_test))
+                    if unique_classes >= 2 and proba.shape[1] == len(set(y_test)):
+                        metrics["roc_auc"] = roc_auc_score(
+                            y_test, proba, average="macro", multi_class="ovr"
+                        )
+                        metrics["roc_auc_weighted"] = roc_auc_score(
+                            y_test, proba, average="weighted", multi_class="ovr"
+                        )
+            except (ValueError, Exception) as e:
+                # ROC AUC calculation failed, skip it
+                pass
+
         return metrics
 
     def get_param_space(self) -> Dict[str, ParamSpace]: