shufa_category.py

# -*- encoding: utf-8 -*-
__author__ = 'Euphoria'

import os
import sys

import argparse
import time
from PIL import Image
import numpy as np

import torch
from torch import nn
from torch.utils.data import TensorDataset, DataLoader, Dataset
from torchvision import transforms
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score, precision_score, recall_score, f1_score

from model.discriminators import Discriminator


inv_writer_dict = {
        '智永': 0, ' 隸書-趙之謙': 1, '張即之': 2, '張猛龍碑': 3, '柳公權': 4, '標楷體-手寫': 5, '歐陽詢-九成宮': 6,
        '歐陽詢-皇甫誕': 7, '沈尹默': 8, '美工-崩雲體': 9, '美工-瘦顏體': 10, '虞世南': 11, '行書-傅山': 12, '行書-王壯為': 13,
        '行書-王鐸': 14, '行書-米芾': 15, '行書-趙孟頫': 16, '行書-鄭板橋': 17, '行書-集字聖教序': 18, '褚遂良': 19, '趙之謙': 20,
        '趙孟頫三門記體': 21, '隸書-伊秉綬': 22, '隸書-何紹基': 23, '隸書-鄧石如': 24, '隸書-金農': 25,  '顏真卿-顏勤禮碑': 26,
        '顏真卿多寶塔體': 27, '魏碑': 28
    }

writer_dict = {v: k for k, v in inv_writer_dict.items()}


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument('--action', choices=['train', 'eval'])

    parser.add_argument('--image_size', type=int, default=256, help="size of your input images")
    parser.add_argument('--embedding_num', type=int, default=30, help="number for distinct embeddings")
    parser.add_argument('--batch_size', type=int, default=32, help='number of examples in batch')
    parser.add_argument('--input_nc', type=int, default=1)
    parser.add_argument('--input_path', type=str, required=True)

    # train
    parser.add_argument('--resume', type=str, default=None)
    parser.add_argument('--lr', type=float, default=3e-4)
    parser.add_argument('--epoch', type=int, default=10)
    parser.add_argument('--save_path', type=str, default='output')

    # eval
    parser.add_argument('--ckpt_path', type=str, default='output/category.pth')

    args = parser.parse_args()
    return args


def collate_fn_val(batch):
    total_img = []
    for img_path in batch:
        img = Image.open(img_path).convert('L').resize((args.image_size, args.image_size), Image.BICUBIC)
        img = transforms.ToTensor()(img)
        img = transforms.Normalize(0.5, 0.5)(img)
        img = img.unsqueeze(dim=0)
        total_img.append(img)
    total_img = torch.cat(total_img, dim=0)
    return total_img


def collate_fn_train(batch):
    total_img = []
    total_font = []
    for img_path, font in batch:
        img = Image.open(img_path).convert('L').resize((args.image_size, args.image_size), Image.BICUBIC)
        img = transforms.ToTensor()(img)
        img = transforms.Normalize(0.5, 0.5)(img)
        img = img.unsqueeze(dim=0)
        total_img.append(img)
        total_font.append(font)
    total_img = torch.cat(total_img, dim=0)
    total_font = torch.LongTensor(total_font)
    return total_img, total_font


class ImgDataset(Dataset):
    def __init__(self, img_list, label_list=None):
        super(ImgDataset).__init__()
        self.img_list = img_list
        if label_list is not None:
            self.label_list = label_list
            assert len(img_list) == len(label_list)
        else:
            self.label_list = None

    def __getitem__(self, index):
        if self.label_list is not None:
            return self.img_list[index], self.label_list[index]
        else:
            return self.img_list[index]

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



def load_val_dataloader(args):
    IMG_EXT = {'.jpg', '.png', '.tif', '.tiff'}
    raw_img_list = []
    for root, dirs, files in os.walk(args.input_path):
        for file in files:
            if os.path.splitext(file)[-1].lower() in IMG_EXT:
                raw_img_list.append(os.path.join(root, file))
    img_list = raw_img_list
    # img_list = torch.repeat_interleave(img_list, repeats=2, dim=1)
    img_dataset = ImgDataset(img_list)
    img_dataloader = DataLoader(img_dataset, batch_size=args.batch_size, shuffle=False, collate_fn=collate_fn_val)
    return raw_img_list, img_dataloader


def eval(args):
    font_names = [s.strip() for s in open(args.type_file, encoding='utf-8').readlines()]
    font_map = {idx: font for idx, font in enumerate(font_names)}

    model = Discriminator(input_nc=args.input_nc, embedding_num=args.embedding_num)
    model_ckpt = torch.load(os.path.join(args.ckpt_path))
    model.load_state_dict(model_ckpt)
    print('load model {}'.format(args.ckpt_path))

    model.to('cuda')

    raw_image_list, img_dataloader = load_val_dataloader(args)

    total_category = []

    for batch in img_dataloader:
        img = batch
        img = img.to('cuda')
        _, catagory_logits = model(img)
        catagory_logits = catagory_logits.detach().cpu()
        catagory_idx = torch.argmax(catagory_logits, dim=-1)
        catagory_idx = catagory_idx.numpy().tolist()
        total_category.extend(catagory_idx)

    for img, catagory_idx in zip(raw_image_list, total_category):
        print('{}: {}'.format(img, font_map[catagory_idx]))


def load_train_dataloader(args, inv_font_map):
    IMG_EXT = {'.jpg', '.png', '.tif', '.tiff'}
    raw_img_list = []
    raw_label_list = []
    for root, dirs, files in os.walk(args.input_path):
        for file in files:
            if os.path.splitext(file)[-1].lower() in IMG_EXT:
                raw_img_list.append(os.path.join(root, file))
                raw_label_list.append(file.split('~')[1])

    img_list = [img_path for img_path in raw_img_list]
    label_list = [inv_font_map[font_name] for font_name in raw_label_list]

    img_train, img_val, label_train, label_val = train_test_split(img_list, label_list, test_size=0.1, random_state=777)

    print('get {} train examples, {} val examples.'.format(len(img_train), len(img_val)))

    # img_train = torch.cat(img_train, dim=0)
    # label_train = torch.LongTensor(label_train)
    train_dataset = ImgDataset(img_train, label_train)
    train_dataloader = DataLoader(train_dataset, batch_size=args.batch_size, shuffle=True, collate_fn=collate_fn_train)

    # img_val = torch.cat(img_val, dim=0)
    # label_val = torch.LongTensor(label_val)
    val_dataset = ImgDataset(img_val, label_val)
    val_dataloader = DataLoader(val_dataset, batch_size=args.batch_size, shuffle=False, collate_fn=collate_fn_train)

    return train_dataloader, val_dataloader


def train(args):

    font_map = writer_dict
    inv_font_map = inv_writer_dict

    if not os.path.isdir(args.save_path):
        os.makedirs(args.save_path)

    model = Discriminator(input_nc=args.input_nc, embedding_num=args.embedding_num)

    if args.resume is not None:
        model_ckpt = torch.load(args.resume)
        try:
            model.load_state_dict(model_ckpt)
        except RuntimeError:
            print('Guess resume from raw discriminator ckpt. Try pop model.0.weight.')
            model_ckpt.pop('model.0.weight')
            try:
                model.load_state_dict(model_ckpt, strict=False)
            except RuntimeError:
                print('Guess resume ckpt and your model have different catagories. Try pop catagory parameters.')
                model_ckpt.pop('catagory.weight')
                model_ckpt.pop('catagory.bias')
                model.load_state_dict(model_ckpt, strict=False)
        print('load model {}'.format(args.resume), flush=True)

    model.to('cuda')
    optimizer = torch.optim.AdamW(params=model.parameters(), lr=args.lr)
    train_dataloader, val_dataloader = load_train_dataloader(args, inv_font_map)

    best_acc = 0
    best_f1 = 0

    loss_criterion = nn.CrossEntropyLoss()

    for epoch_idx in range(args.epoch):
        losses = []
        for batch_idx, batch in enumerate(train_dataloader):
            img, label = batch
            img = img.to('cuda')
            label = label.to('cuda')
            model.zero_grad()
            _, catagory_logits = model(img)
            loss = loss_criterion(catagory_logits, label)
            losses.append(loss.item())
            loss.backward()
            optimizer.step()
            if batch_idx % 500 == 0:
                print('Epoch: {}, Batch: {}, Loss:{:.4f}'.format(epoch_idx, batch_idx, loss.item()), flush=True)
        print('Epoch: {}, Loss:{:.4f}'.format(epoch_idx, np.mean(losses)), flush=True)
        with torch.no_grad():
            model.eval()
            pred = []
            gold = []
            for batch_idx, batch in enumerate(val_dataloader):
                img, label = batch
                img = img.to('cuda')
                gold.extend(label.numpy().tolist())
                _, catagory_logits = model(img)
                catagory_idx = torch.argmax(catagory_logits, dim=-1)
                catagory_idx = catagory_idx.detach().cpu().numpy().tolist()
                pred.extend(catagory_idx)
            acc = accuracy_score(gold, pred)
            pre = precision_score(gold, pred, average='macro')
            rec = recall_score(gold, pred, average='macro')
            f1 = f1_score(gold, pred, average='macro')
            print('Epoch: {}, p: {:.4f}, r: {:.4f}, f1: {:.4f} ACC: {:.4f}'.format(epoch_idx, pre, rec, f1, acc), flush=True)
            if f1 >= best_f1 or acc >= best_acc:
                print('Save best ckpt.', flush=True)
                torch.save(model.state_dict(), os.path.join(args.save_path, 'category_best.pth'))
                best_acc = acc
                best_f1 = f1
            model.train()


if __name__ == '__main__':
    args = parse_args()
    if args.action == 'eval':
        eval(args)
    if args.action == 'train':
        train(args)