Tryag File Manager

//paip/script/lpr/detectron2_LPR/recog_demo.py

import string
import argparse

import torch
import torch.backends.cudnn as cudnn
import torch.utils.data
import torch.nn.functional as F

from utils import CTCLabelConverter, AttnLabelConverter
from dataset import RawDataset, AlignCollate
from model import Model
# import os
import time

cudnn.benchmark = True
cudnn.deterministic = True

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
#device = 'cpu'
class Recognition:
    __instance = None

def __init__(self, ):
        # parser = argparse.Namespace()

if not Recognition.__instance:
            self.eval_data = True
            self.benchmark_all_eval = 'store_true'
            self.workers = 4
            self.batch_size = 192
            self.saved_model = './weight/lpr_weight/TPS-ResNet-None-CTC-Seed504-9599.pth'
            # self.saved_model = './weight/lpr_weight/TPS-ResNet-None-CTC_Best.pth'
            self.image_folder = './test_datatset/save_crop'

""" Data processing """
            self.batch_max_length = 25
            self.imgH = 32
            self.imgW = 100
            self.rgb = False
            self.character = '0123456789가나다라마거너더러머버서어저고노도로모보소오조구누두루무부수우주허하호강원경기남북광주대전산서울세종울인천제충바사아자배'
            self.sensitive = 'store_true'
            self.PAD = 'store_true'

""" Model Architecture """
            self.Transformation = 'TPS'
            self.FeatureExtraction = 'ResNet'
            self.SequenceModeling = 'None'
            self.Prediction = 'CTC'
            self.num_fiducial = 20
            self.input_channel = 1
            self.output_channel = 512
            self.hidden_size = 256
            self.num_gpu = 0

# opt = parser

""" model configuration """
            if 'CTC' in self.Prediction:
                self.converter = CTCLabelConverter(self.character)
            else:
                self.converter = AttnLabelConverter(self.character)
            self.num_class = len(self.converter.character)

if self.rgb:
               self.input_channel = 3

model = Model(self)
            #print('model input parameters', opt.imgH, opt.imgW, opt.num_fiducial, opt.input_channel, opt.output_channel,
            #      opt.hidden_size, opt.num_class, opt.batch_max_length, opt.Transformation, opt.FeatureExtraction,
            #      opt.SequenceModeling, opt.Prediction)
            self.model = torch.nn.DataParallel(model).to(device)

# load model
            print('loading pretrained model from %s' % self.saved_model)
            self.model.load_state_dict(torch.load(self.saved_model, map_location=device))
        else:
            # 인스턴스가 이미 존재 할 때 구현
            self.getInstance()

@classmethod
    def getInstance(cls):
        if not cls.__instance:
            cls.__instance = Recognition()
        return cls.__instance

def recog(self):
        # prepare data. two demo images from https://github.com/bgshih/crnn#run-demo
        AlignCollate_demo = AlignCollate(imgH=self.imgH, imgW=self.imgW, keep_ratio_with_pad=self.PAD)
        demo_data = RawDataset(root=self.image_folder, opt=self)  # use RawDataset
        demo_loader = torch.utils.data.DataLoader(
            demo_data, batch_size=self.batch_size,
            shuffle=False,
            num_workers=int(self.workers),
            collate_fn=AlignCollate_demo, pin_memory=True)

# predict
        self.model.eval()
        with torch.no_grad():
            for image_tensors, image_path_list in demo_loader:
                print(image_tensors.shape)
                batch_size = image_tensors.size(0)
                image = image_tensors.to(device)
                # For max length prediction
                length_for_pred = torch.IntTensor([self.batch_max_length] * batch_size).to(device)
                text_for_pred = torch.LongTensor(batch_size, self.batch_max_length + 1).fill_(0).to(device)

if 'CTC' in self.Prediction:
                    preds = self.model(image, text_for_pred)

# Select max probabilty (greedy decoding) then decode index to character
                    preds_size = torch.IntTensor([preds.size(1)] * batch_size)
                    _, preds_index = preds.max(2)
                    # preds_index = preds_index.view(-1)
                    preds_str = self.converter.decode(preds_index, preds_size)

else:
                    preds = self.model(image, text_for_pred, is_train=False)

# select max probabilty (greedy decoding) then decode index to character
                    _, preds_index = preds.max(2)
                    preds_str = self.converter.decode(preds_index, length_for_pred)

# log = open(f'./log_demo_result.txt', 'a')
                dashed_line = '-' * 80
                head = f'{"image_path":25s}\t{"predicted_labels":25s}\tconfidence score'

print(f'{dashed_line}\n{head}\n{dashed_line}')
                # log.write(f'{dashed_line}\n{head}\n{dashed_line}\n')

preds_prob = F.softmax(preds, dim=2)
                preds_max_prob, _ = preds_prob.max(dim=2)

for img_name, pred, pred_max_prob in zip(image_path_list, preds_str, preds_max_prob):
                    if 'Attn' in self.Prediction:
                        pred_EOS = pred.find('[s]')
                        pred = pred[:pred_EOS]  # prune after "end of sentence" token ([s])
                        pred_max_prob = pred_max_prob[:pred_EOS]

# calculate confidence score (= multiply of pred_max_prob)
                    confidence_score = pred_max_prob.cumprod(dim=0)[-1]

print(f'{img_name:25s}\t{pred:25s}\t{confidence_score:0.4f}')
                    # log.write(f'{img_name:25s}\t{pred:25s}\t{confidence_score:0.4f}\n')

# log.close()

if __name__ == '__main__':

# opt.num_gpu = torch.cuda.device_count()
    recog = Recognition.getInstance()
    recog.recog()