mv get_gpu_id func to utility
This commit is contained in:
parent
96a3f90ced
commit
d83b1de17e
|
@ -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):
|
||||
|
|
Loading…
Reference in New Issue