Isaacgonzales
commited on
Commit
•
8f3d671
1
Parent(s):
bd33281
update model
Browse files
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)
|