|
- #%%
- import os, sys
-
- os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
- os.environ["CUDA_VISIBLE_DEVICES"] = "1"
- import json
- import cv2
- import torch
- import numpy as np
- import matplotlib
- matplotlib.use("Agg")
-
- sys.path.append("../")
- from config import system_configs
- from nnet.py_factory import NetworkFactory
- from db.datasets import datasets
- from test.images import kp_detection_one
-
- torch.backends.cudnn.benchmark = True
-
-
- class Detector():
- def __init__(self) -> None:
- args = {
- 'cfg_file': 'LSTR',
- 'split': 'validation',
- 'testiter':192000,
- 'debug': True,
- 'suffix': None,
- 'modality': 'images',
- # 'image_root':'./',
- 'batch': 216,
- 'debugEnc': False,
- 'debugDec': False,
- }
-
- if args['suffix'] is None:
- cfg_file = os.path.join(system_configs.config_dir, args['cfg_file'] + ".json")
- else:
- cfg_file = os.path.join(system_configs.config_dir, args['cfg_file'] + "-{}.json".format(args['suffix']))
- print("config_file: {}".format(cfg_file))
-
- with open(cfg_file, "r") as f:
- configs = json.load(f)
-
- configs["system"]["snapshot_name"] = args['cfg_file']
- system_configs.update_config(configs["system"])
-
- print("loading all datasets...")
- dataset = system_configs.dataset
-
- testiter = args['testiter']
- test_iter = system_configs.max_iter if testiter is None else testiter
- print("loading parameters at iteration: {}".format(test_iter))
-
- print("building neural network...")
- nnet = NetworkFactory()
-
- print("loading parameters...")
- nnet.load_params(test_iter)
- nnet.cuda()
- nnet.eval_mode()
-
- self.model = nnet
-
- # evaluator = Evaluator(db, result_dir)
- # test_file = "test.images"
- # kp_detection_one = importlib.import_module(test_file).kp_detection_one
-
- testing_db = datasets[dataset](configs["db"], 'val') # /Data/TuSimple/LSTR/db/tusimple.py create an object
- self.testing_db = testing_db
-
-
- def predict(self, img):
- """图片检测推理
-
- Args:
- img (ndarray): 维度为(H,W,C)的图片,通道类型为RGB
-
- Returns:
- [type]: [description]
- """
- assert img is not None, "check img input"
- lanes_info={}
- kp_detection_one(self.testing_db, self.model, img, debug=False, evaluator=None, lanes_info=lanes_info)
-
- return lanes_info
-
-
- def draw_pred_lanes_on_img(self, img, pred_res=None, out_path=None):
- """[summary]
-
- Args:
- img (ndarray): 维度为(H,W,C)的图片,通道类型为RGB
- pred_res (dict): 图片检测推理结果
- out_path (path, optional): . Defaults to None.
-
- Returns:
- [type]: [description]
- """
- assert img is not None, "check img input"
- if pred_res is None:
- pred_res = self.predict(img)
- assert 'pred' in pred_res, "dict not include 'file' or 'pred' "
-
- img_h, img_w, _ = img.shape
- overlay = img.copy()
- color = (0, 255, 0)
-
- for i, lane in enumerate(pred_res['pred']):
- lane = lane[1:] # remove conf
- lower, upper = lane[0], lane[1]
- lane = lane[2:] # remove upper, lower positions
-
- # generate points from the polynomial
- ys = np.linspace(lower, upper, num=50)
- points = np.zeros((len(ys), 2), dtype=np.int32)
- points[:, 1] = (ys * img_h).astype(int)
- points[:, 0] = ((lane[0] / (ys - lane[1]) ** 2 + lane[2] / (ys - lane[1]) + lane[3] + lane[4] * ys - lane[5]) * img_w).astype(int)
- points = points[(points[:, 0] > 0) & (points[:, 0] < img_w)]
-
- # draw lane with a polyline on the overlay
- for current_point, next_point in zip(points[:-1], points[1:]):
- overlay = cv2.line(overlay, tuple(current_point), tuple(next_point), color=color, thickness=15)
-
- # draw lane ID
- if len(points) > 0:
- cv2.putText(img, str(i), tuple(points[len(points)//2]), fontFace=cv2.FONT_HERSHEY_SIMPLEX, fontScale=1,
- color=color,
- thickness=3)
-
- # Add lanes overlay
- w = 0.6
- img = ((1. - w) * img + w * overlay).astype(np.uint8)
-
- if out_path is not None:
- cv2.imwrite(out_path, img)
-
- return img
-
-
- #%%
- if __name__ == '__main__':
-
- fn = './images/demo.jpg'
- img_np = cv2.imread(fn)
-
- det = Detector()
- ans = det.predict(img_np)
- img_ans = det.draw_pred_lanes_on_img(img_np, out_path='./images/demo_res.jpg')
|