Isaacgonzales commited on
Commit
8f3d671
1 Parent(s): bd33281

update model

Browse files
Files changed (1) hide show
  1. model.py +2 -2
model.py CHANGED
@@ -12,7 +12,7 @@ from time import process_time
12
 
13
  device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
14
 
15
- model_recog = load_from_checkpoint("weights/parseq/last.ckpt").eval().to(device)
16
  img_transform = SceneTextDataModule.get_transform(model_recog.hparams.img_size)
17
 
18
  model = torch.load('weights/best_model.pth', map_location=torch.device(device))
@@ -42,7 +42,7 @@ def prediction(image_path):
42
 
43
  image = cv2.cvtColor(result, cv2.COLOR_BGR2RGB)
44
  im_pil = Image.fromarray(image)
45
- image = img_transform(im_pil).unsqueeze(0).to(device)
46
 
47
  p = model_recog(image).softmax(-1)
48
  pred, p = model_recog.tokenizer.decode(p)
 
12
 
13
  device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
14
 
15
+ model_recog = load_from_checkpoint("weights/parseq/last.ckpt").eval().to(device).float()
16
  img_transform = SceneTextDataModule.get_transform(model_recog.hparams.img_size)
17
 
18
  model = torch.load('weights/best_model.pth', map_location=torch.device(device))
 
42
 
43
  image = cv2.cvtColor(result, cv2.COLOR_BGR2RGB)
44
  im_pil = Image.fromarray(image)
45
+ image = img_transform(im_pil).float().unsqueeze(0).to(device)
46
 
47
  p = model_recog(image).softmax(-1)
48
  pred, p = model_recog.tokenizer.decode(p)