# 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)