Lab Report: Object Detection Task Using LGG Dataset

from torchvision import transforms

# Define transformations
data_transforms = transforms.Compose([
    transforms.ToTensor(),
    transforms.RandomHorizontalFlip(0.5),
    transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.1)
])

import torchvision
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor

# Load Faster R-CNN pre-trained on COCO
model = torchvision.models.detection.fasterrcnn_resnet50_fpn(weights=True)

# Modify the box predictor to match our dataset (1 class + background)
num_classes = 2   # 1 (LGG) + 1 (background)
in_features = model.roi_heads.box_predictor.cls_score.in_features
model.roi_heads.box_predictor = FastRCNNPredictor(in_features, num_classes)

from tqdm import tqdm

def train_one_epoch(model, dataloader, optimizer, device="cpu"):
    """
    Train the Faster R-CNN model for one epoch.

    Args:
        model (torch.nn.Module): Faster R-CNN model to train.
        dataloader (torch.utils.data.DataLoader): DataLoader for training data.
        optimizer (torch.optim.Optimizer): Optimizer for training.
        device (str or torch.device): Device to run the training on (default "cpu").

    Returns:
        float: Average training loss for the epoch.
    """
    model.train()  # Set the model to training mode
    total_loss = 0  # Track total loss
    for images, targets in tqdm(dataloader, desc="Training"):
        # Move data to device
        images = [img.to(device) for img in images]
        targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
        # Forward pass and compute loss
        loss_dict = model(images, targets)  # Loss dictionary
        loss = sum(loss for loss in loss_dict.values())  # Total loss
        # Backward pass and optimization
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        total_loss += loss.item()  # Accumulate loss
    # Compute average loss
    average_loss = total_loss / len(dataloader)
    # print(f"Training Loss: {average_loss:.4f}")
    return average_loss

def evaluate_manual_mAP(model, dataloader, device="cpu"):
    """
    Evaluate the Faster R-CNN model on the validation set and compute mAP manually.
    Args:
        model (torch.nn.Module): Faster R-CNN model to evaluate.
        dataloader (torch.utils.data.DataLoader): DataLoader for validation data.
        device (str or torch.device): Device to run the evaluation on (default "cpu").
    Returns:
        dict: Simplified metrics including mAP, precision, and recall.
    """
    model.eval()  # Set model to evaluation mode
    all_predictions = []
    all_targets = []

    with torch.no_grad():
        for images, targets in tqdm(dataloader, desc="Evaluating"):
            # Move data to device
            images = [img.to(device) for img in images]
            targets = [{k: v.to(device) for k, v in t.items()} for t in targets]
            # Forward pass to get predictions
            outputs = model(images)
            # Store predictions and targets
            all_predictions.extend(outputs)
            all_targets.extend(targets)

    # Compute mAP manually (example implementation)
    metrics = compute_metrics(all_predictions, all_targets)
    # print(f"Validation - mAP: {metrics['map']:.4f}, Precision: {metrics['precision']:.4f}, Recall: {metrics['recall']:.4f}")
    return metrics