<jupyter_start><jupyter_text>Load relevant libraries<jupyter_code>import argparse
import math
import os
import random
import sys
import time
import glob
import yaml

import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import Dataset, DataLoader
from torch.optim import SGD, Adam, lr_scheduler
from torchvision.ops import FeaturePyramidNetwork
from torchvision import transforms
from torchvision import models
import tqdm
from PIL import Image
from swin_transformer_pytorch import SwinTransformer
import matplotlib.pyplot as plt

import warnings
warnings.filterwarnings("ignore")

#!pip install git+https://github.com/facebookresearch/fvcore.git
#!pip install swin-transformer-pytorch
#!pip install ipywidgets<jupyter_output><empty_output><jupyter_text>Dataloaders<jupyter_code>#############################################
dataset_folder = "datasets/gtacar"
batch_size = 8
num_workers = 4

transform = transforms.Compose([
    transforms.ConvertImageDtype(torch.float),
    transforms.CenterCrop((638,1914)),
    transforms.Resize((256,768)),
    # transforms.Resize((224,672)),
    # transforms.Resize((281,512)),
    # transforms.Pad((0,115,0,116))
])
#############################################

train_images = "images/train"
train_labels = "labels/train"
val_images = "images/val"
val_labels = "labels/val"

class ImageDataset(Dataset):
    def __init__(self, dataset_folder, image_path, label_path, transform=None, nc=3):
        super().__init__()
        self.dataset_folder = dataset_folder
        self.image_filenames = glob.glob(os.path.join(dataset_folder, image_path, "*"))
        self.label_path = label_path
        self.transform = transform
        self.nc = nc

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

    def __getitem__(self, idx):
        X = transforms.functional.pil_to_tensor(Image.open(self.image_filenames[idx]))
        label_filename = self.image_filenames[idx].split("/")[-1].split(".")[0] + ".txt"
        label_path = os.path.join(self.dataset_folder, self.label_path, label_filename)

        with open(label_path) as f:
            line_data = f.readlines()[0].split(" ")
            label = torch.tensor(int(line_data[0]),dtype=torch.int64)
            bbox = [float(x) for x in line_data[1:]]
        # y = torch.zeros((3,))
        # y[label] = 1

        mask = torch.zeros(X.shape[1:])
        h,w = X.shape[1:]
        mask[int(h*(bbox[1]-bbox[3])):int(h*(bbox[1]+bbox[3])),int(w*(bbox[0]-bbox[2])):int(w*(bbox[0]+bbox[2]))] = 1
        mask = mask.unsqueeze(0)

        if not self.transform is None:
            X = self.transform(X)
            mask = self.transform(mask)
        
        mask = mask.squeeze(0)
        mask = mask.numpy()
        cols, rows = np.nonzero(mask)
        if len(cols)==0: 
            bbox = torch.zeros(4, dtype=torch.float32)
        else:
            top_row = np.min(rows)
            left_col = np.min(cols)
            bottom_row = np.max(rows)
            right_col = np.max(cols)
            bbox = torch.tensor([left_col, top_row, right_col, bottom_row], dtype=torch.float32)

        # im = X.squeeze().permute(1,2,0).numpy()
        # plt.imshow(im)
        # plt.show()
        # plt.imshow(mask)
        # plt.show()

        return X, label, bbox
    
    def get_class_weights(self):
        weights = torch.zeros((self.nc,))
        for filename in self.image_filenames:
            label_filename = filename.split("/")[-1].split(".")[0] + ".txt"
            label_path = os.path.join(self.dataset_folder, self.label_path, label_filename)
            
            with open(label_path) as f:
                y = torch.tensor(int(f.readlines()[0].split(" ")[0])-1).to(torch.int64)

            weights[y] += 1
        
        nmax = torch.max(weights)
        weights = nmax / weights
        return weights

tr_dataset = ImageDataset(dataset_folder, train_images, train_labels, transform)
weights = tr_dataset.get_class_weights().to(device)
print(len(tr_dataset))
tr_dataset = DataLoader(tr_dataset, batch_size, shuffle=True, num_workers=num_workers, pin_memory=True)
val_dataset = ImageDataset(dataset_folder, val_images, val_labels, transform)
print(len(val_dataset))
val_dataset = DataLoader(val_dataset, batch_size, shuffle=False, num_workers=num_workers, pin_memory=True)<jupyter_output>6058
1515<jupyter_text>Model<jupyter_code>############################
lr = 1e-3
############################

#model = torch.hub.load('zhanghang1989/ResNeSt', 'resnest269', pretrained=False, dilated=False, num_classes=3).to(device)
# model = SwinTransformer(
#     hidden_dim=96,
#     layers=(2, 2, 6, 2),
#     heads=(3, 6, 12, 24),
#     channels=3,
#     num_classes=3,
#     head_dim=32,
#     window_size=8,
#     downscaling_factors=(4, 2, 2, 2),
#     relative_pos_embedding=True
# ).to(device)
# model = SwinTransformer(
#     hidden_dim=96,
#     layers=(2, 2, 6, 2),
#     heads=(3, 6, 12, 24),
#     channels=3,
#     num_classes=3,
#     head_dim=32,
#     window_size=7,
#     downscaling_factors=(4, 2, 2, 2),
#     relative_pos_embedding=True
# ).to(device)

class BBReg_SwinTransformer(nn.Module):
    def __init__(self):
        super(BBReg_SwinTransformer, self).__init__()
        swin_transformer = SwinTransformer(
            hidden_dim=96,
            layers=(2, 2, 6, 2),
            heads=(3, 6, 12, 24),
            channels=3,
            num_classes=3,
            head_dim=32,
            window_size=8,
            downscaling_factors=(4, 2, 2, 2),
            relative_pos_embedding=True
        )
        self.swin_transformer = nn.Sequential(*list(swin_transformer.children())[:-1])
        
        # Prediction head
        self.clf = nn.Sequential(
            nn.BatchNorm1d(768),
            nn.Linear(768,128),
            nn.ReLU(),
            nn.Linear(128,3),
            nn.Softmax()
        )
        self.bboxreg = nn.Sequential(
            nn.BatchNorm1d(768),
            nn.Linear(768,4)
        )


    def forward(self, x):
        x = self.swin_transformer(x)
        x = x.mean(dim=[2, 3])
        cls = self.clf(x)
        bbox = self.bboxreg(x) 
        return cls, bbox

# model = models.detection.faster_rcnn.fasterrcnn_resnet50_fpn(pretrained=True)
# in_features = model.roi_heads.box_predictor.cls_score.in_features
# model.roi_heads.box_predictor = models.detection.faster_rcnn.FastRCNNPredictor(in_features, 2)
# model = model.to(device)
model = BBReg_SwinTransformer().to(device)
criterion = torch.nn.CrossEntropyLoss(weight=weights)
bbox_criterion = torch.nn.L1Loss()
# criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.AdamW(model.parameters(), lr=lr)

# Test forward pass
# for X, Y, B in tr_dataset:
#     # X = list(x.to(device) for x in X)
#     # targets = [{"boxes" : b.unsqueeze(0).to(device), "labels" : y.unsqueeze(0).to(torch.int64).to(device)} if y == torch.tensor(-1) else {"boxes" : torch.zeros((0,4)).to(device), "labels" : torch.zeros((0,1)).to(torch.int64).to(device)} for b,y in zip(B,Y)]
#     # print(targets[1]["labels"].dtype)
#     # print(targets)
#     X = X.to(device)
#     outputs = model(X)
#     # outputs = model(X,targets)
#     print(outputs)
#     break<jupyter_output><empty_output><jupyter_text>Train<jupyter_code>######################################
num_epochs = 50
######################################

def train_epoch(data_loader, model, criterion, optimizer, device):
    """
    Train the `model` for one epoch of data from `data_loader`
    Use `optimizer` to optimize the specified `criterion`
    """
    
    model = model.train()

    correct, total = 0, 0
    running_loss = []

    for X, Y, B in tqdm.tqdm(data_loader):
        X = X.to(device)
        Y = Y.to(device)
        B = B.to(device)
        # X = list(x.to(device) for x in X)
        # targets = [{"boxes" : b.unsqueeze(0).to(device), "labels" : y.unsqueeze(0).to(device)} if not torch.all(b == torch.zeros((1,4))) else {"boxes" : torch.zeros((0,4)).to(device), "labels" : torch.zeros((0,1)).to(device)} for b,y in zip(B,Y)]
        # targets = [{"boxes" : b.unsqueeze(0).to(device), "labels" : y.unsqueeze(0).to(torch.int64).to(device)} if y == torch.tensor(-1) else {"boxes" : torch.zeros((0,4)).to(device), "labels" : torch.zeros((0,1)).to(torch.int64).to(device)} for b,y in zip(B,Y)]
        # clear parameter gradients
        optimizer.zero_grad()

        # forward + backward + optimize
        output, bbox = model(X)
        loss = criterion(output, Y) + bbox_criterion(bbox,B) / 1000
        loss.backward()
        optimizer.step()

        # loss_dict = model(X, targets)
        # loss = sum(loss for loss in loss_dict.values())

        predicted = torch.argmax(output, dim=1)
        total += Y.shape[0]
        correct += (predicted == Y).sum().item()

        # calculate loss and accuracy
        running_loss.append(loss.item())

    train_loss = np.mean(running_loss)
    train_acc = correct / total

    return train_loss, train_acc

def evaluate_epoch(val_loader, model, criterion, device):
    model.eval()

    with torch.no_grad():
        Y_true, Y_pred = [], []
        correct, total = 0, 0
        running_loss = []
        for X, Y, B in tqdm.tqdm(val_loader):
            X = X.to(device)
            Y = Y.to(device)

            output, bbox = model(X)
            predicted = torch.argmax(output, dim=1)
            
            total += Y.shape[0]
            correct += (predicted == Y).sum().item()

            # loss calculation
            running_loss.append(criterion(output, Y).item())
            
        val_loss = np.mean(running_loss)
        val_acc = correct / total

    return val_loss, val_acc

def train(model, criterion, optimizer, tr_loader, val_loader, device):
    stats = {"epoch" : [], "tr_acc" : [], "tr_loss" : [], "val_acc" : [], "val_loss" : []}
    best_val_acc = 0.64
    for epoch in range(1, num_epochs+1):
        print("Epoch {}:".format(epoch))
        stats["epoch"].append(epoch)

        # Train model
        tr_loss, tr_acc = train_loss, train_acc = train_epoch(tr_loader, model, criterion, optimizer, device)

        print("Train loss = {}, train accuracy = {}".format(tr_loss, tr_acc))
        stats["tr_acc"].append(tr_acc)
        stats["tr_loss"].append(tr_loss)

        # Evaluate model
        val_loss, val_acc = evaluate_epoch(val_loader, model, criterion, device)

        print("Validation loss = {}, validation accuracy = {}".format(val_loss, val_acc))
        stats["val_acc"].append(val_acc)
        stats["val_loss"].append(val_loss)

        if val_acc > best_val_acc:
            torch.save({
                'epoch': epoch,
                'model_state_dict': model.state_dict(),
            }, "epoch{}.valacc{}.pth".format(epoch, val_acc))
            best_val_acc = val_acc


stats = train(model, criterion, optimizer, tr_dataset, val_dataset, device)<jupyter_output>Epoch 1:<jupyter_text>Inference<jupyter_code>def inference(model, folder_path, output_path, transform, device):
    with open(output_path, "w") as f:
        f.write("guid/image,label\n")

        model.eval()
        image_paths = glob.glob(os.path.join(folder_path, "*"))
        for image_path in tqdm.tqdm(image_paths):
            im = transforms.functional.pil_to_tensor(Image.open(image_path))
            im = transform(im).to(device)
            im = im.unsqueeze(0)
            pred, bbox = model(im)
            prediction = torch.argmax(pred).item()

            guid = image_path.split("/")[-1].split(".")[0]
            im_no = image_path.split("/")[-1].split(".")[1]

            f.write(guid + "/" + im_no + "," + str(prediction) + "\n")

checkpoint = torch.load("epoch27.valacc0.7148514851485148.pth")
model.load_state_dict(checkpoint['model_state_dict'])   
inference(model, os.path.join(dataset_folder,"images/test"), "output.txt", transform, device)
def thresh_evaluate_epoch(val_loader, model, criterion, thresh, device):
    model.eval()

    with torch.no_grad():
        Y_true, Y_pred = [], []
        correct, total = 0, 0
        running_loss = []
        for X, Y, B in tqdm.tqdm(val_loader):
            X = X.to(device)
            Y = Y.to(device)
            print(Y)

            output, bbox = model(X)
            predicted = torch.argmax(output, dim=1)
            print(output)
            maximums = torch.max(output,dim=1)
            break
            predicted[torch.max(output,dim=1) < torch.ones_like(predicted)*thresh] = 0
            
            total += Y.shape[0]
            correct += (predicted == Y).sum().item()

            # loss calculation
            running_loss.append(criterion(output, Y).item())
            
        val_loss = np.mean(running_loss)
        val_acc = correct / total

    return val_loss, val_acc

print(thresh_evaluate_epoch(val_dataset, model, criterion, 0.4, device))<jupyter_output>0%|                                                   | 0/190 [00:00<?, ?it/s]
{'fraction_non_alphanumeric': 0.15731726162880544, 'fraction_numerical': 0.021480451244011746, 'mean_word_length': 2.7989433519225124, 'pattern_counts': {'":': 0, '<': 20, '<?xml version=': 0, '>': 19, 'https://': 1, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '5', 'dir': 'jupyter-structured-clean-dedup', 'id': '515773', 'license': 'MIT', 'n_tokens_mistral': 4546, 'n_tokens_neox': 4343, 'n_words': 802, 'path': 'swin/swin_model.ipynb', 'repo_name': 'bruceli-rw0/rob535-perception', 'size': '12942'}