You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
176 lines
6.3 KiB
Python
176 lines
6.3 KiB
Python
# -*- 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="Crop_img path for predict"
|
|
)
|
|
parser.add_argument(
|
|
"--output_folder", default="./test/json", type=str, help="Crop_img path for json"
|
|
)
|
|
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="json 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,
|
|
)
|