requirement.txt

  1. sklearn
  2. pandas
  3. efficientnet_pytorch

数据分析

  1. import os
  2. import cv2
  5. img_path = "./dataset/train/images/"
  7. img_list = [img_path+ i for i in os.listdir(img_path)]
  9. shape = []
  10. for img in img_list:
  11.     arr = cv2.imread(img)
  12.     shape.append(arr.shape)

训练

主要修改模型的结构 class VisitNet(nn.Module)

  1. # -*- coding: utf-8 -*-
  2. import os, sys, glob, argparse
  3. import pandas as pd
  4. import numpy as np
  5. from tqdm import tqdm
  7. import time, datetime
  8. import pdb, traceback
  10. import cv2
  11. # import imagehash
  12. from PIL import Image
  14. from sklearn.model_selection import train_test_split, StratifiedKFold, KFold
  16. from efficientnet_pytorch import EfficientNet
  17. # model = EfficientNet.from_pretrained('efficientnet-b4') 
  19. import torch
  20. torch.manual_seed(0)
  21. torch.backends.cudnn.deterministic = False
  22. torch.backends.cudnn.benchmark = True
  24. import torchvision.models as models
  25. import torchvision.transforms as transforms
  26. import torchvision.datasets as datasets
  27. import torch.nn as nn
  28. import torch.nn.functional as F
  29. import torch.optim as optim
  30. from torch.autograd import Variable
  31. from torch.utils.data.dataset import Dataset
  33. # input dataset
  34. """
  35. train_set 目录结构
  36. ├── 0
  37. |    ├── xx.jpg
  38. │   ├── xx.jpg
  39. ├── 1
  40. |    ├── xx.jpg
  41. │   ├── xx.jpg
  42. ├── 2
  43. |    ├── xx.jpg
  44. │   ├── xx.jpg
  45. ├── 3
  46. |    ├── xx.jpg
  47. │   ├── xx.jpg
  49. """
  50. train_jpg = glob.glob('./train_dataset/*/*')
  51. train_jpg = np.array(train_jpg)
  53. print(train_jpg)
  55. if os.path.exists("log.txt"):
  56.     os.remove("log.txt")
  58. class QRDataset(Dataset):
  59.     def __init__(self, img_path, transform=None):
  60.         self.img_path = img_path
  61.         if transform is not None:
  62.             self.transform = transform
  63.         else:
  64.             self.transform = None
  66.     def __getitem__(self, index):
  67.         start_time = time.time()
  68.         img = Image.open(self.img_path[index]).convert('RGB')
  70.         if self.transform is not None:
  71.             img = self.transform(img)
  72.         return img,torch.from_numpy(np.array(int(self.img_path[index].split("/")[-2])))
  74.     def __len__(self):
  75.         return len(self.img_path)
  77. def accuracy(output, target, topk=(1,)):
  78.     """Computes the accuracy over the k top predictions for the specified values of k"""
  79.     with torch.no_grad():
  80.         maxk = max(topk)
  81.         batch_size = target.size(0)
  83.         _, pred = output.topk(maxk, 1, True, True)
  84.         pred = pred.t()
  85.         correct = pred.eq(target.view(1, -1).expand_as(pred))
  87.         res = []
  88.         for k in topk:
  89.             correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
  90.             res.append(correct_k.mul_(100.0 / batch_size))
  91.         return res
  93. class AverageMeter(object):
  94.     """Computes and stores the average and current value"""
  95.     def __init__(self, name, fmt=':f'):
  96.         self.name = name
  97.         self.fmt = fmt
  98.         self.reset()
  100.     def reset(self):
  101.         self.val = 0
  102.         self.avg = 0
  103.         self.sum = 0
  104.         self.count = 0
  106.     def update(self, val, n=1):
  107.         self.val = val
  108.         self.sum += val * n
  109.         self.count += n
  110.         self.avg = self.sum / self.count
  112.     def __str__(self):
  113.         fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'
  114.         return fmtstr.format(**self.__dict__)
  116. class ProgressMeter(object):
  117.     def __init__(self, num_batches, *meters):
  118.         self.batch_fmtstr = self._get_batch_fmtstr(num_batches)
  119.         self.meters = meters
  120.         self.prefix = ""
  123.     def pr2int(self, batch):
  124.         entries = [self.prefix + self.batch_fmtstr.format(batch)]
  125.         entries += [str(meter) for meter in self.meters]
  126.         print('\t'.join(entries))
  127.         with open("log.txt", "a+") as f:
  128.             f.write('\t'.join(entries)+"\n")
  130.     def _get_batch_fmtstr(self, num_batches):
  131.         num_digits = len(str(num_batches // 1))
  132.         fmt = '{:' + str(num_digits) + 'd}'
  133.         return '[' + fmt + '/' + fmt.format(num_batches) + ']'
  135. class VisitNet(nn.Module):
  136.     def __init__(self):
  137.         super(VisitNet, self).__init__()
  139.         model = models.resnet18(True)
  140.         model.avgpool = nn.AdaptiveAvgPool2d(1)
  141.         model.fc = nn.Linear(512, 4)
  142.         self.resnet = model
  144. #         model = EfficientNet.from_pretrained('efficientnet-b4') 
  145. #         model._fc = nn.Linear(1792, 4)
  146. #         self.resnet = model
  148.     def forward(self, img):        
  149.         out = self.resnet(img)
  150.         return out
  152. def validate(val_loader, model, criterion):
  153.     batch_time = AverageMeter('Time', ':6.3f')
  154.     losses = AverageMeter('Loss', ':.4e')
  155.     top1 = AverageMeter('Acc@1', ':6.2f')
  156.     top5 = AverageMeter('Acc@2', ':6.2f')
  157.     progress = ProgressMeter(len(val_loader), batch_time, losses, top1, top5)
  159.     # switch to evaluate mode
  160.     model.eval()
  162.     with torch.no_grad():
  163.         end = time.time()
  164.         for i, (input, target) in enumerate(val_loader):
  165.             input = input.cuda()
  166.             target = target.cuda()
  168.             # compute output
  169.             output = model(input)
  170.             loss = criterion(output, target)
  172.             # measure accuracy and record loss
  173.             acc1, acc5 = accuracy(output, target, topk=(1, 2))
  174.             losses.update(loss.item(), input.size(0))
  175.             top1.update(acc1[0], input.size(0))
  176.             top5.update(acc5[0], input.size(0))
  178.             # measure elapsed time
  179.             batch_time.update(time.time() - end)
  180.             end = time.time()
  182.         # TODO: this should also be done with the ProgressMeter
  183.         print(' * Acc@1 {top1.avg:.3f} Acc@5 {top5.avg:.3f}'
  184.               .format(top1=top1, top5=top5))
  185.         with open("log.txt", "a+") as f:
  186.             f.write(' * Acc@1 {top1.avg:.3f} Acc@5 {top5.avg:.3f}'
  187.               .format(top1=top1, top5=top5)+"\n")
  188.         return top1
  190. def predict(test_loader, model, tta=10):
  191.     # switch to evaluate mode
  192.     model.eval()
  194.     test_pred_tta = None
  195.     for _ in range(tta):
  196.         test_pred = []
  197.         with torch.no_grad():
  198.             end = time.time()
  199.             for i, (input, target) in enumerate(test_loader):
  200.                 input = input.cuda()
  201.                 target = target.cuda()
  203.                 # compute output
  204.                 output = model(input, path)
  205.                 output = output.data.cpu().numpy()
  207.                 test_pred.append(output)
  208.         test_pred = np.vstack(test_pred)
  210.         if test_pred_tta is None:
  211.             test_pred_tta = test_pred
  212.         else:
  213.             test_pred_tta += test_pred
  215.     return test_pred_tta
  217. def train(train_loader, model, criterion, optimizer, epoch):
  218.     batch_time = AverageMeter('Time', ':6.3f')
  219.     # data_time = AverageMeter('Data', ':6.3f')
  220.     losses = AverageMeter('Loss', ':.4e')
  221.     top1 = AverageMeter('Acc@1', ':6.2f')
  222.     # top5 = AverageMeter('Acc@5', ':6.2f')
  223.     progress = ProgressMeter(len(train_loader), batch_time, losses, top1)
  225.     # switch to train mode
  226.     model.train()
  228.     end = time.time()
  229.     for i, (input, target) in enumerate(train_loader):
  230.         input = input.cuda(non_blocking=True)
  231.         target = target.cuda(non_blocking=True)
  233.         # compute output
  234.         output = model(input)
  235.         loss = criterion(output, target)
  237.         # measure accuracy and record loss
  238.         acc1, acc5 = accuracy(output, target, topk=(1, 2))
  239.         losses.update(loss.item(), input.size(0))
  240.         top1.update(acc1[0], input.size(0))
  241.         # top5.update(acc5[0], input.size(0))
  243.         # compute gradient and do SGD step
  244.         optimizer.zero_grad()
  245.         loss.backward()
  246.         optimizer.step()
  248.         # measure elapsed time
  249.         batch_time.update(time.time() - end)
  250.         end = time.time()
  252.         if i % 100 == 0:
  253.             progress.pr2int(i)
  255. skf = KFold(n_splits=2, random_state=233, shuffle=True)
  256. for flod_idx, (train_idx, val_idx) in enumerate(skf.split(train_jpg, train_jpg)):
  257.     print(flod_idx, train_idx, val_idx)
  259.     train_loader = torch.utils.data.DataLoader(
  260.         QRDataset(train_jpg[train_idx],
  261.                 transforms.Compose([
  262.                             # transforms.RandomGrayscale(),
  263.                             transforms.Resize((512, 512)),
  264.                             transforms.RandomAffine(10),
  265.                             transforms.ColorJitter(hue=.05, saturation=.05),
  266.                             transforms.RandomCrop((224, 224)),
  267.                             transforms.RandomHorizontalFlip(),
  268.                             transforms.RandomVerticalFlip(),
  269.                             transforms.ToTensor(),
  270.                             transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
  271.             ])
  272.         ), batch_size=2, shuffle=True, num_workers=8, pin_memory=True
  273.     )
  274.     val_loader = torch.utils.data.DataLoader(
  275.         QRDataset(train_jpg[val_idx],
  276.                 transforms.Compose([
  277.                             transforms.Resize((512, 512)),
  278.                             # transforms.Resize((124, 124)),
  279.                             # transforms.RandomCrop((88, 88)),
  280.                             transforms.ToTensor(),
  281.                             transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
  282.             ])
  283.         ), batch_size=1, shuffle=False, num_workers=8, pin_memory=True
  284.     )
  286.     model = VisitNet().cuda()
  287.     # model = nn.DataParallel(model).cuda()
  288.     criterion = nn.CrossEntropyLoss().cuda()
  289.     optimizer = torch.optim.Adam(model.parameters())
  290.     scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=25, gamma=0.85)
  291.     best_acc = 0.0
  292.     for epoch in range(120):
  293.         scheduler.step()
  294.         print('Epoch: ', epoch)
  295.         with open("log.txt", "a+") as f:
  296.             f.write('Epoch: '+str(epoch)+"\n")
  297.         train(train_loader, model, criterion, optimizer, epoch)
  298.         val_acc = validate(val_loader, model, criterion)
  300.         if val_acc.avg.item() > best_acc:
  301.             best_acc = val_acc.avg.item()
  302.             torch.save(model.state_dict(), './efficientnet{0}.pt'.format(flod_idx))

预测

# -*- coding: utf-8 -*-
import os, sys, glob, argparse
import pandas as pd
import numpy as np
from tqdm import tqdm

import time, datetime
import pdb, traceback

import cv2
# import imagehash
from PIL import Image

from sklearn.model_selection import train_test_split, StratifiedKFold, KFold

from efficientnet_pytorch import EfficientNet
# model = EfficientNet.from_pretrained('efficientnet-b4') 

import torch
torch.manual_seed(0)
torch.backends.cudnn.deterministic = False
torch.backends.cudnn.benchmark = True

import torchvision.models as models
import torchvision.transforms as transforms
import torchvision.datasets as datasets
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.autograd import Variable
from torch.utils.data.dataset import Dataset

class QRDataset(Dataset):
    def __init__(self, img_path, transform=None):
        self.img_path = img_path
        if transform is not None:
            self.transform = transform
        else:
            self.transform = None

    def __getitem__(self, index):
        start_time = time.time()
        img = Image.open(self.img_path[index]).convert('RGB')

        if self.transform is not None:
            img = self.transform(img)
        return img, torch.from_numpy(np.array([1]))#,torch.from_numpy(np.array(int(self.img_path[index].split("/")[-2])))

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

class VisitNet(nn.Module):
    def __init__(self):
        super(VisitNet, self).__init__()

#         model = models.resnet18(True)
#         model.avgpool = nn.AdaptiveAvgPool2d(1)
#         model.fc = nn.Linear(512, 4)
#         self.resnet = model

        model = EfficientNet.from_pretrained('efficientnet-b4') 
        model._fc = nn.Linear(1792, 4)
        self.resnet = model

    def forward(self, img):        
        out = self.resnet(img)
        return out

def predict(test_loader, model, tta=10):
    # switch to evaluate mode
    model.eval()

    test_pred_tta = None
    for _ in range(tta):
        test_pred = []
        with torch.no_grad():
            end = time.time()
            for i, (input, target) in enumerate(test_loader):
                input = input.cuda()
                target = target.cuda()

                # compute output
                output = model(input)
                output = output.data.cpu().numpy()

                test_pred.append(output)
        test_pred = np.vstack(test_pred)

        if test_pred_tta is None:
            test_pred_tta = test_pred
        else:
            test_pred_tta += test_pred

    return test_pred_tta

upload_csv_file = "./dataset/test/upload.csv"
result = pd.read_csv(upload_csv_file) 
test_jpg = list(result.iloc[:, 0].values)
test_jpg = list(map(lambda x:"./dataset/test/"+x, test_jpg))
test_jpg = np.array(test_jpg)


test_pred = None
for model_path in ["efficientnet0.pt", "efficientnet1.pt"]:#['resnet18_fold{}.pt'.format(i) for i in [0,2,3,4,6]]:
    print(model_path)
    test_loader = torch.utils.data.DataLoader(
        QRDataset(test_jpg,
                transforms.Compose([
                            transforms.Resize((224, 224)),
#                             transforms.RandomHorizontalFlip(),
#                             transforms.RandomVerticalFlip(),
                            transforms.ToTensor(),
                            transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
            ])
        ), batch_size=64, shuffle=False, num_workers=16, pin_memory=True
    )


    model = VisitNet().cuda()
    model.load_state_dict(torch.load(model_path))
    # model = nn.DataParallel(model).cuda()
    if test_pred is None:
        test_pred = predict(test_loader, model, 5)
    else:
        test_pred += predict(test_loader, model, 5)


test_csv = pd.DataFrame()
test_csv[0] = list(result.iloc[:, 0].values)
test_csv[1] = np.argmax(test_pred, 1)
test_csv.to_csv('tmp.csv', index=None, header=None)