DEBUG:rasterio.env:Entered env context: <rasterio.env.Env object at 0x316fec6d0>
DEBUG:rasterio._io:Output nodata value read from file: -32768.0
DEBUG:rasterio._io:Output nodata values: [-32768.0]
DEBUG:rasterio._io:all_valid: False
DEBUG:rasterio._io:mask_flags: ([<MaskFlags.nodata: 8>],)
DEBUG:rasterio._io:Jump straight to _read()
DEBUG:rasterio._io:Window: Window(col_off=6901, row_off=9339, width=512, height=512)
DEBUG:rasterio._io:IO window xoff=6901.0 yoff=9339.0 width=512.0 height=512.0
DEBUG:rasterio.env:Exiting env context: <rasterio.env.Env object at 0x316fec6d0>
DEBUG:rasterio.env:Cleared existing <rasterio._env.GDALEnv object at 0x3183b72b0> options
DEBUG:rasterio._env:Stopped GDALEnv <rasterio._env.GDALEnv object at 0x3183b72b0>.
DEBUG:rasterio.env:Exiting outermost env
DEBUG:rasterio.env:Exited env context: <rasterio.env.Env object at 0x316fec6d0>
[DEBUG] Channel 'vh' image shape before transformations: (512, 512)
[DEBUG] Channel 'vh' image shape after individual_transforms: (512, 512)
[DEBUG] Channel 'vh' image shape after normalization: (512, 512)
[DEBUG] Channel 'vv' image shape before transformations: (512, 512)
[DEBUG] Channel 'vv' image shape after individual_transforms: (512, 512)
[DEBUG] Channel 'vv' image shape after normalization: (512, 512)
[DEBUG] Number of channels: 2
[DEBUG] Channel 0 shape: (512, 512)
[DEBUG] Channel 1 shape: (512, 512)
[DEBUG] Combined image shape after np.dstack: (512, 512, 2)
[DEBUG] Keypoints being passed to albumentations:
(257.0, 263.0, 0.0) [<class 'float'>, <class 'float'>, <class 'float'>]
Error executing job with overrides: ['+config=configs/multilabel_circle_net.yaml']
Traceback (most recent call last):
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 529, in <module>
    main()
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/main.py", line 94, in decorated_main
    _run_hydra(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 394, in _run_hydra
    _run_app(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 457, in _run_app
    run_and_report(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 223, in run_and_report
    raise ex
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 220, in run_and_report
    return func()
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/utils.py", line 458, in <lambda>
    lambda: hydra.run(
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/_internal/hydra.py", line 132, in run
    _ = ret.return_value
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/core/utils.py", line 260, in return_value
    raise self._return_value
  File "/opt/anaconda3/envs/sam3/lib/python3.9/site-packages/hydra/core/utils.py", line 186, in run_job
    ret.return_value = task_function(task_cfg)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 525, in main
    CenternetPipeline(cfg).train()
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 416, in train
    model = self.build_model(config, self.loaders)
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/train_multilabel.py", line 207, in build_model
    model = nas.search(config.nas.num_iterations, loaders['train'], loaders['valid'])
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/nas.py", line 35, in search
    scores = [self.evaluate_model(model, train_loader, val_loader) for model in population]
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/nas.py", line 35, in <listcomp>
    scores = [self.evaluate_model(model, train_loader, val_loader) for model in population]
  File "/Users/samerkharboush/Downloads/SAR-MPS-OptunaNasFbnYoloTransformers/xview3/nas.py", line 109, in evaluate_model
    raise ValueError("Batch does not contain 'image/images' or 'targets/labels' keys.")
ValueError: Batch does not contain 'image/images' or 'targets/labels' keys.
(sam3) samerkharboush@Samers-MBP SAR-MPS-OptunaNasFbnYoloTransformers % 


here's full nas.py file as I think it's not fully correct
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from skopt import gp_minimize
from skopt.space import Real, Integer, Categorical
import random
import math
import copy
from torch import Tensor
from torch.optim import Adam
from torch.utils.data import DataLoader
from xview3.common import AdaptiveVesselDetectionModule
from xview3.centernet.bboxer.multilabel_circle_coder import MultilabelCircleNetCoder

class ProgressiveNAS:
    def __init__(self, base_model, search_space):
        self.base_model = base_model
        self.search_space = search_space
        self.box_coder = MultilabelCircleNetCoder(
            image_size=(1024, 1024),
            output_stride=2,
            max_objects=512,
            heatmap_encoding="umich",
            labels_encoding="circle",
            ignore_value=-1
        )
        self.population_size = 20
        self.tournament_size = 5

    def search(self, num_iterations, train_loader, val_loader):
        population = [self.mutate_model(self.base_model) for _ in range(self.population_size)]
        
        for _ in range(num_iterations):
            scores = [self.evaluate_model(model, train_loader, val_loader) for model in population]
            
            new_population = []
            for _ in range(self.population_size):
                parent1 = self.tournament_selection(population, scores)
                parent2 = self.tournament_selection(population, scores)
                child = self.crossover(parent1, parent2)
                child = self.mutate_model(child)
                new_population.append(child)
            
            population = new_population

        best_model = max(population, key=lambda model: self.evaluate_model(model, train_loader, val_loader))
        return best_model

    def tournament_selection(self, population, scores):
        tournament = random.sample(list(zip(population, scores)), self.tournament_size)
        return max(tournament, key=lambda x: x[1])[0]

    def crossover(self, parent1, parent2):
        child = copy.deepcopy(parent1)
        for name, param in child.named_parameters():
            if random.random() < 0.5:
                param.data = parent2.state_dict()[name].clone()
        return child


    def mutate_model(self, model):
        new_model = copy.deepcopy(model)
        for param in new_model.parameters():
            if random.random() < 0.1:
                param.data += torch.randn_like(param.data) * 0.1
        
        # Structure mutation
        if random.random() < 0.2:
            new_layer = nn.Linear(64, 64)  # Example layer, adjust as needed
            if hasattr(new_model, 'encoder') and hasattr(new_model.encoder, 'add_module'):
                new_model.encoder.add_module(f"extra_layer_{id(new_layer)}", new_layer)
            elif hasattr(new_model, 'decoder') and hasattr(new_model.decoder, 'add_module'):
                new_model.decoder.add_module(f"extra_layer_{id(new_layer)}", new_layer)
        
        return new_model

    def evaluate_model(self, model, train_loader, val_loader):
        # Set the model to evaluation mode
        model.eval()
        total_loss = 0.0
        total_correct = 0
        total_samples = 0

        with torch.no_grad():
            for batch_idx, batch in enumerate(train_loader):
                print(f"[DEBUG] Batch index: {batch_idx}")
                print(f"[DEBUG] Type of batch: {type(batch)}")
                if isinstance(batch, dict):
                    print(f"[DEBUG] Batch keys: {list(batch.keys())}")
                elif isinstance(batch, (list, tuple)):
                    print(f"[DEBUG] Batch length: {len(batch)}")
                    for i, item in enumerate(batch):
                        print(f"[DEBUG] Batch[{i}] type: {type(item)}")
                        if isinstance(item, torch.Tensor):
                            print(f"[DEBUG] Batch[{i}] shape: {item.shape}")
                        else:
                            print(f"[DEBUG] Batch[{i}]: {item}")
                else:
                    print("[DEBUG] Batch is neither dict nor list/tuple")
                    print(f"[DEBUG] Batch content: {batch}")
                    continue  # Skip this batch or handle accordingly

                # Adjusting to batch structure
                if isinstance(batch, dict):
                    inputs = batch.get("image") or batch.get("images")
                    targets = batch.get("targets") or batch.get("labels")
                    if inputs is None or targets is None:
                        raise ValueError("Batch does not contain 'image/images' or 'targets/labels' keys.")
                elif isinstance(batch, (list, tuple)):
                    if len(batch) >= 2:
                        inputs, targets = batch[0], batch[1]
                    else:
                        raise ValueError("Batch tuple does not have at least two elements (inputs and targets).")
                else:
                    raise ValueError("Batch is in an unexpected format.")

                # Move data to device
                device = next(model.parameters()).device
                inputs = inputs.to(device)
                targets = targets.to(device)

                # Proceed with evaluation
                outputs = model(inputs)
                loss = self.loss_fn(outputs, targets)

                # Update metrics
                total_loss += loss.item() * inputs.size(0)
                # Assuming classification task
                _, predicted = outputs.max(1)
                total_correct += predicted.eq(targets).sum().item()
                total_samples += inputs.size(0)

        # Return evaluation score (e.g., accuracy)
        accuracy = total_correct / total_samples
        return accuracy

    def calculate_loss(self, outputs, targets):
        # Implement your loss calculation here
        # This is a placeholder implementation
        loss = 0
        for key in outputs:
            if key in targets:
                loss += F.mse_loss(outputs[key], targets[key])
        return loss


class NAS:
    def __init__(self, search_space, evaluation_function):
        self.search_space = search_space
        self.evaluation_function = evaluation_function

    def search(self, num_iterations, train_loader, val_loader):
        population = self.initialize_population()
        
        for _ in range(num_iterations):
            scores = [self.evaluate_model(model, train_loader, val_loader) for model in population]
            best_model = population[np.argmax(scores)]
            
            new_population = [best_model]
            while len(new_population) < len(population):
                parent1, parent2 = self.selection(population, scores)
                child = self.crossover(parent1, parent2)
                child = self.mutate(child)
                new_population.append(child)
            
            population = new_population
        
        return population[np.argmax([self.evaluate_model(model, train_loader, val_loader) for model in population])]


    def sample_architecture(self):
        architecture = {}
        for key, value_range in self.search_space.items():
            if isinstance(value_range, list):
                architecture[key] = random.choice(value_range)
            elif isinstance(value_range, tuple):
                architecture[key] = random.uniform(*value_range)
        return architecture

    def apply_to_model(self, model, architecture):
        for name, module in model.named_modules():
            if isinstance(module, nn.Conv2d):
                if 'num_filters' in architecture:
                    in_channels = module.in_channels
                    out_channels = architecture['num_filters']
                    kernel_size = architecture.get('kernel_size', module.kernel_size)
                    new_conv = nn.Conv2d(in_channels, out_channels, kernel_size, padding=kernel_size[0]//2)
                    setattr(model, name, new_conv)
            elif isinstance(module, nn.ReLU):
                if architecture.get('activation') == 'leaky_relu':
                    setattr(model, name, nn.LeakyReLU())
        
        if 'num_layers' in architecture:
            current_layers = len(list(model.children()))
            if architecture['num_layers'] > current_layers:
                for _ in range(architecture['num_layers'] - current_layers):
                    model.add_module(f'extra_layer_{_}', nn.Linear(64, 64))  # Adjust as needed
            elif architecture['num_layers'] < current_layers:
                for _ in range(current_layers - architecture['num_layers']):
                    model = nn.Sequential(*list(model.children())[:-1])
        
        return model

# Example usage
base_model = nn.Sequential(
    nn.Conv2d(3, 64, 3, padding=1),
    nn.ReLU(),
    nn.Conv2d(64, 64, 3, padding=1),
    nn.ReLU(),
    nn.MaxPool2d(2)
)

search_space = {
    'num_layers': (2, 5),
    'num_filters': (32, 128),
    'kernel_size': (3, 5),
    'activation': ['relu', 'leaky_relu']
}

nas = ProgressiveNAS(base_model, search_space)
# Ensure to provide train_loader and val_loader when calling search
# best_model = nas.search(num_iterations=10, train_loader=train_loader, val_loader=val_loader)