未验证 提交 3629d5a3 编写于 作者: D Double_V 提交者: GitHub

support set gpu_id when inference (#8926)

* support min_area_rect crop

* add check_install

* fix requirement.txt

* fix check_install

* add lanms-neo for drrg

* fix

* fix doc

* fix

* support set gpu_id when inference

* fix #8855

* fix #8855
上级 9059f747
......@@ -29,7 +29,7 @@ def gen_rec_label(input_path, out_label):
def gen_det_label(root_path, input_dir, out_label):
with open(out_label, 'w') as out_file:
for label_file in os.listdir(input_dir):
img_path = root_path + label_file[3:-4] + ".jpg"
img_path = os.path.join(root_path, label_file[3:-4] + ".jpg")
label = []
with open(
os.path.join(input_dir, label_file), 'r',
......
......@@ -42,6 +42,7 @@ def init_args():
parser.add_argument("--min_subgraph_size", type=int, default=15)
parser.add_argument("--precision", type=str, default="fp32")
parser.add_argument("--gpu_mem", type=int, default=500)
parser.add_argument("--gpu_id", type=int, default=0)
# params for text detector
parser.add_argument("--image_dir", type=str)
......@@ -219,7 +220,7 @@ def create_predictor(args, mode, logger):
logger.warning(
"GPU is not found in current device by nvidia-smi. Please check your device or ignore it if run on jetson."
)
config.enable_use_gpu(args.gpu_mem, 0)
config.enable_use_gpu(args.gpu_mem, args.gpu_id)
if args.use_tensorrt:
config.enable_tensorrt_engine(
workspace_size=1 << 30,
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册