Tryag File Manager

//paip/script/lpr/detectron2_LPR/cnn_lpr_inference.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, Lpr_tensor
from model import Model
from PIL import Image

cudnn.benchmark = True
cudnn.deterministic = True

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

class LicensePlate_Recognition:
    __instance = None

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

if not LicensePlate_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.image_folder = './LPD_test/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 = LicensePlate_Recognition()
        return cls.__instance

def Predict_LPR(self, lpr_result, license_plate_type):
        try:
            if license_plate_type == 3:
                preds_str = ['-']
                
                return preds_str
            else:
                #self.imgH, self.imgW, channel = lpr_result
                lpr_candi = Image.fromarray(lpr_result)
                AlignCollate_demo = Lpr_tensor(keep_ratio_with_pad=self.PAD)

image_tensors, image_path_list = AlignCollate_demo(lpr_candi)
                image_tensors = torch.Tensor(image_tensors)

self.model.eval()
                with torch.no_grad():
                    #batch_size = image_tensors.size(0)
                    batch_size = 1
                    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')
                return preds_str

except Exception as e :
            raise Exception("Error = ", e)

if __name__ == '__main__':
    # opt.num_gpu = torch.cuda.device_count()
    recog = LicensePlate_Recognition.getInstance()
    recog.recog()