<jupyter_start><jupyter_text><jupyter_code>!pip -q install vit_pytorch linformer
!unzip -uq "bbb.zip"
from __future__ import print_function

import glob
from itertools import chain
import os
import random
import zipfile
import shutil, os

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from linformer import Linformer
from PIL import Image
from sklearn.model_selection import train_test_split
from torch.optim.lr_scheduler import StepLR
from torch.utils.data import DataLoader, Dataset
from torchvision import datasets, transforms
from tqdm.notebook import tqdm

from vit_pytorch.efficient import ViT

print(f"Torch: {torch.__version__}")
# Training settings
batch_size = 64
epochs = 50
lr = 3e-5
gamma = 0.7
seed = 42
def seed_everything(seed):
    random.seed(seed)
    os.environ['PYTHONHASHSEED'] = str(seed)
    np.random.seed(seed)
    torch.manual_seed(seed)
    torch.cuda.manual_seed(seed)
    torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True

seed_everything(seed)
# example of cropping an image
path = 'bbb'
file = '10129-AUG-1.jpg'
im = Image.open(os.path.join(path,file))
with open('bbb/10129-AUG-1.txt','r') as f:
  text = f.read()
  x = float(text.split()[1])*im.size[0]
  y = float(text.split()[2])*im.size[1]
  w = float(text.split()[3])*im.size[0]
  h = float(text.split()[4])*im.size[1]
cropped_image = im.crop((x-w/2, y-h/2, x+w/2, y+h/2))

cropped_image

# rm -rf bbb
# cropping all images and store in cropped folder
if not os.path.exists('cropped'): os.mkdir('cropped') 
path = 'bbb'
files = os.listdir(path)
for index, file in enumerate(files):
    if file.endswith('.jpg'):
      im = Image.open(os.path.join(path,file))
      with open(os.path.join(path, ''.join([file.split(".")[0],'.txt'])), 'r') as f:
        text = f.read()
        x = float(text.split()[1])*im.size[0]
        y = float(text.split()[2])*im.size[1]
        w = float(text.split()[3])*im.size[0]
        h = float(text.split()[4])*im.size[1]
      cropped_image = im.crop((x-w/2, y-h/2, x+w/2, y+h/2))
      cropped_image.save(os.path.join('cropped', ''.join([file.split(".")[0],'.jpg'])))  
      shutil.copy(os.path.join(path, ''.join([file.split(".")[0],'.txt'])), 'cropped')
device = 'cuda'
os.makedirs('data', exist_ok=True)
data_dir = 'bbb'
cropped_dir = 'cropped'
# data_dir = 'cropped'
# train_dir = 'data/train'
# test_dir = 'data/test'
# with zipfile.ZipFile('train.zip') as train_zip:
#     train_zip.extractall('data')  
# with zipfile.ZipFile('test.zip') as test_zip:
#     test_zip.extractall('data')
train_list = glob.glob(os.path.join(data_dir,'*.jpg'))

labels=[]

for root in train_list:
  path = os.path.dirname(root)
  name = root.split("/")[-1].split(".")[0]
  with open(os.path.join(path, ''.join([name,'.txt'])), 'r') as f:
    text = f.read()
    label = int(text.split()[0])+1
    labels.append(label)


train_list, test_list = train_test_split(train_list, 
                        test_size=0.33,
                        stratify=labels,
                       random_state=seed)
# train_list = glob.glob(os.path.join(train_dir,'*.jpg'))
# test_list = glob.glob(os.path.join(test_dir, '*.jpg'))
print(f"Train Data: {len(train_list)}")
print(f"Test Data: {len(test_list)}")
print(labels)
print(train_list)
labels = []
for root in train_list:
  path = os.path.dirname(root)
  name = root.split("/")[-1].split(".")[0]
  with open(os.path.join(path, ''.join([name,'.txt'])), 'r') as f:
    text = f.read()
    label = int(text.split()[0])+1
    labels.append(label)

random_idx = np.random.randint(1, len(train_list), size=9)
fig, axes = plt.subplots(3, 3, figsize=(16, 12))

for idx, ax in enumerate(axes.ravel()):
    img = Image.open(train_list[idx])
    ax.set_title(labels[idx])
    ax.imshow(img)
# split dataset
train_list, valid_list = train_test_split(train_list, 
                        test_size=0.2,
                        stratify=labels,
                       random_state=seed)


print(f"Train Data: {len(train_list)}")
print(f"Validation Data: {len(valid_list)}")
print(f"Test Data: {len(test_list)}")
train_transforms = transforms.Compose(
    [
        transforms.Resize((224, 224)),
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
    ]
)

val_transforms = transforms.Compose(
    [
        transforms.Resize((224, 224)),
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
    ]
)


test_transforms = transforms.Compose(
    [
        transforms.Resize((224, 224)),
        transforms.RandomResizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
    ]
)
class trucksDataset(Dataset):
    def __init__(self, file_list, transform=None):
        self.file_list = file_list
        self.transform = transform

    def __len__(self):
        self.filelength = len(self.file_list)
        return self.filelength

    def __getitem__(self, idx):
        img_path = self.file_list[idx]
        img = Image.open(img_path)
        img_transformed = self.transform(img)

        name = img_path.split("/")[-1].split(".")[0]
        path = os.path.dirname(img_path)
        with open(os.path.join(path, ''.join([name,'.txt'])), 'r') as f:
          text = f.read()
          label = int(text.split()[0])
        # label = img_path.split("/")[-1].split(".")[0]
        # label = 1 if label == "dog" else 0

        return img_transformed, label
train_data = trucksDataset(train_list, transform=train_transforms)
valid_data = trucksDataset(valid_list, transform=test_transforms)
test_data = trucksDataset(test_list, transform=test_transforms)

train_loader = DataLoader(dataset = train_data, batch_size=batch_size, shuffle=True )
valid_loader = DataLoader(dataset = valid_data, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(dataset = test_data, batch_size=batch_size, shuffle=True)

print(len(train_data), len(train_loader))
print(len(valid_data), len(valid_loader))
efficient_transformer = Linformer(
    dim=128,
    seq_len=49+1,  # 7x7 patches + 1 cls-token
    depth=12,
    heads=8,
    k=64
)
model = ViT(
    dim=128,
    image_size=224,
    patch_size=32,
    num_classes=13,
    transformer=efficient_transformer,
    channels=3,
).to(device)

# loss function
criterion = nn.CrossEntropyLoss()
# optimizer
optimizer = optim.Adam(model.parameters(), lr=lr)
# scheduler
scheduler = StepLR(optimizer, step_size=1, gamma=gamma)
epoch50 = []
accuracy_train = []
accuracy_test = []
for epoch in range(epochs):
    epoch_loss = 0
    epoch_accuracy = 0

    for data, label in tqdm(train_loader):
        data = data.to(device)
        label = label.to(device)

        output = model(data)
        loss = criterion(output, label)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        acc = (output.argmax(dim=1) == label).float().mean()
        epoch_accuracy += acc / len(train_loader)
        epoch_loss += loss / len(train_loader)

    with torch.no_grad():
        epoch_val_accuracy = 0
        epoch_val_loss = 0
        for data, label in valid_loader:
            data = data.to(device)
            label = label.to(device)

            val_output = model(data)
            val_loss = criterion(val_output, label)

            acc = (val_output.argmax(dim=1) == label).float().mean()
            epoch_val_accuracy += acc / len(valid_loader)
            epoch_val_loss += val_loss / len(valid_loader)

    print(
        f"Epoch : {epoch+1} - loss : {epoch_loss:.4f} - acc: {epoch_accuracy:.4f} - val_loss : {epoch_val_loss:.4f} - val_acc: {epoch_val_accuracy:.4f}\n"
    )
    epoch50.append(epoch)
    accuracy_train.append(epoch_accuracy)
    accuracy_test.append(epoch_val_accuracy)
epoch50 = [x+1 for x in epoch50]


plt.plot(epoch50, accuracy_train,label='training')
plt.plot(epoch50, accuracy_test,label='test')
plt.legend(loc='lower left')
plt.xlabel('Epochs')
plt.ylabel('accuracy')
plt.title('Model Accuracy')<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.14623707050291285, 'fraction_numerical': 0.01830935679467364, 'mean_word_length': 3.052023121387283, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '294975', 'license': 'MIT', 'n_tokens_mistral': 3110, 'n_tokens_neox': 2992, 'n_words': 548, 'path': 'Truck_Classifier.ipynb', 'repo_name': 'mashihan123/ViT-Truck-Images', 'size': '8411'}