| |
| |
| |
|
|
| 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): |
| """ |
| 对传入的图像进行预测，支持图像地址,opencv 读取图片，偏慢 |
| :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) |
| |
| 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 |
| ) |
| 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) |
|
|
| |
| 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, |
| ) |
|
|
