提交 d83b1de1 编写于 作者: L LDOUBLEV

mv get_gpu_id func to utility

上级 96a3f90c
......@@ -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
......
......@@ -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):
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册