# -*- coding: utf-8 -*- # @Time : 2019/8/24 12:06 # @Author : zhoujun import os import sys import pathlib __dir__ = pathlib.Path(os.path.abspath(__file__)) sys.path.append(str(__dir__)) sys.path.append(str(__dir__.parent.parent)) import time import cv2 import paddle from data_loader import get_transforms from models import build_model from post_processing import get_post_processing def resize_image(img, short_size): height, width, _ = img.shape if height < width: new_height = short_size new_width = new_height / height * width else: new_width = short_size new_height = new_width / width * height new_height = int(round(new_height / 32) * 32) new_width = int(round(new_width / 32) * 32) resized_img = cv2.resize(img, (new_width, new_height)) return resized_img class PaddleModel: def __init__(self, model_path, post_p_thre=0.7, gpu_id=None): ''' 初始化模型 :param model_path: 模型地址(可以是模型的参数或者参数和计算图一起保存的文件) :param gpu_id: 在哪一块gpu上运行 ''' self.gpu_id = gpu_id if self.gpu_id is not None and isinstance( self.gpu_id, int) and paddle.device.is_compiled_with_cuda(): paddle.device.set_device("gpu:{}".format(self.gpu_id)) else: paddle.device.set_device("cpu") checkpoint = paddle.load(model_path) config = checkpoint['config'] config['arch']['backbone']['pretrained'] = False self.model = build_model(config['arch']) self.post_process = get_post_processing(config['post_processing']) self.post_process.box_thresh = post_p_thre self.img_mode = config['dataset']['train']['dataset']['args'][ 'img_mode'] self.model.set_state_dict(checkpoint['state_dict']) self.model.eval() self.transform = [] for t in config['dataset']['train']['dataset']['args']['transforms']: if t['type'] in ['ToTensor', 'Normalize']: self.transform.append(t) self.transform = get_transforms(self.transform) def predict(self, img_path: str, is_output_polygon=False, short_size: int=1024): ''' 对传入的图像进行预测，支持图像地址,opecv 读取图片，偏慢 :param img_path: 图像地址 :param is_numpy: :return: ''' assert os.path.exists(img_path), 'file is not exists' img = cv2.imread(img_path, 1 if self.img_mode != 'GRAY' else 0) if self.img_mode == 'RGB': img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB) h, w = img.shape[:2] img = resize_image(img, short_size) # 将图片由(w,h)变为(1,img_channel,h,w) tensor = self.transform(img) tensor = tensor.unsqueeze_(0) batch = {'shape': [(h, w)]} with paddle.no_grad(): start = time.time() preds = self.model(tensor) box_list, score_list = self.post_process( batch, preds, is_output_polygon=is_output_polygon) box_list, score_list = box_list[0], score_list[0] if len(box_list) > 0: if is_output_polygon: idx = [x.sum() > 0 for x in box_list] box_list = [box_list[i] for i, v in enumerate(idx) if v] score_list = [score_list[i] for i, v in enumerate(idx) if v] else: idx = box_list.reshape(box_list.shape[0], -1).sum( axis=1) > 0 # 去掉全为0的框 box_list, score_list = box_list[idx], score_list[idx] else: box_list, score_list = [], [] t = time.time() - start return preds[0, 0, :, :].detach().cpu().numpy(), box_list, score_list, t def save_depoly(net, input, save_path): input_spec = [ paddle.static.InputSpec( shape=[None, 3, None, None], dtype="float32") ] net = paddle.jit.to_static(net, input_spec=input_spec) # save static model for inference directly paddle.jit.save(net, save_path) def init_args(): import argparse parser = argparse.ArgumentParser(description='DBNet.paddle') parser.add_argument('--model_path', default=r'model_best.pth', type=str) parser.add_argument( '--input_folder', default='./test/input', type=str, help='img path for predict') parser.add_argument( '--output_folder', default='./test/output', type=str, help='img path for output') parser.add_argument('--gpu', default=0, type=int, help='gpu for inference') parser.add_argument( '--thre', default=0.3, type=float, help='the thresh of post_processing') parser.add_argument( '--polygon', action='store_true', help='output polygon or box') parser.add_argument('--show', action='store_true', help='show result') parser.add_argument( '--save_result', action='store_true', help='save box and score to txt file') args = parser.parse_args() return args if __name__ == '__main__': import pathlib from tqdm import tqdm import matplotlib.pyplot as plt from utils.util import show_img, draw_bbox, save_result, get_image_file_list args = init_args() print(args) # 初始化网络 model = PaddleModel(args.model_path, post_p_thre=args.thre, gpu_id=args.gpu) img_folder = pathlib.Path(args.input_folder) for img_path in tqdm(get_image_file_list(args.input_folder)): preds, boxes_list, score_list, t = model.predict( img_path, is_output_polygon=args.polygon) img = draw_bbox(cv2.imread(img_path)[:, :, ::-1], boxes_list) if args.show: show_img(preds) show_img(img, title=os.path.basename(img_path)) plt.show() # 保存结果到路径 os.makedirs(args.output_folder, exist_ok=True) img_path = pathlib.Path(img_path) output_path = os.path.join(args.output_folder, img_path.stem + '_result.jpg') pred_path = os.path.join(args.output_folder, img_path.stem + '_pred.jpg') cv2.imwrite(output_path, img[:, :, ::-1]) cv2.imwrite(pred_path, preds * 255) save_result( output_path.replace('_result.jpg', '.txt'), boxes_list, score_list, args.polygon)