def load_refinenet_model(cuda: bool = False):
# get refine net path
home_path = str(Path.home())
weight_path = os.path.join(home_path, ".craft_text_detector", "weights",
"craft_refiner_CTW1500.pth")
# load refine net
from craft_text_detector.models.refinenet import RefineNet
refine_net = RefineNet() # initialize
# check if weights are already downloaded, if not download
url = REFINENET_GDRIVE_URL
if os.path.isfile(weight_path) is not True:
print("Craft text refiner weight will be downloaded to {}".format(
weight_path))
file_utils.download(url=url, save_path=weight_path)
# arange device
if cuda:
refine_net.load_state_dict(copyStateDict(
torch_utils.load(weight_path)))
refine_net = refine_net.cuda()
refine_net = torch_utils.DataParallel(refine_net)
torch_utils.cudnn_benchmark = False
else:
refine_net.load_state_dict(
copyStateDict(torch_utils.load(weight_path, map_location="cpu")))
refine_net.eval()
return refine_net
def load_refinenet_model(
cuda: bool = False,
weight_path: Optional[Union[str, Path]] = None
):
# get refine net path
if weight_path is None:
home_path = Path.home()
weight_path = Path(
home_path,
".craft_text_detector",
"weights",
"craft_refiner_CTW1500.pth"
)
weight_path = Path(weight_path).resolve()
weight_path.parent.mkdir(exist_ok=True, parents=True)
weight_path = str(weight_path)
# load refine net
from craft_text_detector.models.refinenet import RefineNet
refine_net = RefineNet() # initialize
# check if weights are already downloaded, if not download
url = REFINENET_GDRIVE_URL
if not os.path.isfile(weight_path):
print("Craft text refiner weight will be downloaded to {}".format(weight_path))
file_utils.download(url=url, save_path=weight_path)
# arange device
if cuda:
refine_net.load_state_dict(copyStateDict(torch_utils.load(weight_path)))
refine_net = refine_net.cuda()
refine_net = torch_utils.DataParallel(refine_net)
torch_utils.cudnn_benchmark = False
else:
refine_net.load_state_dict(
copyStateDict(torch_utils.load(weight_path, map_location="cpu"))
)
refine_net.eval()
return refine_net
def load_refinenet_model(cuda: bool = False):
'''
Load a refine net model, the refine net is used to make better
heatmaps based on the outputs and intermediate features of craft.
It is also trained using the same semisupervised objective function.
Args:
cuda (bool): if True, to use GPU for compute else cpu
Returns
A pytorch model object
'''
# get refine net path
home_path = str(Path.home())
weight_path = os.path.join(home_path, ".craft_text_detector", "weights",
"craft_refiner_CTW1500.pth")
# load refine net
from craft_text_detector.models.refinenet import RefineNet
refine_net = RefineNet() # initialize
# check if weights are already downloaded, if not download
url = REFINENET_GDRIVE_URL
if os.path.isfile(weight_path) is not True:
print("Craft text refiner weight will be downloaded to {}".format(
weight_path))
file_utils.download(url=url, save_path=weight_path)
# arange device
if cuda:
refine_net.load_state_dict(copyStateDict(
torch_utils.load(weight_path)))
refine_net = refine_net.cuda()
refine_net = torch_utils.DataParallel(refine_net)
torch_utils.cudnn_benchmark = False
else:
refine_net.load_state_dict(
copyStateDict(torch_utils.load(weight_path, map_location="cpu")))
refine_net.eval()
return refine_net
def load_craftnet_model(cuda: bool = False):
'''
Loads and returns the craftnet model with the state value as defined in the
state dictionary
Args:
cuda (bool): if True, to use GPU for compute else cpu
Returns
A pytorch model object
'''
# get craft net path
home_path = str(Path.home())
weight_path = os.path.join(home_path, ".craft_text_detector", "weights",
"craft_mlt_25k.pth")
# load craft net
from craft_text_detector.models.craftnet import CraftNet
craft_net = CraftNet() # initialize
# check if weights are already downloaded, if not download
url = CRAFT_GDRIVE_URL
if os.path.isfile(weight_path) is not True:
print("Craft text detector weight will be downloaded to {}".format(
weight_path))
file_utils.download(url=url, save_path=weight_path)
# arange device
if cuda:
craft_net.load_state_dict(copyStateDict(torch_utils.load(weight_path)))
craft_net = craft_net.cuda()
craft_net = torch_utils.DataParallel(craft_net)
torch_utils.cudnn_benchmark = False
else:
craft_net.load_state_dict(
copyStateDict(torch_utils.load(weight_path, map_location="cpu")))
craft_net.eval()
return craft_net
