code/train.py

import os
os.environ['PYTORCH_CUDA_ALLOC_CONF'] = 'expandable_segments:True'
os.environ['NCCL_TIMEOUT'] = '120'
import torch
import torch.nn as nn
from dataset.hpadataset import HPADatasetMIL, HPADatasetDownsample, SeededSampler, ResumableDistributedSampler, HPADatasetMIL_url
from torch.utils.data import DataLoader
import pandas as pd
from tqdm import tqdm
import numpy as np
import time
import configparser
import argparse
import wandb
import copy
from sklearn.model_selection import train_test_split
from datetime import datetime
from os.path import join
import matplotlib.pyplot as plt
import seaborn as sns
from sklearn.metrics import confusion_matrix
import torch.cuda.amp  # Import Automatic Mixed Precision
import re
import random

"""
Be sure to init conda env:
conda activate hpa

Note: If running on a non-H100 GPU, there might be an error for pyarrow.
Reinstall pyarrow with following command will resolve the issue:
conda install -c conda-forge pyarrow
"""

"""
CrossEntropyLoss is better than BCEWithLogitsLoss for multi-class classification
because it directly handles mutually exclusive classes by computing probabilities
over all classes, whereas BCEWithLogitsLoss treats each class independently.

If use BCEWithLogitsLoss for multi-class classification, it treats each class as
independent and computes a separate binary classification for each class. This is
problematic because, in multi-class classification, only one class can be correct
for each input, and the classes are mutually exclusive. BCEWithLogitsLoss doesn't
enforce this exclusivity, leading to incorrect probability distributions where
multiple classes could be predicted as "active" at once, which is not appropriate
for multi-class problems where exactly one class should be correct.
"""
criterion_intensity = nn.CrossEntropyLoss()  # For intensity
criterion_location = nn.CrossEntropyLoss()   # For location
criterion_quantity = nn.CrossEntropyLoss()   # For quantity
criterion_tissue = nn.CrossEntropyLoss()     # For tissue type (58 classes)
criterion_malignancy = nn.CrossEntropyLoss() # For tumor vs non-tumor (binary)

# ---------------------------------------------------------------- data locations
# HPA10M metadata and RLE tissue masks (environment variables; see README).
DATA_ROOT = os.environ.get("ISIGHT_DATA_ROOT", "")
DATA_TRAIN_META = os.environ.get("ISIGHT_TRAIN_META",
                                 os.path.join(DATA_ROOT, "remaining_training_metadata.feather"))
DATA_TEST_META = os.environ.get("ISIGHT_TEST_META",
                                os.path.join(DATA_ROOT, "validation_metadata.feather"))
DATA_RLE_DIR = os.environ.get("ISIGHT_RLE_DIR", os.path.join(DATA_ROOT, "rle_masks"))
# only used by the `simple_downsample` model version
DATA_IMAGE_DIR = os.environ.get("ISIGHT_IMAGE_DIR", os.path.join(DATA_ROOT, "hpa10m"))
DATA_RLE_INDEX = os.environ.get("ISIGHT_RLE_INDEX",
                                os.path.join(DATA_ROOT, "rle_mask_index.json"))
# Step ReduceLROnPlateau once per epoch instead of once per batch (default: per batch).
SCHEDULER_PER_EPOCH = os.environ.get("SCHEDULER_PER_EPOCH", "0") == "1"


def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)

def custom_collate_fn(batch):
    # If no valid samples, return None
    batch_valid = [batch_item for batch_item in batch if batch_item is not None]
    
    if len(batch_valid) == 0:
        return None
    images, metadata, query_input, caption_output, snomed_text, snomed_code, image_url, staining_intensity, staining_location, staining_quantity, malignancy, tissue_one_hot, cell_type_one_hot = zip(*batch_valid)
    # Note: here images is a list of (N, 3, 224, 224), where N is the number of patches per image, and is different for each image.
    
    # Convert to tensors and stack
    staining_intensity = torch.stack([torch.tensor(sublist, dtype=torch.float) for sublist in staining_intensity])
    staining_location = torch.stack([torch.tensor(sublist, dtype=torch.float) for sublist in staining_location])
    staining_quantity = torch.stack([torch.tensor(sublist, dtype=torch.float) for sublist in staining_quantity])
    malignancy = torch.stack([torch.tensor(sublist, dtype=torch.float) for sublist in malignancy])
    tissue_one_hot = torch.stack([torch.tensor(sublist, dtype=torch.float) for sublist in tissue_one_hot])
    cell_type_one_hot = torch.stack([torch.tensor(sublist, dtype=torch.float) for sublist in cell_type_one_hot])
    
    return images, metadata, query_input, caption_output, snomed_text, snomed_code, image_url, staining_intensity, staining_location, staining_quantity, malignancy, tissue_one_hot, cell_type_one_hot


def train(config):
    # Setup
    if config.distributed:
        # Set the device and get the rank of the current process
        device = torch.device('cuda', config.local_rank)
        torch.cuda.set_device(device)
        rank = torch.distributed.get_rank()
    else:
        device = "cuda" if torch.cuda.is_available() else "cpu"
        rank = 0  # Default rank for non-distributed training

    # Initialize AMP scaler
    scaler = torch.cuda.amp.GradScaler()  # For half-precision scaling

    base_model_name = config.base_model_name

    # Model initialization (unchanged)
    if config.model_version == "v1":
        from model.patch_encoder_with_clam import create_clam_vit
        model, optimizer, scheduler = create_clam_vit(
            base_model_name=base_model_name,
            device=device,
            lr=config.learning_rate,
            weight_decay=config.weight_decay,
            freeze_vit=config.freeze_vit,
            freeze_query_features_encoder=config.freeze_query_features_encoder,
            use_cell_type_embedding=config.use_cell_type_embedding,
        )
    elif config.model_version == "v2":
        from model.patch_encoder_with_clam_v2 import create_clam_vit
        model, optimizer, scheduler = create_clam_vit(
            base_model_name=base_model_name, 
            device=device,
            lr=config.learning_rate,
            weight_decay=config.weight_decay,
            freeze_vit=config.freeze_vit,
            freeze_query_features_encoder=config.freeze_query_features_encoder,
            use_cell_type_embedding=config.use_cell_type_embedding
        )
    elif config.model_version == "v3":
        from model.patch_encoder_with_clam_v3 import create_clam_vit
        model, optimizer, scheduler = create_clam_vit(
            base_model_name=base_model_name,
            training_config=config,
            device=device,
            lr=config.learning_rate,
            weight_decay=config.weight_decay,
            freeze_vit=config.freeze_vit,
            freeze_query_features_encoder=config.freeze_query_features_encoder,
            use_cell_type_embedding=config.use_cell_type_embedding
        )
    elif config.model_version == "v3_all_tokens":
        from model.patch_encoder_with_clam_v3_all_tokens import create_clam_vit
        model, optimizer, scheduler = create_clam_vit(
            base_model_name=base_model_name,
            training_config=config,
            device=device,
            lr=config.learning_rate,
            weight_decay=config.weight_decay,
            freeze_vit=config.freeze_vit,
            freeze_query_features_encoder=config.freeze_query_features_encoder,
            use_cell_type_embedding=config.use_cell_type_embedding
        )
    elif config.model_version == "simple_downsample":
        from model.simple_downsample import create_vit
        model, optimizer, scheduler = create_vit(
            base_model_name=base_model_name, 
            device=device,
            lr=config.learning_rate,
            weight_decay=config.weight_decay,
            freeze_vit=config.freeze_vit,
            freeze_query_features_encoder=config.freeze_query_features_encoder,
            use_cell_type_embedding=config.use_cell_type_embedding
        )
    model.to(device)

    # Wrap the model with DistributedDataParallel if using distributed training
    if config.distributed:
        model = torch.nn.parallel.DistributedDataParallel(
            model,
            device_ids=[config.local_rank],
            output_device=config.local_rank,
            find_unused_parameters=True
        )

    if base_model_name == "vinid/plip":
        patch_size = 224
    elif base_model_name == "openai/clip-vit-large-patch14-336":
        patch_size = 336
    else:
        raise ValueError(f"Model {base_model_name} not supported.")

    processor = model.module.patch_processor if config.distributed else model.patch_processor

    start_epoch = 0
    start_batch_idx = 0
    num_images_so_far = 0
    shuffle_seed = None

    if config.load_checkpoint:
        print(f"Loading checkpoint from {config.checkpoint_path}")
        checkpoint = torch.load(config.checkpoint_path, map_location=device)
        if config.distributed:
            model.module.load_state_dict(checkpoint['model_state_dict'])
        else:
            model.load_state_dict(checkpoint['model_state_dict'])
        optimizer.load_state_dict(checkpoint['optimizer_state_dict'])
        scheduler.load_state_dict(checkpoint['scheduler_state_dict'])
        
        # Extract epoch, batch_idx, and shuffle_seed from checkpoint
        start_epoch = checkpoint['epoch']
        start_batch_idx = checkpoint['batch_idx'] + 1  # Start from next batch
        shuffle_seed = checkpoint['shuffle_seed']
        print(f"Resuming training from epoch {start_epoch}, batch {start_batch_idx}, with shuffle_seed {shuffle_seed}")

    else:
        shuffle_seed = random.randint(0, 2**32 - 1)

    # Modify the dataset creation

    # # use 2000 validation for training validation
    # validation_set = hpa10m_index[hpa10m_index['split'] == 'validation']
    # hpa10m_test = validation_set.sample(n=2000, random_state=42)
    # remaining_validation = validation_set.drop(hpa10m_test.index).reset_index(drop=True)

    # # add the rest of val to the training set
    # hpa10m_train = pd.concat([hpa10m_index[hpa10m_index['split'] == 'train'], remaining_validation]).reset_index(drop=True)
    # hpa10m_test = hpa10m_test.reset_index(drop=True)

    # data locations: environment variables, see README
    hpa10m_train = pd.read_feather(DATA_TRAIN_META)
    hpa10m_test = pd.read_feather(DATA_TEST_META)
    overlap = pd.merge(hpa10m_train, hpa10m_test, on="name")
    if rank == 0:
        print(f"Overlap: {len(overlap)} (suppose to be 0)")
    # hpa10m_train = hpa10m_train[~hpa10m_train['url'].isin(overlap['url'])]

    if config.use_cell_type != "All":
        hpa10m_train = hpa10m_train.loc[hpa10m_train["cell_type"] == config.use_cell_type, ].reset_index(drop=True)
        # split 95% into train, 5% into validation
        hpa10m_train, hpa10m_test = train_test_split(hpa10m_train, test_size=0.01, random_state=42)
        hpa10m_train = hpa10m_train.reset_index(drop=True)
        hpa10m_test = hpa10m_test.reset_index(drop=True)
        if rank == 0:
            print(f"Using {config.use_cell_type} only. Which subset to {len(hpa10m_train)} images.")
    # unique_tissue_types = [val.lower().replace("cancer","").replace("tissue","").strip() for val in hpa10m_train["tissue"].unique()]
    
    # `datadir` is only read by the `simple_downsample` model version
    datadir = DATA_IMAGE_DIR
    hdf5_base_dir = DATA_RLE_DIR
    rle_map_path = DATA_RLE_INDEX

    if config.model_version == "simple_downsample":
        dataset = HPADatasetDownsample(hpa10m_train, datadir, data_split="train", target_size=patch_size, processor=processor)
    else:
        # dataset = HPADatasetMIL(hpa10m_train, datadir, data_split="train", patch_size=patch_size, processor=processor)
        dataset = HPADatasetMIL_url(hpa10m_train, rle_map_path, hdf5_base_dir, data_split="train", patch_size=patch_size, processor=processor)

    # Calculate the number of batches per GPU
    num_gpus = torch.distributed.get_world_size() if config.distributed else 1
    num_batches_per_gpu = len(dataset) // (config.batch_size * num_gpus)

    # Adjust start_batch_idx if it is too large
    if start_batch_idx >= num_batches_per_gpu:
        print(f"Warning: start_batch_idx ({start_batch_idx}) is larger than available batches per GPU ({num_batches_per_gpu}). Resetting to 0.")
        start_batch_idx = 0

    # Create test dataset and dataloader
    if config.model_version == "simple_downsample":
        test_dataset = HPADatasetDownsample(hpa10m_test, datadir, data_split="test", target_size=patch_size, processor=processor)
    else:
        test_dataset = HPADatasetMIL_url(hpa10m_test, rle_map_path, hdf5_base_dir, data_split="test", patch_size=patch_size, processor=processor)
    
    if config.distributed:
        test_sampler = torch.utils.data.distributed.DistributedSampler(test_dataset, shuffle=False)
        test_shuffle = False
    else:
        test_sampler = None
        test_shuffle = False

    test_dataloader = DataLoader(
        test_dataset,
        batch_size=config.batch_size,
        shuffle=test_shuffle,
        num_workers=config.num_workers,
        collate_fn=custom_collate_fn,
        pin_memory=False,
        prefetch_factor=8,
        persistent_workers=True,
        sampler=test_sampler
    )

    # Training parameters
    num_epochs = config.num_epochs
    learning_rate = config.learning_rate  # lr = 1e-3 does not reduce the loss.
    optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)

    for epoch in range(start_epoch, num_epochs):
        epoch_loss = 0
        epoch_sched_loss = 0.0

        if epoch == start_epoch and config.load_checkpoint:
            current_seed = shuffle_seed
        else:
            current_seed = random.randint(0, 2**32 - 1)
            shuffle_seed = current_seed  # Update the shuffle_seed

        # Set seeds for reproducibility
        set_seed(current_seed + epoch)  # Adding epoch to seed for variety
    
        if config.distributed:
            train_sampler = ResumableDistributedSampler(
                dataset,
                num_replicas=num_gpus,
                rank=rank,
                shuffle=config.shuffle_training,
                seed=current_seed
            )
            train_sampler.set_epoch(epoch)
            if epoch == start_epoch and start_batch_idx > 0:
                skip_samples = start_batch_idx * config.batch_size
                train_sampler.set_start_index(skip_samples)
        else:
            train_sampler = SeededSampler(dataset, shuffle=config.shuffle_training, seed=current_seed)
            if epoch == start_epoch and start_batch_idx > 0:
                skip_samples = start_batch_idx * config.batch_size
                train_sampler.set_start_index(skip_samples)

        # Create DataLoader
        dataloader = DataLoader(
            dataset,
            batch_size=config.batch_size,
            sampler=train_sampler,
            num_workers=config.num_workers,
            collate_fn=custom_collate_fn,
            pin_memory=False,
            prefetch_factor=8,
            persistent_workers=True
        )
        num_batches = len(dataloader)
        if epoch == start_epoch:
            num_images_so_far = (epoch * num_batches + start_batch_idx) * config.batch_size
        else:
            num_images_so_far = epoch * num_batches * config.batch_size


        for batch_idx, batch in enumerate(tqdm(dataloader, disable=(rank != 0))):
            # Skip batches if resuming from a specific batch index
            # if epoch == start_epoch and batch_idx < start_batch_idx:
            #     continue

            num_images_so_far += config.batch_size
            if batch is None:
                if rank == 0:
                    print("Batch is None - Maybe invalid image.")
                continue
            patches, metadata, query_input, caption_output, snomed_text, snomed_code, image_url, staining_intensity, staining_location, staining_quantity, malignancy, tissue_one_hot, cell_type_one_hot = batch
            # Patches is a tuple of K patches, each patch shape: (N, 3, 224, 224)
            patches = [patch.to(device) for patch in patches]
            # Move tensors to device
            staining_intensity = staining_intensity.to(device)
            staining_location = staining_location.to(device)
            staining_quantity = staining_quantity.to(device)
            tissue_one_hot = tissue_one_hot.to(device)
            malignancy = malignancy.to(device)
            
            # Backward pass and optimization
            optimizer.zero_grad()

            st = time.time()

            # Mixed precision forward pass
            with torch.cuda.amp.autocast():
                # Forward pass
                intensity_out, location_out, quantity_out, tissue_out, malignancy_out, A_raw = model(
                    patches, query_input, cell_type_one_hot, phase="train"
                )
                
                # Compute losses
                intensity_loss = criterion_intensity(intensity_out, torch.argmax(staining_intensity, dim=1))
                location_loss = criterion_location(location_out, torch.argmax(staining_location, dim=1))
                quantity_loss = criterion_quantity(quantity_out, torch.argmax(staining_quantity, dim=1))
                tissue_loss = criterion_tissue(tissue_out, torch.argmax(tissue_one_hot, dim=1))
                malignancy_loss = criterion_malignancy(malignancy_out, torch.argmax(malignancy, dim=1))

                # Combine losses
                total_loss = intensity_loss + location_loss + quantity_loss + tissue_loss + malignancy_loss

            # Backward pass and optimization with scaling
            if config.amp:
                scaler.scale(total_loss).backward()
                scaler.step(optimizer)
                scaler.update()
            else:
                total_loss.backward()
                optimizer.step()

            # LR scheduler: step on the loss averaged across ranks so all ranks step together;
            # per batch by default, per epoch with SCHEDULER_PER_EPOCH=1.
            _sched_loss = total_loss.detach()
            if config.distributed:
                torch.distributed.all_reduce(_sched_loss, op=torch.distributed.ReduceOp.SUM)
                _sched_loss = _sched_loss / torch.distributed.get_world_size()
            if not SCHEDULER_PER_EPOCH:
                scheduler.step(_sched_loss.item())
            epoch_sched_loss += _sched_loss.item()

            epoch_loss += total_loss.item()

            et = time.time()
            training_time = et - st

            # Logging and printing only in the main process
            if batch_idx % 100 == 0:
                if rank == 0:
                    # print("="*100)
                    # Print last evaluation performances
                    # tqdm.write("Last Evaluation Performances:")
                    # tqdm.write(f"  Test Loss: {wandb.run.summary.get('test_loss', -1):.4f}")
                    # tqdm.write(f"  Test Intensity Accuracy: {wandb.run.summary.get('test_intensity_accuracy', -1):.4f}")
                    # tqdm.write(f"  Test Location Accuracy: {wandb.run.summary.get('test_location_accuracy', -1):.4f}")
                    # tqdm.write(f"  Test Quantity Accuracy: {wandb.run.summary.get('test_quantity_accuracy', -1):.4f}")
                    # tqdm.write(f"  Test Tissue Accuracy: {wandb.run.summary.get('test_tissue_accuracy', -1):.4f}")
                    # tqdm.write(f"  Test Malignancy Accuracy: {wandb.run.summary.get('test_malignancy_accuracy', -1):.4f}")
                    # tqdm.write("")

                    # # Update tqdm with losses
                    # tqdm.write(f"Epoch {epoch+1}/{num_epochs}, Batch {batch_idx+start_batch_idx+1}/{len(dataloader)}", end="  ")
                    # tqdm.write(f"  Intensity Loss: {intensity_loss.item():.4f}", end="  ")
                    # tqdm.write(f"  Location Loss: {location_loss.item():.4f}", end="  ")
                    # tqdm.write(f"  Quantity Loss: {quantity_loss.item():.4f}", end="  ")
                    # tqdm.write(f"  Tissue Loss: {tissue_loss.item():.4f}", end="  ")
                    # tqdm.write(f"  Malignancy Loss: {malignancy_loss.item():.4f}", end="  ")
                    # tqdm.write(f"  Total Loss: {total_loss.item():.4f}")

                    # print(f"Number of images learned: {len(patches)}, Training: {training_time:.2f} seconds")

                    # Log losses per iteration
                    wandb.log({
                        "epoch": epoch + 1,
                        "batch": batch_idx + start_batch_idx + 1,
                        "intensity_loss": intensity_loss.item(),
                        "location_loss": location_loss.item(),
                        "quantity_loss": quantity_loss.item(),
                        "tissue_loss": tissue_loss.item(),
                        "malignancy_loss": malignancy_loss.item(),
                        "total_loss": total_loss.item(),
                        "training_time": training_time,
                        "num_images_so_far": num_images_so_far,
                        "test_loss": wandb.run.summary.get('test_loss', -1),
                        "test_intensity_accuracy": wandb.run.summary.get('test_intensity_accuracy', -1),
                        "test_location_accuracy": wandb.run.summary.get('test_location_accuracy', -1),
                        "test_quantity_accuracy": wandb.run.summary.get('test_quantity_accuracy', -1),
                        "test_tissue_accuracy": wandb.run.summary.get('test_tissue_accuracy', -1),
                        "test_malignancy_accuracy": wandb.run.summary.get('test_malignancy_accuracy', -1),
                        "lr": scheduler.get_last_lr()[0]
                    })

            # Evaluation step
            if config.batch_size == 16:
                eval_per_step = 33
            elif config.batch_size == 24:
                eval_per_step = 22
            elif config.batch_size == 2:
                eval_per_step = 264
            elif config.batch_size == 1:
                # eval_per_step = 5280
                eval_per_step = 10000
            else:
                eval_per_step = 50  # Default value

            if batch_idx % eval_per_step == 0:
                if rank == 0:
                    # Evaluate on test set after each epoch
                    os.makedirs(os.path.join(config.checkpoint_dir, "validation_output"), exist_ok=True)
                    save_name = os.path.join(config.checkpoint_dir, "validation_output", f"test_results_epoch_{epoch}_step_{batch_idx + start_batch_idx}.csv")
                    test_loss, test_accuracies = evaluate(model, test_dataloader, device, save_name)
                    print(f"Epoch {epoch+1}/{num_epochs} - Test Loss: {test_loss:.4f}")
                    for k, v in test_accuracies.items():
                        print(f"Test {k.capitalize()} Accuracy: {v:.4f}")

                    # Log test metrics
                    wandb.log({
                        "epoch": epoch + 1,
                        "batch": batch_idx + start_batch_idx + 1,
                        "num_images_so_far": num_images_so_far,
                        "test_loss": test_loss,
                        **{f"test_{k}_accuracy": v for k, v in test_accuracies.items()}
                    })

                    # Save checkpoint
                    checkpoint_path = os.path.join(config.checkpoint_dir, f"checkpoint_epoch_{epoch}_step_{batch_idx + start_batch_idx}.pth")
                    torch.save({
                        'epoch': epoch,
                        'batch_idx': batch_idx + start_batch_idx,
                        'shuffle_seed': shuffle_seed,
                        'model_state_dict': model.module.state_dict() if config.distributed else model.state_dict(),
                        'optimizer_state_dict': optimizer.state_dict(),
                        'scheduler_state_dict': scheduler.state_dict(),
                        'test_loss': test_loss,
                        'test_accuracies': test_accuracies,
                        'num_images_so_far': num_images_so_far
                    }, checkpoint_path)
                    print(f"Saved checkpoint to {checkpoint_path}")

        # Log epoch loss
        if rank == 0:
            wandb.log({
                "epoch": epoch + 1,
                "epoch_loss": epoch_loss / len(dataloader)
            })

            # Save checkpoint at the end of each epoch
            checkpoint_path = os.path.join(config.checkpoint_dir, f"checkpoint_epoch_{epoch}.pth")
            torch.save({
                'epoch': epoch + 1,
                'batch_idx': -1,  # Indicate end of epoch
                'shuffle_seed': shuffle_seed,
                'model_state_dict': model.module.state_dict() if config.distributed else model.state_dict(),
                'optimizer_state_dict': optimizer.state_dict(),
                'scheduler_state_dict': scheduler.state_dict(),
                'test_loss': test_loss,
                'test_accuracies': test_accuracies,
                'num_images_so_far': num_images_so_far
            }, checkpoint_path)
            print(f"Saved checkpoint to {checkpoint_path}")
            print(f"Epoch {epoch+1}/{num_epochs} completed")

    if rank == 0:
        print("Training completed!")
        wandb.finish()


def evaluate(model, dataloader, device, save_name):
    if config.distributed and torch.distributed.get_rank() != 0:
        return
    model.eval()
    total_loss = 0
    correct_predictions = {
        'intensity': 0,
        'location': 0,
        'quantity': 0,
        'tissue': 0,
        'malignancy': 0
    }
    total_samples = 0

    # Lists to store all outputs and ground truth
    all_outputs = {
        'intensity': [],
        'location': [],
        'quantity': [],
        'tissue': [],
        'malignancy': []
    }
    all_ground_truth = {
        'intensity': [],
        'location': [],
        'quantity': [],
        'tissue': [],
        'malignancy': []
    }

    all_predicted_classes = {
        'intensity': [],
        'location': [],
        'quantity': [],
        'tissue': [],
        'malignancy': []
    }

    all_image_urls = []

    with torch.no_grad():
        for batch in tqdm(dataloader, desc="Evaluating"):
            if batch is None:
                continue
            patches, metadata, query_input, caption_output, snomed_text, snomed_code, image_url, staining_intensity, staining_location, staining_quantity, malignancy, tissue_one_hot, cell_type_one_hot = batch
            
            patches = [patch.to(device) for patch in patches]
            staining_intensity = staining_intensity.to(device)
            staining_location = staining_location.to(device)
            staining_quantity = staining_quantity.to(device)
            tissue_one_hot = tissue_one_hot.to(device)
            malignancy = malignancy.to(device)

            # Mixed precision inference
            with torch.cuda.amp.autocast():
                intensity_out, location_out, quantity_out, tissue_out, malignancy_out, A_raw = model(patches, query_input, cell_type_one_hot, phase="test")

                # Compute losses
                intensity_loss = criterion_intensity(intensity_out, torch.argmax(staining_intensity, dim=1))
                location_loss = criterion_location(location_out, torch.argmax(staining_location, dim=1))
                quantity_loss = criterion_quantity(quantity_out, torch.argmax(staining_quantity, dim=1))
                tissue_loss = criterion_tissue(tissue_out, torch.argmax(tissue_one_hot, dim=1))
                malignancy_loss = criterion_malignancy(malignancy_out, torch.argmax(malignancy, dim=1))
                

                total_loss += (intensity_loss + location_loss + quantity_loss + tissue_loss + malignancy_loss).item()

            # Calculate correct predictions
            correct_predictions['intensity'] += (torch.argmax(intensity_out, dim=1) == torch.argmax(staining_intensity, dim=1)).sum().item()
            correct_predictions['location'] += (torch.argmax(location_out, dim=1) == torch.argmax(staining_location, dim=1)).sum().item()
            correct_predictions['quantity'] += (torch.argmax(quantity_out, dim=1) == torch.argmax(staining_quantity, dim=1)).sum().item()
            correct_predictions['tissue'] += (torch.argmax(tissue_out, dim=1) == torch.argmax(tissue_one_hot, dim=1)).sum().item()
            correct_predictions['malignancy'] += (torch.argmax(malignancy_out, dim=1) == torch.argmax(malignancy, dim=1)).sum().item()

            total_samples += staining_intensity.size(0)

            # Store outputs and ground truth
            all_outputs['intensity'].extend(intensity_out.cpu().numpy())
            all_outputs['location'].extend(location_out.cpu().numpy())
            all_outputs['quantity'].extend(quantity_out.cpu().numpy())
            all_outputs['tissue'].extend(tissue_out.cpu().numpy())
            all_outputs['malignancy'].extend(malignancy_out.cpu().numpy())

            all_predicted_classes['intensity'].extend(torch.argmax(intensity_out, dim=1).cpu().numpy())
            all_predicted_classes['location'].extend(torch.argmax(location_out, dim=1).cpu().numpy())
            all_predicted_classes['quantity'].extend(torch.argmax(quantity_out, dim=1).cpu().numpy())
            all_predicted_classes['tissue'].extend(torch.argmax(tissue_out, dim=1).cpu().numpy())
            all_predicted_classes['malignancy'].extend(torch.argmax(malignancy_out, dim=1).cpu().numpy())

            all_ground_truth['intensity'].extend(staining_intensity.cpu().numpy())
            all_ground_truth['location'].extend(staining_location.cpu().numpy())
            all_ground_truth['quantity'].extend(staining_quantity.cpu().numpy())
            all_ground_truth['tissue'].extend(tissue_one_hot.cpu().numpy())
            all_ground_truth['malignancy'].extend(malignancy.cpu().numpy())

            all_image_urls.extend(image_url)

    avg_loss = total_loss / len(dataloader)
    accuracies = {k: v / total_samples for k, v in correct_predictions.items()}

    # Save all outputs and ground truth to a CSV file
    results_df = pd.DataFrame({
        'image_url': all_image_urls,
        'intensity_out': all_outputs['intensity'],
        'intensity_true': all_ground_truth['intensity'],
        'location_out': all_outputs['location'],
        'location_true': all_ground_truth['location'],
        'quantity_out': all_outputs['quantity'],
        'quantity_true': all_ground_truth['quantity'],
        'tissue_out': all_outputs['tissue'],
        'tissue_true': all_ground_truth['tissue'],
        'malignancy_out': all_outputs['malignancy'],
        'malignancy_true': all_ground_truth['malignancy']
    })
    # Save the results
    results_df.to_csv(save_name, index=False)
    print(f"Test results saved to {save_name}")

    # Plot confusion matrices
    def plot_confusion_matrix(y_true, y_pred, title, ax):
        cm = confusion_matrix(y_true, y_pred)
        sns.heatmap(cm, annot=True, fmt='d', cmap='Blues', ax=ax)
        ax.set_title(title)
        ax.set_ylabel('True label')
        ax.set_xlabel('Predicted label')

        # Get the class names from hpadataset.py
        if 'intensity' in title.lower():
            class_names = ['negative', 'weak', 'moderate', 'strong']
        elif 'location' in title.lower():
            class_names = ['none', 'cytoplasmic/membranous', 'nuclear', 'cytoplasmic/membranous,nuclear']
        elif 'quantity' in title.lower():
            class_names = ['none', '<25%', '25%-75%', '>75%']
        elif 'malignancy' in title.lower():
            class_names = ['normal', 'cancer']
        elif 'tissue' in title.lower():
            class_names = ['adipose', 'adrenal gland', 'appendix', 'bone marrow', 'breast',
                           'bronchus', 'carcinoid', 'caudate', 'cerebellum',
                           'cerebral cortex', 'cervical', 'cervix', 'colon', 'colorectal',
                           'duodenum', 'endometrial', 'endometrium', 'epididymis',
                           'esophagus', 'fallopian tube', 'gallbladder', 'glioma',
                           'head and neck', 'heart muscle', 'hippocampus', 'kidney', 'liver',
                           'lung', 'lymph node', 'lymphoma', 'melanoma', 'nasopharynx',
                           'oral mucosa', 'ovarian', 'ovary', 'pancreas', 'pancreatic',
                           'parathyroid gland', 'placenta', 'prostate', 'rectum', 'renal',
                           'salivary gland', 'seminal vesicle', 'skeletal muscle', 'skin',
                           'small intestine', 'smooth muscle', 'soft', 'spleen', 'stomach',
                           'testis', 'thyroid', 'thyroid gland', 'tonsil', 'urinary bladder',
                           'urothelial', 'vagina']
        else:
            class_names = [str(i) for i in range(cm.shape[0])]

        ax.set_xticks(np.arange(len(class_names)) + 0.5)
        ax.set_yticks(np.arange(len(class_names)) + 0.5)
        ax.set_xticklabels(class_names, rotation=90, ha='right')
        ax.set_yticklabels(class_names, rotation=0)

    # Plot for staining intensity, location, quantity, and malignancy
    fig, axes = plt.subplots(2, 2, figsize=(20, 20))
    tasks = ['intensity', 'location', 'quantity', 'malignancy']
    for i, task in enumerate(tasks):
        y_true = np.argmax(all_ground_truth[task], axis=1)
        y_pred = np.argmax(all_outputs[task], axis=1)
        plot_confusion_matrix(y_true, y_pred, f'Confusion Matrix - {task.capitalize()}', axes[i//2, i%2])

    plt.tight_layout()
    cm_path = save_name.replace('.csv', '_confusion_matrices.png')
    plt.savefig(cm_path)
    plt.close()

    # Plot for tissue type
    plt.figure(figsize=(30, 30))  # Increased figure size for better readability
    y_true = np.argmax(all_ground_truth['tissue'], axis=1)
    y_pred = np.argmax(all_outputs['tissue'], axis=1)
    plot_confusion_matrix(y_true, y_pred, 'Confusion Matrix - Tissue Type', plt.gca())
    
    tissue_cm_path = save_name.replace('.csv', '_tissue_confusion_matrix.png')
    plt.savefig(tissue_cm_path)
    plt.close()

    return avg_loss, accuracies


def parse_config():
    parser = argparse.ArgumentParser(description='Train HPA-VLM model')
    # parser.add_argument('--config', type=str, default='training_config/config_9.ini', help='Path to the configuration file')
    parser.add_argument('--config', type=str, default='config/config.ini', help='Path to the configuration file')
    # Add arguments for distributed training
    parser.add_argument('--distributed', action='store_true', help='Use distributed training')
    parser.add_argument('--local_rank', type=int, default=int(os.environ.get('LOCAL_RANK', 0)), help='Local rank for distributed training')
    args = parser.parse_args()

    config = configparser.ConfigParser()
    config.read(args.config)
    
    # Convert config values to desired data types without defaults
    def convert_value(convert_func, key):
        try:
            return convert_func(config['DEFAULT'][key])
        except (ValueError, TypeError) as e:
            raise ValueError(f"Error converting {key}: {str(e)}. Please check your config file.")

    # Create a dictionary to hold the converted values
    converted_config = {
        'base_model_name': config['DEFAULT']['base_model_name'],
        'shuffle_training': convert_value(lambda x: x.lower() == 'true', 'shuffle_training'),
        'amp': convert_value(lambda x: x.lower() == 'true', 'amp'),
        'freeze_vit': convert_value(lambda x: x.lower() == 'true', 'freeze_vit'),
        'freeze_query_features_encoder': convert_value(lambda x: x.lower() == 'true', 'freeze_query_features_encoder'),
        'batch_size': convert_value(int, 'batch_size'),
        'num_workers': convert_value(int, 'num_workers'),
        'num_epochs': convert_value(int, 'num_epochs'),
        'learning_rate': convert_value(float, 'learning_rate'),
        'weight_decay': convert_value(float, 'weight_decay'),
        'use_cell_type_embedding': convert_value(lambda x: x.lower() == 'true', 'use_cell_type_embedding'),
        'use_cell_type': config['DEFAULT']['use_cell_type'],
        'model_version': config['DEFAULT']['model_version'],
        'checkpoint_dir': config['DEFAULT'].get('checkpoint_dir', 'checkpoints'),
        'load_checkpoint': eval(config['DEFAULT'].get('load_checkpoint')),
        'checkpoint_path': config['DEFAULT'].get('checkpoint_path', ''),
        'distributed': args.distributed,
        'local_rank': args.local_rank,
    }
    
    # Create a run name using base_model_name and learning_rate
    run_name = f"{converted_config['base_model_name']}_lr={converted_config['learning_rate']}"
    if converted_config['freeze_vit']:
        run_name += "_freeze_vit"
    if converted_config['freeze_query_features_encoder']:
        run_name += "_freeze_query_encoder"
    if converted_config['use_cell_type_embedding']:
        run_name += "use_cell_type_embedding"
    if converted_config['load_checkpoint']:
        run_name += "_resumed"
    run_name += f"_rank{converted_config['local_rank']}"

    # Initialize wandb with config and run name
    if not converted_config['distributed'] or converted_config['local_rank'] == 0:
        wandb.init(
            project=os.environ.get("WANDB_PROJECT", "isight-slide"),
            config=converted_config,
            name=run_name
        )

    return argparse.Namespace(**converted_config)

if __name__ == "__main__":
    config = parse_config()
    set_seed(1001)
    if config.distributed:
        torch.distributed.init_process_group(backend='nccl')
        torch.cuda.set_device(config.local_rank)
    # **Modified code to reflect resumed training in checkpoint_dir**
    if config.load_checkpoint:
        # Use the same run directory as the checkpoint
        run_dir = os.path.dirname(config.checkpoint_path)
        print(f"Resuming from run directory: {run_dir}")
    else:
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        run_dir = os.path.join(config.checkpoint_dir, f"run_{timestamp}")
    os.makedirs(run_dir, exist_ok=True)
    os.makedirs(os.path.join(run_dir, "wandb"), exist_ok=True)
    # Set wandb directory
    os.environ["WANDB_DIR"] = os.path.join(run_dir, "wandb")
    # Update config with new checkpoint directory
    config.checkpoint_dir = run_dir
    print("-------------------------------------------------------------")
    print(config.checkpoint_dir)
    print("-------------------------------------------------------------")

    train(config)