def __init__(self, model_path='./checkpoints/CTPN.pth'):
self.model = CTPN_Model()
self.use_gpu = torch.cuda.is_available()
if self.use_gpu:
self.model.cuda()
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
self.model.load_state_dict(
torch.load(model_path, map_location=device)['model_state_dict'])
for p in self.model.parameters():
p.requires_grad = False
self.model.eval()
self.prob_thresh = 0.5
from ctpn_model import CTPN_Model
from ctpn_utils import gen_anchor, bbox_transfor_inv, clip_box, filter_bbox,nms, TextProposalConnectorOriented
from ctpn_utils import resize
import config
prob_thresh = 0.8
width = 600
#device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
device = torch.device('cpu')
#weights = os.path.join(config.checkpoints_dir, 'trained weights file.pth.tar')
weights = config.model_path
model = CTPN_Model()
model.load_state_dict(torch.load(weights, map_location=device)['model_state_dict'])
model.to(device)
model.eval()
def dis(image):
cv2.imshow('image', image)
cv2.waitKey(0)
cv2.destroyAllWindows()
filenames = [os.path.join(config.img_path, file) for file in os.listdir(config.img_path)]
print(filenames)
