# Install the ultralytics package for YOLOv8 (if not already installed)
# !pip install ultralytics

# Import necessary libraries
import pandas as pd
import numpy as np
import torch
from torch.utils.data import Dataset, DataLoader
from tqdm.notebook import tqdm
from torchvision import transforms
from PIL import Image
import torch.nn as nn

# Suppress warnings from YOLO and transformers
import warnings
warnings.filterwarnings("ignore")

# Mount Google Drive if needed
from google.colab import drive
drive.mount('/content/drive')

# Load the CSV file
csv_file = '/content/scenario32_dev.csv'  # Replace with your actual path
df = pd.read_csv(csv_file)

# Import YOLOv8 model
from ultralytics import YOLO

# Suppress YOLO logging output
import logging
logging.getLogger('ultralytics').setLevel(logging.ERROR)

# Load the pre-trained YOLOv8 segmentation model
yolo_model = YOLO('yolov8l-seg.pt')  # Using 'yolov8l-seg.pt' for better segmentation

# Create a custom dataset class
class BeamDataset(Dataset):
    def __init__(self, dataframe, transform=None, label_mapping=None, augment=False):
        self.data = dataframe.reset_index(drop=True)
        self.image_paths = self.data['unit1_rgb']
        self.labels = self.data['unit1_beam']
        self.transform = transform
        self.label_mapping = label_mapping
        self.augment = augment

        # Map labels to indices
        self.labels = self.labels.map(self.label_mapping).astype(int)

        if self.augment and self.transform is not None:
            self.transform = transforms.Compose([
                transforms.RandomHorizontalFlip(),
                transforms.RandomRotation(15),
                transforms.ColorJitter(brightness=0.2, contrast=0.2, saturation=0.2, hue=0.2),
                self.transform
            ])

    def __len__(self):
        return len(self.data)

    def __getitem__(self, idx):
        img_path = self.image_paths.iloc[idx]

        # Load image
        image = Image.open(img_path).convert('RGB')

        # Perform YOLOv8 segmentation without verbose output
        results = yolo_model(img_path, verbose=False)
        # Get the segmentation mask
        masks = results[0].masks.data.cpu().numpy()  # Shape: (num_objects, H, W)

        # If no objects detected, create an empty mask
        if len(masks) == 0:
            seg_mask = np.zeros((image.height, image.width), dtype=np.uint8)
        else:
            # Combine all masks into a single mask
            seg_mask = np.sum(masks, axis=0)
            seg_mask = np.clip(seg_mask, 0, 1).astype(np.uint8)
            # Resize to match image size if necessary
            seg_mask = Image.fromarray(seg_mask * 255).resize(image.size)

        # Apply transforms
        if self.transform:
            image = self.transform(image)
            seg_mask = self.transform(seg_mask.convert('RGB'))
        else:
            transform = transforms.ToTensor()
            image = transform(image)
            seg_mask = transform(seg_mask.convert('RGB'))

        label = self.labels.iloc[idx]
        return {'image': image, 'seg_mask': seg_mask, 'label': label}

# Define the transforms
from transformers import ViTImageProcessor

# Use a smaller model for practical GPU memory usage
model_name = 'google/vit-large-patch16-224-in21k'
image_processor = ViTImageProcessor.from_pretrained(model_name)

common_transforms = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=image_processor.image_mean, std=image_processor.image_std)
])

# Create label mapping
all_labels = df['unit1_beam'].unique()
label_mapping = {label: idx for idx, label in enumerate(sorted(all_labels))}
num_labels = len(label_mapping)

# Handle class imbalance by computing class weights
from sklearn.utils.class_weight import compute_class_weight

class_weights = compute_class_weight(class_weight='balanced', classes=np.unique(df['unit1_beam']), y=df['unit1_beam'])
class_weights = torch.tensor(class_weights, dtype=torch.float).to('cuda' if torch.cuda.is_available() else 'cpu')

# Split the dataset into training and testing sets
from sklearn.model_selection import train_test_split

train_df, test_df = train_test_split(df, test_size=0.2, random_state=42, stratify=df['unit1_beam'])

train_dataset = BeamDataset(train_df, transform=common_transforms, label_mapping=label_mapping, augment=True)
test_dataset = BeamDataset(test_df, transform=common_transforms, label_mapping=label_mapping, augment=False)

train_loader = DataLoader(train_dataset, batch_size=16, shuffle=True, num_workers=2, pin_memory=True)
test_loader = DataLoader(test_dataset, batch_size=16, shuffle=False, num_workers=2, pin_memory=True)

# Load the ViT model
from transformers import ViTModel

# Load pre-trained ViT models for images and segmentation masks
vit_image = ViTModel.from_pretrained(model_name)
vit_segmentation = ViTModel.from_pretrained(model_name)

# Fine-tune the last few layers of ViT models
for name, param in vit_image.named_parameters():
    if 'encoder.layer.11' not in name and 'encoder.layer.10' not in name:
        param.requires_grad = False

for name, param in vit_segmentation.named_parameters():
    if 'encoder.layer.11' not in name and 'encoder.layer.10' not in name:
        param.requires_grad = False

# Define the combined model
class BeamPredictionModel(nn.Module):
    def __init__(self, num_labels, vit_image, vit_segmentation):
        super(BeamPredictionModel, self).__init__()
        self.vit_image = vit_image
        self.vit_segmentation = vit_segmentation
        self.dropout = nn.Dropout(0.3)
        self.classifier = nn.Linear(vit_image.config.hidden_size * 2, num_labels)

    def forward(self, image, seg_mask):
        # Pass through ViT models
        image_outputs = self.vit_image(pixel_values=image).pooler_output
        seg_outputs = self.vit_segmentation(pixel_values=seg_mask).pooler_output

        # Concatenate embeddings
        combined_representation = torch.cat(
            (image_outputs, seg_outputs), dim=1
        )

        # Classification
        x = self.dropout(combined_representation)
        logits = self.classifier(x)
        return logits

# Initialize the model
model = BeamPredictionModel(num_labels, vit_image, vit_segmentation)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
model.to(device)

# Define loss function with class weights
criterion = nn.CrossEntropyLoss(weight=class_weights)

# Define optimizer with weight decay
optimizer = torch.optim.AdamW(model.parameters(), lr=2e-5, weight_decay=1e-4)

# Define learning rate scheduler
from torch.optim.lr_scheduler import ReduceLROnPlateau

scheduler = ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=3, verbose=True)

# Use mixed precision
scaler = torch.cuda.amp.GradScaler()

# Import matplotlib for plotting
import matplotlib.pyplot as plt

# Initialize lists to store metrics
train_losses = []
val_losses = []
val_top1_accuracies = []
val_top2_accuracies = []
val_top3_accuracies = []
val_dba_scores = []

# Early Stopping parameters
patience = 5
best_val_loss = float('inf')
epochs_no_improve = 0

# Training loop
num_epochs = 20  # Increased number of epochs

for epoch in range(num_epochs):
    model.train()
    total_loss = 0
    progress_bar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{num_epochs}")

    for batch in progress_bar:
        images = batch['image'].to(device, non_blocking=True)
        seg_masks = batch['seg_mask'].to(device, non_blocking=True)
        labels = batch['label'].to(device, non_blocking=True)

        optimizer.zero_grad()
        with torch.cuda.amp.autocast():
            logits = model(images, seg_masks)
            loss = criterion(logits, labels)

        scaler.scale(loss).backward()
        scaler.step(optimizer)
        scaler.update()

        total_loss += loss.item()
        progress_bar.set_postfix(loss=loss.item())

    avg_loss = total_loss / len(train_loader)
    train_losses.append(avg_loss)
    print(f"Epoch [{epoch+1}/{num_epochs}], Training Loss: {avg_loss:.4f}")

    # Evaluate on the test set
    model.eval()
    total_correct_top1 = 0
    total_correct_top2 = 0
    total_correct_top3 = 0
    total_samples = 0
    total_dba = 0
    total_val_loss = 0

    with torch.no_grad():
        for batch in test_loader:
            images = batch['image'].to(device, non_blocking=True)
            seg_masks = batch['seg_mask'].to(device, non_blocking=True)
            labels = batch['label'].to(device, non_blocking=True)

            logits = model(images, seg_masks)
            val_loss = criterion(logits, labels)
            total_val_loss += val_loss.item()

            # Get top-k predictions
            _, preds_topk = logits.topk(3, dim=1)
            correct_topk = preds_topk.eq(labels.view(-1, 1))  # Shape: (batch_size, 3)

            # Compute Top-1 accuracy
            total_correct_top1 += correct_topk[:, 0].sum().item()

            # Compute Top-2 accuracy
            total_correct_top2 += correct_topk[:, :2].any(dim=1).sum().item()

            # Compute Top-3 accuracy
            total_correct_top3 += correct_topk.any(dim=1).sum().item()

            # Compute DBA score
            probs = torch.softmax(logits, dim=1)
            batch_size = labels.size(0)
            for i in range(batch_size):
                prob = probs[i]
                label = labels[i]
                sorted_probs, indices = prob.sort(descending=True)
                rank = (indices == label).nonzero(as_tuple=True)[0].item() + 1  # Rank starts from 1
                total_dba += 1.0 / rank

            total_samples += labels.size(0)

    avg_val_loss = total_val_loss / len(test_loader)
    val_losses.append(avg_val_loss)

    top1_accuracy = total_correct_top1 / total_samples
    top2_accuracy = total_correct_top2 / total_samples
    top3_accuracy = total_correct_top3 / total_samples
    avg_dba = total_dba / total_samples

    val_top1_accuracies.append(top1_accuracy)
    val_top2_accuracies.append(top2_accuracy)
    val_top3_accuracies.append(top3_accuracy)
    val_dba_scores.append(avg_dba)

    print(f"Validation Loss: {avg_val_loss:.4f}")
    print(f"Validation Top-1 Accuracy: {top1_accuracy:.4f}")
    print(f"Validation Top-2 Accuracy: {top2_accuracy:.4f}")
    print(f"Validation Top-3 Accuracy: {top3_accuracy:.4f}")
    print(f"Validation DBA Score: {avg_dba:.4f}")

    # Step the scheduler
    scheduler.step(avg_val_loss)

    # Early Stopping Check
    if avg_val_loss < best_val_loss:
        best_val_loss = avg_val_loss
        epochs_no_improve = 0
        # Save the best model
        torch.save(model.state_dict(), 'best_beam_prediction_model.pth')
    else:
        epochs_no_improve += 1
        if epochs_no_improve >= patience:
            print(f"Early stopping triggered after {patience} epochs without improvement.")
            break

# Load the best model
model.load_state_dict(torch.load('best_beam_prediction_model.pth'))

# Plotting Training and Validation Loss over Epochs
plt.figure(figsize=(10, 6))
plt.plot(range(1, len(train_losses)+1), train_losses, label='Training Loss')
plt.plot(range(1, len(val_losses)+1), val_losses, label='Validation Loss')
plt.xlabel('Epoch')
plt.ylabel('Loss')
plt.title('Training and Validation Loss over Epochs')
plt.legend()
plt.grid(True)
plt.show()

# Plotting Top-1, Top-2, Top-3 Accuracy over Epochs
plt.figure(figsize=(10, 6))
plt.plot(range(1, len(val_top1_accuracies)+1), val_top1_accuracies, label='Top-1 Accuracy')
plt.plot(range(1, len(val_top2_accuracies)+1), val_top2_accuracies, label='Top-2 Accuracy')
plt.plot(range(1, len(val_top3_accuracies)+1), val_top3_accuracies, label='Top-3 Accuracy')
plt.xlabel('Epoch')
plt.ylabel('Accuracy')
plt.title('Top-1, Top-2, Top-3 Accuracy over Epochs')
plt.legend()
plt.grid(True)
plt.show()

# Plotting DBA Score over Epochs
plt.figure(figsize=(10, 6))
plt.plot(range(1, len(val_dba_scores)+1), val_dba_scores, label='DBA Score', marker='o')
plt.xlabel('Epoch')
plt.ylabel('DBA Score')
plt.title('DBA Score over Epochs')
plt.legend()
plt.grid(True)
plt.show()


the above code gives this error :

RuntimeError                              Traceback (most recent call last)
<ipython-input-4-cd53479c1665> in <cell line: 206>()
    209     progress_bar = tqdm(train_loader, desc=f"Epoch {epoch+1}/{num_epochs}")
    210 
--> 211     for batch in progress_bar:
    212         images = batch['image'].to(device, non_blocking=True)
    213         seg_masks = batch['seg_mask'].to(device, non_blocking=True)

5 frames
/usr/local/lib/python3.10/dist-packages/tqdm/notebook.py in __iter__(self)
    248         try:
    249             it = super().__iter__()
--> 250             for obj in it:
    251                 # return super(tqdm...) will not catch exception
    252                 yield obj

/usr/local/lib/python3.10/dist-packages/tqdm/std.py in __iter__(self)
   1179 
   1180         try:
-> 1181             for obj in iterable:
   1182                 yield obj
   1183                 # Update and possibly print the progressbar.

/usr/local/lib/python3.10/dist-packages/torch/utils/data/dataloader.py in __next__(self)
    699                 # TODO(https://github.com/pytorch/pytorch/issues/76750)
    700                 self._reset()  # type: ignore[call-arg]
--> 701             data = self._next_data()
    702             self._num_yielded += 1
    703             if (

/usr/local/lib/python3.10/dist-packages/torch/utils/data/dataloader.py in _next_data(self)
   1463             else:
   1464                 del self._task_info[idx]
-> 1465                 return self._process_data(data)
   1466 
   1467     def _try_put_index(self):

/usr/local/lib/python3.10/dist-packages/torch/utils/data/dataloader.py in _process_data(self, data)
   1489         self._try_put_index()
   1490         if isinstance(data, ExceptionWrapper):
-> 1491             data.reraise()
   1492         return data
   1493 

/usr/local/lib/python3.10/dist-packages/torch/_utils.py in reraise(self)
    713             # instantiate since we don't know how to
    714             raise RuntimeError(msg) from None
--> 715         raise exception
    716 
    717 

RuntimeError: Caught RuntimeError in DataLoader worker process 0.
Original Traceback (most recent call last):
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/_utils/worker.py", line 351, in _worker_loop
    data = fetcher.fetch(index)  # type: ignore[possibly-undefined]
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/_utils/fetch.py", line 52, in fetch
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "/usr/local/lib/python3.10/dist-packages/torch/utils/data/_utils/fetch.py", line 52, in <listcomp>
    data = [self.dataset[idx] for idx in possibly_batched_index]
  File "<ipython-input-4-cd53479c1665>", line 67, in __getitem__
    results = yolo_model(img_path, verbose=False)
  File "/usr/local/lib/python3.10/dist-packages/ultralytics/engine/model.py", line 176, in __call__
    return self.predict(source, stream, **kwargs)
  File "/usr/local/lib/python3.10/dist-packages/ultralytics/engine/model.py", line 547, in predict
    self.predictor.setup_model(model=self.model, verbose=is_cli)
  File "/usr/local/lib/python3.10/dist-packages/ultralytics/engine/predictor.py", line 306, in setup_model
    device=select_device(self.args.device, verbose=verbose),
  File "/usr/local/lib/python3.10/dist-packages/ultralytics/utils/torch_utils.py", line 218, in select_device
    s += f"{'' if i == 0 else space}CUDA:{d} ({get_gpu_info(i)})\n"  # bytes to MB
  File "/usr/local/lib/python3.10/dist-packages/ultralytics/utils/torch_utils.py", line 129, in get_gpu_info
    properties = torch.cuda.get_device_properties(index)
  File "/usr/local/lib/python3.10/dist-packages/torch/cuda/__init__.py", line 523, in get_device_properties
    _lazy_init()  # will define _get_device_properties
  File "/usr/local/lib/python3.10/dist-packages/torch/cuda/__init__.py", line 305, in _lazy_init
    raise RuntimeError(
RuntimeError: Cannot re-initialize CUDA in forked subprocess. To use CUDA with multiprocessing, you must use the 'spawn' start method