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 #8855pull/8930/head
parent
9059f747e6
commit
3629d5a365
|
@ -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,
|
||||
|
|
Loading…
Reference in New Issue