diff --git a/tools/infer/predict_det.py b/tools/infer/predict_det.py index 6a9ae98793be4c6e831bda1b0605e1be7c41973b..62e7ff73964e3c387a19e6f6a1c212ccfdbf8c91 100755 --- a/tools/infer/predict_det.py +++ b/tools/infer/predict_det.py @@ -101,7 +101,7 @@ class TextDetector(object): if args.benchmark: import auto_log pid = os.getpid() - gpu_id = self.get_infer_gpuid() + gpu_id = utility.get_infer_gpuid() self.autolog = auto_log.AutoLogger( model_name="det", model_precision=args.precision, @@ -118,19 +118,6 @@ class TextDetector(object): warmup=2, logger=logger) - def get_infer_gpuid(self): - cmd = "nvidia-smi" - res = os.popen(cmd).readlines() - if len(res) == 0: - return None - cmd = "env | grep CUDA_VISIBLE_DEVICES" - env_cuda = os.popen(cmd).readlines() - if len(env_cuda) == 0: - return 0 - else: - gpu_id = env_cuda[0].strip().split("=")[1] - return int(gpu_id[0]) - def order_points_clockwise(self, pts): """ reference from: https://github.com/jrosebr1/imutils/blob/master/imutils/perspective.py diff --git a/tools/infer/utility.py b/tools/infer/utility.py index 1c82280099f17f6d3bf848669e47439505f10576..527a811d620efac33ece9cdbd4b6196e18a8497d 100755 --- a/tools/infer/utility.py +++ b/tools/infer/utility.py @@ -280,6 +280,20 @@ def create_predictor(args, mode, logger): return predictor, input_tensor, output_tensors, config +def get_infer_gpuid(): + cmd = "nvidia-smi" + res = os.popen(cmd).readlines() + if len(res) == 0: + return None + cmd = "env | grep CUDA_VISIBLE_DEVICES" + env_cuda = os.popen(cmd).readlines() + if len(env_cuda) == 0: + return 0 + else: + gpu_id = env_cuda[0].strip().split("=")[1] + return int(gpu_id[0]) + + def draw_e2e_res(dt_boxes, strs, img_path): src_im = cv2.imread(img_path) for box, str in zip(dt_boxes, strs):