# 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 = '/kaggle/input/scenario32/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

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)

        # Initialize YOLOv8 model as None; it will be loaded per worker
        self.yolo_model = None

        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):
        # Initialize the YOLOv8 model on CPU in the worker process
        if self.yolo_model is None:
            self.yolo_model = YOLO('yolov8l-seg.pt')
            self.yolo_model.to('cpu')  # Ensure the model runs on CPU

        img_path = self.image_paths.iloc[idx]

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

        # Perform YOLOv8 segmentation on CPU without verbose output
        results = self.yolo_model(img_path, verbose=False, device='cpu')

        # 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()


convert this code to kaggle code