text_plain_output_1.png
from torch.utils.data import Dataset, DataLoader

import csv

import cv2

import numpy as np

import numpy as np # linear algebra

import random

import torch

import torch.nn as nn

TRAIN_PATH = '/kaggle/input/captcha-hacker-2023-spring/dataset/train'
TEST_PATH = '/kaggle/input/captcha-hacker-2023-spring/dataset/test'
device = 'cpu'

alphabets = 'abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ0123456789'
alphabets2index = {alphabet: i for i, alphabet in enumerate(alphabets)}

class Task1Dataset(Dataset):
    def __init__(self, data, root, return_filename=False):
        self.data = [sample for sample in data if sample[0].startswith('task1')]
        self.return_filename = return_filename
        self.root = root
    def __getitem__(self, index):
        filename, label = self.data[index]
        img = cv2.imread(f'{self.root}/{filename}')
        img = cv2.resize(img, (32, 32))
        img = np.mean(img, axis=2)
        if self.return_filename:
            return (torch.FloatTensor((img - 128) / 128), filename)
        else:
            return (torch.FloatTensor((img - 128) / 128), alphabets2index[label])
    def __len__(self):
        return len(self.data)

class Model(nn.Module):
    def __init__(self):
        super().__init__()
        self.layers = nn.Sequential(nn.Linear(1024, 512), nn.LeakyReLU(), nn.Linear(512, len(alphabets)))
    def forward(self, x):
        batch_size, h, w = x.shape
        x = x.view(batch_size, h * w)
        return self.layers(x)

train_data = []
val_data = []
with open(f'{TRAIN_PATH}/annotations.csv', newline='') as csvfile:
    for row in csv.reader(csvfile, delimiter=','):
        if random.random() < 0.8:
            train_data.append(row)
        else:
            val_data.append(row)
train_ds = Task1Dataset(train_data, root=TRAIN_PATH)
train_dl = DataLoader(train_ds, batch_size=100, num_workers=4, drop_last=True, shuffle=True)
val_ds = Task1Dataset(val_data, root=TRAIN_PATH)
val_dl = DataLoader(val_ds, batch_size=100, num_workers=4, drop_last=False, shuffle=False)

model = Model().to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
loss_fn = nn.CrossEntropyLoss()
for epoch in range(50):
    print(f'Epoch [{epoch}]')
    model.train()
    for image, label in train_dl:
        image = image.to(device)
        label = label.to(device)
        pred = model(image)
        loss = loss_fn(pred, label)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    sample_count = 0
    correct_count = 0
    model.eval()
    for image, label in val_dl:
        image = image.to(device)
        label = label.to(device)
        pred = model(image)
        loss = loss_fn(pred, label)
        pred = torch.argmax(pred, dim=1)
        sample_count += len(image)
        correct_count += (label == pred).sum()
    print('accuracy (validation):', correct_count / sample_count)