text_plain_output_1.png
from torch.utils.data.dataset import Dataset

import cv2

import os

import timm

import torch

conf = {'batch': 16, 'image_dir': '../input/dog-image-dsg/photo/photo', 'image_size': 224, 'tta': 1, 'num_classes': 73, 'num_workers': 2, 'device': 'cuda' if torch.cuda.is_available() else 'cpu'}

modeldef = [{'mdl': 'convnext_base_in22ft1k', 'pth': '../input/d/sinpcw/dog-image-dsg/ex001/ex001/k0/model_best.pth'}, {'mdl': 'convnext_base_in22ft1k', 'pth': '../input/d/sinpcw/dog-image-dsg/ex001/ex001/k1/model_best.pth'}, {'mdl': 'convnext_base_in22ft1k', 'pth': '../input/d/sinpcw/dog-image-dsg/ex001/ex001/k2/model_best.pth'}, {'mdl': 'convnext_base_in22ft1k', 'pth': '../input/d/sinpcw/dog-image-dsg/ex001/ex001/k3/model_best.pth'}, {'mdl': 'convnext_base_in22ft1k', 'pth': '../input/d/sinpcw/dog-image-dsg/ex001/ex001/k4/model_best.pth'}]

def TTA(x, ops):
    if ops == 0:
        y = x
    elif ops == 1:
        y = torch.flip(x, [-1])
    elif ops == 2:
        y = torch.flip(x, [-2])
    elif ops == 3:
        y = torch.flip(x, [-1, -2])
    else:
        raise ValueError()
    return y

class InferDataset(Dataset):
    def __init__(self, image_dir, dataframe, augmentop):
        self.image_dir = image_dir
        self.dataframe = dataframe
        self.augmentop = augmentop
    def __len__(self):
        return len(self.dataframe)
    def __getitem__(self, idx):
        img = cv2.imread(os.path.join(self.image_dir, self.dataframe.iat[idx, 0] + '.jpg'))
        img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img = self.augmentop(force_apply=False, image=img)['image']
        img = img.transpose(2, 0, 1)
        img = torch.from_numpy(img)
        return img

def GetLoader(image_dir, dataframe, augmentop, batch=1, num_workers=2):
    return torch.utils.data.DataLoader(InferDataset(image_dir, dataframe, augmentop), batch_size=batch, shuffle=False, drop_last=False, num_workers=num_workers)

def GetModel(name, num_classes, pth):
    model = timm.create_model(model_name=name, num_classes=num_classes, in_chans=3, pretrained=False)
    state = torch.load(pth, map_location='cpu')
    model.load_state_dict(state, strict=True)
    model.eval()
    return model

def GetModels(config, mdefs):
    models = []
    for i, mdef in enumerate(mdefs):
        mdl = GetModel(mdef['mdl'], config['num_classes'], mdef['pth']).to(conf['device'])
        models.append(mdl)
    return models

infer_models = GetModels(conf, modeldef)