"vscode:/vscode.git/clone" did not exist on "74aaea6f2be01631942231ef936bf6bcb17a8e45"
Unverified Commit 697da27c authored by Xiaomeng Zhao's avatar Xiaomeng Zhao Committed by GitHub
Browse files

Merge pull request #2144 from myhloli/dev

fix(table): add model path for slanet-plus to resolve RapidTableError
parents 99d5c022 e327e9ba
import os
from pathlib import Path
import cv2 import cv2
import numpy as np import numpy as np
import torch import torch
...@@ -17,7 +19,9 @@ class RapidTableModel(object): ...@@ -17,7 +19,9 @@ class RapidTableModel(object):
if torch.cuda.is_available() and table_sub_model_name == "unitable": if torch.cuda.is_available() and table_sub_model_name == "unitable":
input_args = RapidTableInput(model_type=table_sub_model_name, use_cuda=True, device=get_device()) input_args = RapidTableInput(model_type=table_sub_model_name, use_cuda=True, device=get_device())
else: else:
input_args = RapidTableInput(model_type=table_sub_model_name) root_dir = Path(__file__).absolute().parent.parent.parent.parent.parent
slanet_plus_model_path = os.path.join(root_dir, 'resources', 'slanet_plus', 'slanet-plus.onnx')
input_args = RapidTableInput(model_type=table_sub_model_name, model_path=slanet_plus_model_path)
else: else:
raise ValueError(f"Invalid table_sub_model_name: {table_sub_model_name}. It must be one of {sub_model_list}") raise ValueError(f"Invalid table_sub_model_name: {table_sub_model_name}. It must be one of {sub_model_list}")
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment