Commit e327e9ba authored by myhloli's avatar myhloli
Browse files

fix(table): add model path for slanet-plus to resolve RapidTableError

- Import os and pathlib modules to handle file paths
- Define the path to the slanet-plus model
- Update RapidTableInput initialization to include the model path
parent 7b61b418
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