Following this instruction :
1. Strictly adhere to PEP-8 standards, ensuring code is clean, well-structured, and follows Python's style guide for readability.
2. Use appropriate modules like `typing` for type hints, ensuring code clarity, maintainability, and leveraging Python's standard libraries effectively.
3. Focus on time, space, robustness, scalability; write optimized, efficient code that can handle large datasets and scale seamlessly.
4. Implement comprehensive error handling, catching exceptions gracefully, providing meaningful error messages, and ensuring code robustness under unexpected conditions.
5. Ensure the developer agent writes advanced, maintainable Python code, balancing performance and readability while adhering to best practices.



TASK :make below pipeline generalised, training,evaluation,inference,visualization plots and weights save in the ouput folder 


import os
import time
import copy 
import numpy as np
import collections
from PIL import Image
import matplotlib.pyplot as plt
import plotly.gragh_objects as go
from plotly.subplots import make_subplots
import plotly.express as px
import torch
import torch.nn as nn
import torch torch.backends.cudnn as cudnn
import torch.optim as optim
from torch.optim.lr_scheduler import CosineAnealingLR
import torchvision
from torchvision import datasets,models
from torchvision,transforms import v2
cudnn.benchmark=True
plt.ion()
device=torch.device("cuda" if torch.cuda.is_available() "cpu")

NUM_CLASSES=10
IMAGE_SIZE=224
MEAN,STD=[0.4914,0.4822,0.4465],[0.247.0.243.0.261]

train_transforms=v2.Compose(
[
v2.Resize((IMAGE_SIZE,IMAGE_SIZE),
v2.AutoAugment(policy=v2.AutoAugmentPolicy.CIFAR10),
v2.ToTensor(),
v2.Normalize(MEAN,STD)
]
)

validation_transformers=v2.Compose(
[
v2.Resize((IMAGE_SIZE,IMAGE_SIZE)),
v2.ToTensor(),
v2.Normalize(MEAN,STD)

]
)
train_set=torchvision.datasets.CIFAR10(
root="./data",
train=True,
download=True,
transform=train_transforms

)
validation_set=torchvision.datasets.CIFAR10(
root="./data",
train=False,
download=True,
transform=validation_transforms

)

print(f"Train dataset size:{len(train_set:.4f)})
print("validation dataset size :{len(validation_set)})
def plot_class_distribution(dataset,dataset_name):

    labels=[y for _,y in dataset]
    counter=collections.Counter(labels)
    print(f"Class Image Counter for {dataset_name})
    print(counter,"\n")
    plt.bar(counter.keys(),counter.values())
    plt.xlabel("CLASS")
    plt.ylabel("number of images")
    plt.title(f"class Distribution for {dataset_names} Data")
    plt.show()
plot_class_distribution(train_set,"Training")
plot_class_distributions(validation_set,"validation")
batch_size=128
num_gpus=torch.cuda.device_count()
train_loader=torch.utils.data.DataLoader(
train_set,
batch_size=batch_size,
shuffle=True,
num_workers=2*num_gpus,
pin_memory=True
)
validation_loader=torch.utils.data.DataLoader(
validation_set,
batch_size=batch_size,
shuffle=False,
num_memory=2*num_gpus,
pin_memory=True

)
def imshow(tensor,title=Nome):
    image=torchvision.utils.make_grid(tensor).numpy().transpose((1,2,0))
    mean=np.array([0.4914,0.4822,0.4465])
    std=np.array([0.247,0.243,0.261])
    image=std*image+mean
    image=np.clip(image,0,1)
    plt.figure(figsize=(10,10))
    plt.imshow(image)
    if title:
       plt.title(title)
    plt.pause(0.001)
classes=("plane","car","bird","cat","deer","dog","frog","horse","ship","trucK")

# Get a batch of training data
for inputs, labels in train_loader:
    # Display the first 8 images from the batch
    imshow(inputs[:5], title=[classes[x] for x in labels[:5]])

    break  # Break the loop after it finishes



learning_rate=1e-4
def get_lr(opt):
    for param_group in opt.param_groups:
        return param_group["lr"]

def metrics_batch(output,target):
    pred=output.argmax(dim=1,keepdim=True)
    corrects=pred.eq(target.view_as(pred)).sum().item()
    return corrects
def loss_batch(loss_func,model_output,target,optimizer=None):
    loss_value=loss_func(model_output,traget)
    metric_batch=metrics_batch(model_output,target)
    if optimizer is not None:
       optimizer.zero_grad()
       loss_value.backward()
       optimizer.step()
     return loss_value.item(),metric_batch

def loss_epoch(model,loss_func,data_loader,check_id=False,optimizer=None):
    running_loss=0.0
    runnning_metric=0.0
    total_sample=len(data_loader.dataset)
    for inputs,labels in data_loader:
        inputs,labels=inputs.to(device) ,labels.to(device)
        model_output=model(inputs)
        loss_batch_value,metric_batch_value=loss_batch(
        loss_func,model_output,labels,optimizer
)
        running_loss+=loss_batch_value
        if metric_batch_value is not None:
           running_metric+=metric_batch_value
        if check_id:
           break
     average_loss=running_loss/float(total_sample)
     average_metric=(
      running_metric/float(total_samples) if total_samples > 0 else None
)

    return average_loss,average_metric

def train_val(model,paramas,verbose=False):
    epochs=params["epochs"]
    loss_func=params["loss_func"]
    opt=params["optimizer"]
    train_dl=params["val_dl"]
    val_dl=params["val_dl"]
    check_id=params["check_id"]
    lr_scheduler=params["lr_scheduler"]
    path=params["path"]
    loss_history={"train":[],"val":[]}
    metric_history={"train":[],"val":[]}
    best_model_wts=copy.deepcopy(model.state_dict())
    best_loss=float("inf")
    for epoch in range(epochs):
        current_lr=get_lr(opt)
        if verbose:
           print(f"Epoch {epoch}/{epoch-1}, current lr={current_lr}")
           
    

        model.train()
                    train_loss,train_metric=loss_epoch(model,loss_func,train_dl,check_id,opt)
        loss_history["train"].append(train_loss)
        metric_history["train"].append(train_metrics)
        model.eval()
        with torch.no_grad():
             val_loss,val_metric=loss_epoch(model,loss_func,val_dl,check_id)
              if val_loss<best_loss:
                 best_loss=val_loss
                 best_model_wts=copy.deepcopy(model.state_dict())
                 torch.save(model.state_dict(),path)
                 if verbose: 
                    print("copied best weight weight!")
              loss_history["val"].append(val_loss)
              metric_history["val"].append(val_metrics)
              lr_scheduler.step()
              if verbose:
                 print(f"train loss:{train_loss.6f},train_accuracy:{100*train_metric:.2f}.validation loss {val_loss.6f},validation accuracy:{100*valmetric:.2f}
                  
    model.load_state_dict(best_model_wts)
    return model,loss_history,metric_history

 

def plot_out(loss_hist, metric_hist, epochs=None):
    """
    Plot training and validation loss, and training and validation metrics over epochs.

    Args:
        loss_hist (dict): Dictionary containing 'train' and 'val' keys for loss history.
        metric_hist (dict): Dictionary containing 'train' and 'val' keys for metric history.
        epochs (int): Number of epochs.
    """

    # Create subplot with two columns
    fig = make_subplots(
        rows=1, cols=2, subplot_titles=["Model Loss", "Model Accuracy"]
    )

    # Plot Loss History
    for phase, color in zip(["train", "val"], ["#F1C40F", "#232323"]):
        fig.add_trace(
            go.Scatter(
                x=list(range(1, epochs + 1)),
                y=loss_hist[phase],
                name=phase,
                mode="lines",
                line_color=color,
            ),
            row=1,
            col=1,
        )

    # Plot Metric History
    for phase, color in zip(["train", "val"], ["#F1C40F", "#232323"]):
        fig.add_trace(
            go.Scatter(
                x=list(range(1, epochs + 1)),
                y=metric_hist[phase],
                name=phase,
                mode="lines",
                line_color=color,
            ),
            row=1,
            col=2,
        )

    # Update layout and display the plot
    fig.update_layout(
        template="plotly_white",
        showlegend=False,
        title="Loss & Accuracy History",
        height=400,
    )
    fig.update_layout(yaxis2=dict(range=[0.4, 1]))
    fig.show()

# Create the base model from the pre-trained model VGG19
pre_vgg19 = models.vgg19(weights='IMAGENET1K_V1')
pre_vgg19
# change the number of classes
pre_vgg19.classifier[6].out_features = NUM_CLASSES

# freeze convolution weights
for param in pre_vgg19.features.parameters():
    param.requires_grad = False

pre_vgg19
# Send earlier defined model to device
device = torch.device("cuda:0")
pre_vgg19 = pre_vgg19.to(device)

# Loss function, optimizer, LR scheduler
loss_func = nn.CrossEntropyLoss(reduction="sum")
optimizer = optim.Adam(pre_vgg19.parameters(), lr=learning_rate)
lr_scheduler = CosineAnnealingLR(optimizer, T_max=5, eta_min=1e-6)

# Set Training Parameters
training_params = {
    "epochs": 10,
    "optimizer": optimizer,
    "loss_func": loss_func,
    "train_dl": train_loader,
    "val_dl": validation_loader,
    "check_id": False,
    "lr_scheduler": lr_scheduler,
    "path": "pre_vgg19.pt",
}

# Train and validate the model
pre_vgg19, loss_history, metric_history = train_val(
    pre_vgg19, training_params, verbose=False
)

# Plot History
plot_out(loss_history, metric_history, epochs=training_params["epochs"])

def visualize_model_predictions(model, img_path):
    """
    Visualizes model predictions on a single image.

    Args:
        model (torch.nn.Module): The trained model.
        img_path (str): Path to the input image.

    Returns:
        None
    """
    # Set the model to evaluation mode
    was_training = model.training
    model.eval()

    # Load and preprocess the input image
    img = Image.open(img_path)
    img = validation_transforms(img)
    img = img.unsqueeze(0)
    img = img.to(device)

    # Perform forward pass to obtain predictions
    with torch.no_grad():
        outputs = model(img)
        _, preds = torch.max(outputs, 1)

        # Display the input image and predicted class
        ax = plt.subplot(2, 2, 1)
        ax.axis("off")
        ax.set_title(f"Predicted: {classes[preds[0]]}")
        imshow(img.cpu().data[0])

    # Set the model back to its original training mode
    model.train(mode=was_training)
# Visualize model predictions
visualize_model_predictions(pre_vgg19, img_path="/scratch/hemanth/Hemanth/computer-vision-course/notebooks/Unit 2 - Convolutional Neural Networks/List-of-Gaited-Horse-Breeds.jpg")

plt.ioff()
plt.show()