MPKFNet/DclTest.py

# coding=utf-8

'''
Input:(Device, ImagePath, ResultVisualationPath)
DclNumberToSpeciesPath:we need to convert number to related species,so need to find a train.txt to relieve this operation
DclNumberToSpeciesPath:../data_MPKFNet/data/DCL/datasets/HCData/anno/train.txt
ImagePath:../data_MPKFNet/data/DCL/datasets/HCData/data
ResultVisualationPath:../data_MPKFNet/result/Dcl_FineGrainedRecognition

How to use:
Dcl = DclInference(devices, args.ImagePath, args.ResultVisualation)
Dcl.NumbwerToSpecies()
Dcl.parse_args()
Species,  ConfidenceLevel= Dcl.DclInference()

'''
import os
import json
import time
import cv2
import sys
import csv
import pickle
import shutil
import argparse
import pandas as pd
import numpy as np
from math import ceil
from tqdm import tqdm
import PIL.Image as Image
import matplotlib.pyplot as plt

import torch
import torch.nn as nn
from torch.autograd import Variable
from torch.nn import CrossEntropyLoss
from torchvision import datasets, models
import torch.backends.cudnn as cudnn

# Add the absolute path to find the package
#sys.path.append('/data1/Hatcher/model/MPKFNet/Dcl')
from Dcl.transforms import transforms
from Dcl.models.LoadModel import MainModel
from Dcl.DclConfig import LoadConfig, load_data_transformers
from Dcl.util.test_tool import set_text, save_multi_img, cls_base_acc

#os.environ['CUDA_VISIBLE_DEVICES'] = '6'
devices = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

class DclInference:
    def __init__(self, device, DclImagePath, DclResultVisualationPath):
        self.device = device
        self.ImagesPath = []
        self.DclOutputSpecies = []
        self.DclOutputConfidenceLevel = []
        self.DclImagePath = DclImagePath
        self.DataSpeciesClassier = dict()
        self.DclResultVisualationPath = DclResultVisualationPath
        self.DclNumberToSpeciesPath = '../data_MPKFNet/data/DCL/datasets/HCData/anno/train.txt'
        self.DclPreWeightPath = '../data_MPKFNet/weigth/FinedGrainedRecognition/Test/weights_170_999_0.8298_0.9289.pth'

    def DclSoftmax(DcloutputsMatrix, DclPredictedValue):

        InputMax, InputIdx = DcloutputsMatrix.max(dim=1)  # row max
        InputExp = torch.exp(DclPredictedValue - InputMax)
        ConfidenceLevel = InputExp / torch.sum(torch.exp(DcloutputsMatrix - InputMax))
        ConfidenceLevel = ConfidenceLevel.item()
        return ConfidenceLevel

    def NumbwerToSpecies(self):
        Number_last = 0
        with open(self.DclNumberToSpeciesPath) as file:
            Data = file.readlines()
            for line in Data:
                line = line.strip('\n')
                species = line.split(' ')[0]
                Number = line.split(' ')[1]
                # Judge Number
                if Number != Number_last:
                    self.DataSpeciesClassier[Number] = str(species.split('_')[0]) + '_' + str(species.split('_')[1])
                Number_last = Number

    def parse_args(self):
        parser = argparse.ArgumentParser(description='dcl parameters')
        parser.add_argument('--data', dest='dataset', default='HCData', type=str)
        parser.add_argument('--backbone', dest='backbone', default='resnet50', type=str)
        parser.add_argument('--b', dest='batch_size', default=1, type=int)
        parser.add_argument('--nw', dest='num_workers', default=8, type=int)
        parser.add_argument('--ver', dest='version', default='val', type=str)
        parser.add_argument('--size', dest='resize_resolution', default=512, type=int)
        parser.add_argument('--crop', dest='crop_resolution', default=448, type=int)
        parser.add_argument('--ss', dest='save_suffix', default=None, type=str)
        parser.add_argument('--acc_report', dest='acc_report', action='store_true')
        parser.add_argument('--swap_num', default=[7, 7], nargs=2, metavar=('swap1', 'swap2'), type=int, help='specify a range')
        self.args = parser.parse_args()

    def DclInference(self):

        # Traverse Images
        if os.path.exists(self.DclImagePath):
            for File in os.listdir(self.DclImagePath):
                if File.endswith('jpg'):
                    self.ImagesPath.append(os.path.join(self.DclImagePath, File))
        print('Load data is ok')

        # Data Pretreated
        self.Config = LoadConfig(self.args, 's')
        transformers = load_data_transformers(self.args.resize_resolution, self.args.crop_resolution, self.args.swap_num)

        # load model
        self.model = MainModel(self.Config)
        model_dict = self.model.state_dict()
        pretrained_dict = torch.load(self.DclPreWeightPath)
        model_dict.update(pretrained_dict)
        self.model.load_state_dict(model_dict)
        print("load model weight is ok")

        self.model.to(self.device)
        self.model.train(False)
        count_bar = tqdm(total = len(self.ImagesPath), desc='Dcl Inference')
        with torch.no_grad():
            for ImgPath in self.ImagesPath:
                count_bar.update(1)
                Img = cv2.imread(ImgPath)
                Img = Image.fromarray(cv2.cvtColor(Img, cv2.COLOR_BGR2RGB))
                Img = transformers['test_totensor'](Img)
                Img = Img.unsqueeze(0)
                inputs = Variable(Img.cuda())
                outputs = self.model(inputs)
                t1 = time.time()
                top3_val, top3_pos = torch.topk(outputs[0], 1)  # topk返回的是一个元组，第一个元素是每一行最大的三个值（默认是按行取最大值），第二个元素是每一行最大的值的排序
                Species = top3_pos.squeeze()
                ConfidenceLevel = DclInference.DclSoftmax(outputs[0], top3_val[0][0])
                t2 = time.time()
                self.DclOutputSpecies.append(Species)
                self.DclOutputConfidenceLevel.append(ConfidenceLevel)
                print('Post Process Time:{:.2f} ConfidenceLevel:{:.2f}'.format((t2 - t1) * 1000, ConfidenceLevel))
            return self.DclOutputSpecies, self.DclOutputConfidenceLevel

def parse_args():
    parser = argparse.ArgumentParser(description='dcl parameters')
    parser.add_argument('--save', dest='resume',
                        default='../data_MPKFNet/weigth/FinedGrainedRecognition/Test/weights_170_999_0.8298_0.9289.pth', type=str)
    parser.add_argument('--ResultVisualation',
                        default='../data_MPKFNet/result/Dcl_FineGrainedRecognition', type=str)
    parser.add_argument('--ImagePath',
                        default='../data_MPKFNet/data/DCL/datasets/HCData/data', type=str)
    parser.add_argument('--NumberToSpecies',
                        default='../data_MPKFNet/data/DCL/datasets/HCData/anno/train.txt', type=str)
    args = parser.parse_args()
    return args

if __name__ == '__main__':
    args = parse_args()
    Dcl = DclInference(devices, args.ImagePath, args.ResultVisualation)
    Dcl.NumbwerToSpecies()
    Dcl.parse_args()
    Species, ConfidenceLevel = Dcl.DclInference()
    print(Species)