<kaggle_start><code># !curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py
# !python pytorch-xla-env-setup.py --version 1.7 --apt-packages libomp5 libopenblas-dev
# !pip install -U pytorch-lightning albumentations
import pandas as pd
from PIL import Image
import os
import numpy as np
from io import BytesIO
from matplotlib import pyplot as plt
from tqdm.notebook import tqdm
import traceback
import pandas as pd
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
from torchvision import models
import cv2
from pytorch_lightning import LightningModule, Trainer, seed_everything
from pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor
from pytorch_lightning.profiler import AdvancedProfiler
from pytorch_lightning.loggers import TensorBoardLogger
import time
import sys
import gc
import pickle

batch_size = 128
gradient_accum = 1
epoch = 20
lr = 5e-3
wd = 5e-4
MODEL = "tf_efficientnet_b4_ns"
from collections import defaultdict

words = defaultdict(int)
for fn in tqdm(
    os.listdir("../input/dcic-2022-ocr/training_dataset-909058/training_dataset")
):
    for c in fn.split(".")[0]:
        words[c] += 1
sorted(list(words.items())), len(words)
vocab = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"


class Tokenizer:
    def __init__(self):
        self.vocab = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
        self.vocab_lookup = {}
        for i, c in enumerate(self.vocab):
            self.vocab_lookup[c] = i

    def __call__(self, s):
        ret = np.zeros(len(s), dtype=np.int64)
        for i, c in enumerate(s):
            ret[i] = self.vocab_lookup[c]
        return ret

    def untokenize(self, s):
        ret = []
        for i in range(len(s)):
            ret.append(self.vocab[s[i]])
        return "".join(ret)


tokenizer = Tokenizer()
tokenizer.untokenize(tokenizer("123jik"))
imgs = os.listdir("../input/dcic-2022-ocr/training_dataset-909058/training_dataset")
import timm
from torchvision import transforms

cfg = timm.data.resolve_data_config({}, MODEL)
cfg
import aug_lib

aug_train = transforms.Compose(
    [
        timm.data.RandomResizedCropAndInterpolation(
            (128, 256), scale=(0.9, 1.2), ratio=(2, 3), interpolation="random"
        ),
        aug_lib.TrivialAugment(),
        transforms.ToTensor(),
        transforms.Normalize(mean=cfg["mean"], std=cfg["std"]),
    ]
)
aug_val = timm.data.create_transform(
    input_size=(128, 256),
    crop_pct=1,
    is_training=False,
    interpolation="bilinear",
)


class MyDataset(Dataset):
    def __init__(self, imgs, tokenizer, aug):
        self.path = "../input/dcic-2022-ocr/training_dataset-909058/training_dataset"
        self.imgs = imgs
        self.tokenizer = tokenizer
        self.aug = aug

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

    def __getitem__(self, idx):
        img = Image.open(os.path.join(self.path, self.imgs[idx])).convert("RGB")
        img = self.aug(img)
        label = self.imgs[idx].split(".")[0]
        return img, self.tokenizer(label)


def reverse_transform(a):
    r = a[0]
    g = a[1]
    b = a[2]
    return torch.stack([r * 0.229 + 0.485, g * 0.224 + 0.456, b * 0.225 + 0.406], dim=0)


import random
import IPython

ds = MyDataset(imgs, tokenizer, aug_train)
for i in range(5):
    d = ds[random.randrange(0, len(ds))]
    IPython.display.display(transforms.ToPILImage()(reverse_transform(d[0])))
    print(tokenizer.untokenize(d[1]))


def one_hot(x, num_classes, on_value=1.0, off_value=0.0, device="cuda"):
    x = x.long().view(-1, 1)
    return torch.full((x.size()[0], num_classes), off_value, device=device).scatter_(
        1, x, on_value
    )


import torch
import torch.nn as nn
import torch.nn.functional as F


# +
class Model(LightningModule):
    def __init__(self, n_classes, lr, wd):
        super().__init__()
        self.save_hyperparameters()
        self.lr = lr
        self.cnn = timm.create_model(MODEL, pretrained=True)  # resnet 50
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.out = nn.Linear(self.cnn.num_features, n_classes * 4)
        self.crit = nn.CrossEntropyLoss()

    def forward(self, input_):
        x = self.cnn.forward_features(input_)
        x = self.avgpool(x).flatten(1)
        x = self.out(x).view(-1, 4, self.hparams.n_classes)
        return x

    def training_step(self, batch, batch_idx):
        self.train()
        crit = nn.MultiLabelSoftMarginLoss()
        input_tensor, target_tensor = batch
        with torch.no_grad():
            target = one_hot(target_tensor, self.hparams.n_classes).view(
                input_tensor.size(0), -1
            )
        decoder_out = self(input_tensor)  # N x L x C
        loss = crit(decoder_out.flatten(1), target)
        with torch.no_grad():
            target_tensor = target_tensor
            preds = decoder_out.argmax(2)
            acc_char = (preds == target_tensor).float().mean()
            acc_sent = (preds == target_tensor).all(dim=1).float().mean()
            self.log("train_loss", loss)
            self.log("train_acc_char", acc_char)
            self.log("train_acc_sent", acc_sent)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.AdamW(
            self.parameters(), lr=0, weight_decay=self.hparams.wd
        )
        scheduler = {
            "scheduler": torch.optim.lr_scheduler.OneCycleLR(
                optimizer,
                max_lr=self.lr,
                total_steps=self.trainer.max_steps,
                anneal_strategy="linear",
                cycle_momentum=False,
                pct_start=0.1,
            ),
            "interval": "step",
            "frequency": 1,
        }
        return [optimizer], [scheduler]

    def validation_step(self, batch, batch_idx):
        self.eval()
        with torch.no_grad():
            input_tensor, target_tensor = batch
            decoder_out = self(input_tensor)  # N x L x C
            loss = self.crit(decoder_out.permute(0, 2, 1), target_tensor)
            target_tensor = target_tensor
            preds = decoder_out.argmax(2)
            acc_char = (preds == target_tensor).float().mean()
            acc_sent = (preds == target_tensor).all(dim=1).float().mean()
            self.log("val_loss", loss)
            self.log("val_acc_char", acc_char)
            self.log("val_acc_sent", acc_sent)


from sklearn.model_selection import train_test_split, KFold
from sklearn.metrics import accuracy_score, f1_score, roc_auc_score

imgs = np.array(imgs)
imgs
from pytorch_lightning.utilities import rank_zero_only
from pytorch_lightning.loggers import LightningLoggerBase
from pytorch_lightning.loggers.base import rank_zero_experiment
from collections import defaultdict


class ConsoleLogger(LightningLoggerBase):
    def __init__(self):
        super().__init__()
        self.metrics = defaultdict(list)

    @property
    def name(self):
        return ""

    @property
    @rank_zero_experiment
    def experiment(self):
        # Return the experiment object associated with this logger.
        pass

    @property
    def version(self):
        # Return the experiment version, int or str.
        return ""

    @rank_zero_only
    def log_hyperparams(self, params):
        # params is an argparse.Namespace
        # your code to record hyperparameters goes here
        pass

    @rank_zero_only
    def log_metrics(self, metrics, step):
        # metrics is a dictionary of metric names and values
        # your code to record metrics goes here
        print("=" * 80)
        if "epoch" in metrics:
            print(f'epoch: {metrics["epoch"]}, steps: {step}')
        else:
            print(f"steps: {step}")
        for k, v in metrics.items():
            if k == "epoch":
                continue
            if isinstance(v, torch.Tensor):
                v = v.item()
            print(f"{k}: {v}")
            self.metrics[k].append((step, v))


import pytorch_lightning

kf = KFold(n_splits=20, random_state=42, shuffle=True)
val_data = []
fold_idx = 0
for train, val in kf.split(imgs):

    def train_func():
        train_imgs = imgs[train]
        val_imgs = imgs[val]
        ds_train = MyDataset(train_imgs, tokenizer, aug_train)
        train_loader = DataLoader(
            ds_train,
            batch_size=batch_size,
            num_workers=4,
            shuffle=True,
            pin_memory=True,
            drop_last=True,
        )
        ds_val = MyDataset(val_imgs, tokenizer, aug_val)
        val_loader = DataLoader(
            ds_val,
            batch_size=batch_size,
            num_workers=4,
            pin_memory=True,
        )
        max_steps = epoch * int(np.ceil(len(train_loader) / gradient_accum))
        print("*" * 80)
        print("training fold", fold_idx)
        print("train samples:", len(train_imgs), "max steps:", max_steps)
        checkpoint_callback = ModelCheckpoint(
            save_top_k=1,
            verbose=True,
            monitor="val_acc_sent",
            mode="max",
            filename="epoch{epoch}-step{step}-val_acc_sent{val_acc_sent:.2f}",
            auto_insert_metric_name=False,
            save_weights_only=True,
            save_last=True,
        )
        logger = ConsoleLogger()
        tb_logger = TensorBoardLogger(
            save_dir=".", version=f"fold_{fold_idx}", name="lightning_logs"
        )
        trainer = Trainer(
            logger=[logger, tb_logger],
            accumulate_grad_batches=gradient_accum,
            gpus=-1,
            callbacks=[LearningRateMonitor(), checkpoint_callback],
            flush_logs_every_n_steps=50,
            log_every_n_steps=50,
            max_steps=max_steps,
            benchmark=True,
            precision=16,
        )
        model = Model(len(vocab), lr, wd)
        #         lr_finder = trainer.tuner.lr_find(model, train_loader)
        #         lr_finder.plot(suggest=True)
        trainer.fit(model, train_dataloaders=train_loader, val_dataloaders=val_loader)

    train_func()
    gc.collect()
    pytorch_lightning.utilities.memory.garbage_collection_cuda()
    break


class MyTestDataset(Dataset):
    def __init__(self, imgs, aug):
        self.path = "../input/dcic-2022-ocr/test_dataset-220406/test_dataset/"
        self.imgs = imgs
        self.tokenizer = tokenizer
        self.aug = aug

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

    def __getitem__(self, idx):
        img = Image.open(os.path.join(self.path, self.imgs[idx])).convert("RGB")
        img = self.aug(img)
        return img


test_imgs = os.listdir("../input/dcic-2022-ocr/test_dataset-220406/test_dataset")
test_imgs = sorted(test_imgs, key=lambda x: int(x.split(".")[0]))
ds_test = MyTestDataset(test_imgs, aug_val)
test_loader = DataLoader(
    ds_test, batch_size=32, num_workers=4, pin_memory=True, shuffle=False
)
model = Model.load_from_checkpoint(
    "./_lightning_logs/_fold_0/checkpoints/last.ckpt"
).cuda()
tmp = []
with torch.no_grad():
    for batch in tqdm(test_loader):
        batch = batch.cuda()
        model.eval()
        pred = model(batch).argmax(2).cpu().numpy()
        tmp.append(pred)
preds = np.concatenate(tmp, axis=0)
df_preds = pd.DataFrame()
df_preds["num"] = [x.split(".")[0] for x in test_imgs]  # [:len(tags)]
df_preds["tag"] = [tokenizer.untokenize(x) for x in preds]
df_preds.to_csv("result.csv", index=None)
df_preds
<|endoftext|>
[{"Id": 87964295, "ScriptId": 24700958, "ParentScriptVersionId": NaN, "ScriptLanguageId": 9, "AuthorUserId": 3508829, "CreationDate": "02/16/2022 06:26:06", "VersionNumber": 3.0, "Title": "dcic ocr effnet multilabel", "EvaluationDate": "02/16/2022", "IsChange": true, "TotalLines": 409.0, "LinesInsertedFromPrevious": 16.0, "LinesChangedFromPrevious": 0.0, "LinesUnchangedFromPrevious": 393.0, "LinesInsertedFromFork": 32.0, "LinesDeletedFromFork": 85.0, "LinesChangedFromFork": 0.0, "LinesUnchangedFromFork": 377.0, "TotalVotes": 0}]
# !curl https://raw.githubusercontent.com/pytorch/xla/master/contrib/scripts/env-setup.py -o pytorch-xla-env-setup.py
# !python pytorch-xla-env-setup.py --version 1.7 --apt-packages libomp5 libopenblas-dev
# !pip install -U pytorch-lightning albumentations
import pandas as pd
from PIL import Image
import os
import numpy as np
from io import BytesIO
from matplotlib import pyplot as plt
from tqdm.notebook import tqdm
import traceback
import pandas as pd
import numpy as np
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
from torchvision import models
import cv2
from pytorch_lightning import LightningModule, Trainer, seed_everything
from pytorch_lightning.callbacks import ModelCheckpoint, LearningRateMonitor
from pytorch_lightning.profiler import AdvancedProfiler
from pytorch_lightning.loggers import TensorBoardLogger
import time
import sys
import gc
import pickle

batch_size = 128
gradient_accum = 1
epoch = 20
lr = 5e-3
wd = 5e-4
MODEL = "tf_efficientnet_b4_ns"
from collections import defaultdict

words = defaultdict(int)
for fn in tqdm(
    os.listdir("../input/dcic-2022-ocr/training_dataset-909058/training_dataset")
):
    for c in fn.split(".")[0]:
        words[c] += 1
sorted(list(words.items())), len(words)
vocab = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"


class Tokenizer:
    def __init__(self):
        self.vocab = "0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ"
        self.vocab_lookup = {}
        for i, c in enumerate(self.vocab):
            self.vocab_lookup[c] = i

    def __call__(self, s):
        ret = np.zeros(len(s), dtype=np.int64)
        for i, c in enumerate(s):
            ret[i] = self.vocab_lookup[c]
        return ret

    def untokenize(self, s):
        ret = []
        for i in range(len(s)):
            ret.append(self.vocab[s[i]])
        return "".join(ret)


tokenizer = Tokenizer()
tokenizer.untokenize(tokenizer("123jik"))
imgs = os.listdir("../input/dcic-2022-ocr/training_dataset-909058/training_dataset")
import timm
from torchvision import transforms

cfg = timm.data.resolve_data_config({}, MODEL)
cfg
import aug_lib

aug_train = transforms.Compose(
    [
        timm.data.RandomResizedCropAndInterpolation(
            (128, 256), scale=(0.9, 1.2), ratio=(2, 3), interpolation="random"
        ),
        aug_lib.TrivialAugment(),
        transforms.ToTensor(),
        transforms.Normalize(mean=cfg["mean"], std=cfg["std"]),
    ]
)
aug_val = timm.data.create_transform(
    input_size=(128, 256),
    crop_pct=1,
    is_training=False,
    interpolation="bilinear",
)


class MyDataset(Dataset):
    def __init__(self, imgs, tokenizer, aug):
        self.path = "../input/dcic-2022-ocr/training_dataset-909058/training_dataset"
        self.imgs = imgs
        self.tokenizer = tokenizer
        self.aug = aug

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

    def __getitem__(self, idx):
        img = Image.open(os.path.join(self.path, self.imgs[idx])).convert("RGB")
        img = self.aug(img)
        label = self.imgs[idx].split(".")[0]
        return img, self.tokenizer(label)


def reverse_transform(a):
    r = a[0]
    g = a[1]
    b = a[2]
    return torch.stack([r * 0.229 + 0.485, g * 0.224 + 0.456, b * 0.225 + 0.406], dim=0)


import random
import IPython

ds = MyDataset(imgs, tokenizer, aug_train)
for i in range(5):
    d = ds[random.randrange(0, len(ds))]
    IPython.display.display(transforms.ToPILImage()(reverse_transform(d[0])))
    print(tokenizer.untokenize(d[1]))


def one_hot(x, num_classes, on_value=1.0, off_value=0.0, device="cuda"):
    x = x.long().view(-1, 1)
    return torch.full((x.size()[0], num_classes), off_value, device=device).scatter_(
        1, x, on_value
    )


import torch
import torch.nn as nn
import torch.nn.functional as F


# +
class Model(LightningModule):
    def __init__(self, n_classes, lr, wd):
        super().__init__()
        self.save_hyperparameters()
        self.lr = lr
        self.cnn = timm.create_model(MODEL, pretrained=True)  # resnet 50
        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
        self.out = nn.Linear(self.cnn.num_features, n_classes * 4)
        self.crit = nn.CrossEntropyLoss()

    def forward(self, input_):
        x = self.cnn.forward_features(input_)
        x = self.avgpool(x).flatten(1)
        x = self.out(x).view(-1, 4, self.hparams.n_classes)
        return x

    def training_step(self, batch, batch_idx):
        self.train()
        crit = nn.MultiLabelSoftMarginLoss()
        input_tensor, target_tensor = batch
        with torch.no_grad():
            target = one_hot(target_tensor, self.hparams.n_classes).view(
                input_tensor.size(0), -1
            )
        decoder_out = self(input_tensor)  # N x L x C
        loss = crit(decoder_out.flatten(1), target)
        with torch.no_grad():
            target_tensor = target_tensor
            preds = decoder_out.argmax(2)
            acc_char = (preds == target_tensor).float().mean()
            acc_sent = (preds == target_tensor).all(dim=1).float().mean()
            self.log("train_loss", loss)
            self.log("train_acc_char", acc_char)
            self.log("train_acc_sent", acc_sent)
        return loss

    def configure_optimizers(self):
        optimizer = torch.optim.AdamW(
            self.parameters(), lr=0, weight_decay=self.hparams.wd
        )
        scheduler = {
            "scheduler": torch.optim.lr_scheduler.OneCycleLR(
                optimizer,
                max_lr=self.lr,
                total_steps=self.trainer.max_steps,
                anneal_strategy="linear",
                cycle_momentum=False,
                pct_start=0.1,
            ),
            "interval": "step",
            "frequency": 1,
        }
        return [optimizer], [scheduler]

    def validation_step(self, batch, batch_idx):
        self.eval()
        with torch.no_grad():
            input_tensor, target_tensor = batch
            decoder_out = self(input_tensor)  # N x L x C
            loss = self.crit(decoder_out.permute(0, 2, 1), target_tensor)
            target_tensor = target_tensor
            preds = decoder_out.argmax(2)
            acc_char = (preds == target_tensor).float().mean()
            acc_sent = (preds == target_tensor).all(dim=1).float().mean()
            self.log("val_loss", loss)
            self.log("val_acc_char", acc_char)
            self.log("val_acc_sent", acc_sent)


from sklearn.model_selection import train_test_split, KFold
from sklearn.metrics import accuracy_score, f1_score, roc_auc_score

imgs = np.array(imgs)
imgs
from pytorch_lightning.utilities import rank_zero_only
from pytorch_lightning.loggers import LightningLoggerBase
from pytorch_lightning.loggers.base import rank_zero_experiment
from collections import defaultdict


class ConsoleLogger(LightningLoggerBase):
    def __init__(self):
        super().__init__()
        self.metrics = defaultdict(list)

    @property
    def name(self):
        return ""

    @property
    @rank_zero_experiment
    def experiment(self):
        # Return the experiment object associated with this logger.
        pass

    @property
    def version(self):
        # Return the experiment version, int or str.
        return ""

    @rank_zero_only
    def log_hyperparams(self, params):
        # params is an argparse.Namespace
        # your code to record hyperparameters goes here
        pass

    @rank_zero_only
    def log_metrics(self, metrics, step):
        # metrics is a dictionary of metric names and values
        # your code to record metrics goes here
        print("=" * 80)
        if "epoch" in metrics:
            print(f'epoch: {metrics["epoch"]}, steps: {step}')
        else:
            print(f"steps: {step}")
        for k, v in metrics.items():
            if k == "epoch":
                continue
            if isinstance(v, torch.Tensor):
                v = v.item()
            print(f"{k}: {v}")
            self.metrics[k].append((step, v))


import pytorch_lightning

kf = KFold(n_splits=20, random_state=42, shuffle=True)
val_data = []
fold_idx = 0
for train, val in kf.split(imgs):

    def train_func():
        train_imgs = imgs[train]
        val_imgs = imgs[val]
        ds_train = MyDataset(train_imgs, tokenizer, aug_train)
        train_loader = DataLoader(
            ds_train,
            batch_size=batch_size,
            num_workers=4,
            shuffle=True,
            pin_memory=True,
            drop_last=True,
        )
        ds_val = MyDataset(val_imgs, tokenizer, aug_val)
        val_loader = DataLoader(
            ds_val,
            batch_size=batch_size,
            num_workers=4,
            pin_memory=True,
        )
        max_steps = epoch * int(np.ceil(len(train_loader) / gradient_accum))
        print("*" * 80)
        print("training fold", fold_idx)
        print("train samples:", len(train_imgs), "max steps:", max_steps)
        checkpoint_callback = ModelCheckpoint(
            save_top_k=1,
            verbose=True,
            monitor="val_acc_sent",
            mode="max",
            filename="epoch{epoch}-step{step}-val_acc_sent{val_acc_sent:.2f}",
            auto_insert_metric_name=False,
            save_weights_only=True,
            save_last=True,
        )
        logger = ConsoleLogger()
        tb_logger = TensorBoardLogger(
            save_dir=".", version=f"fold_{fold_idx}", name="lightning_logs"
        )
        trainer = Trainer(
            logger=[logger, tb_logger],
            accumulate_grad_batches=gradient_accum,
            gpus=-1,
            callbacks=[LearningRateMonitor(), checkpoint_callback],
            flush_logs_every_n_steps=50,
            log_every_n_steps=50,
            max_steps=max_steps,
            benchmark=True,
            precision=16,
        )
        model = Model(len(vocab), lr, wd)
        #         lr_finder = trainer.tuner.lr_find(model, train_loader)
        #         lr_finder.plot(suggest=True)
        trainer.fit(model, train_dataloaders=train_loader, val_dataloaders=val_loader)

    train_func()
    gc.collect()
    pytorch_lightning.utilities.memory.garbage_collection_cuda()
    break


class MyTestDataset(Dataset):
    def __init__(self, imgs, aug):
        self.path = "../input/dcic-2022-ocr/test_dataset-220406/test_dataset/"
        self.imgs = imgs
        self.tokenizer = tokenizer
        self.aug = aug

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

    def __getitem__(self, idx):
        img = Image.open(os.path.join(self.path, self.imgs[idx])).convert("RGB")
        img = self.aug(img)
        return img


test_imgs = os.listdir("../input/dcic-2022-ocr/test_dataset-220406/test_dataset")
test_imgs = sorted(test_imgs, key=lambda x: int(x.split(".")[0]))
ds_test = MyTestDataset(test_imgs, aug_val)
test_loader = DataLoader(
    ds_test, batch_size=32, num_workers=4, pin_memory=True, shuffle=False
)
model = Model.load_from_checkpoint(
    "./_lightning_logs/_fold_0/checkpoints/last.ckpt"
).cuda()
tmp = []
with torch.no_grad():
    for batch in tqdm(test_loader):
        batch = batch.cuda()
        model.eval()
        pred = model(batch).argmax(2).cpu().numpy()
        tmp.append(pred)
preds = np.concatenate(tmp, axis=0)
df_preds = pd.DataFrame()
df_preds["num"] = [x.split(".")[0] for x in test_imgs]  # [:len(tags)]
df_preds["tag"] = [tokenizer.untokenize(x) for x in preds]
df_preds.to_csv("result.csv", index=None)
df_preds
