update model

model.py

@@ -10,10 +10,12 @@ import torch
 import cv2
 from time import process_time
 
-model_recog = load_from_checkpoint("weights/parseq/last.ckpt").eval().to("cpu")
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+model_recog = load_from_checkpoint("weights/parseq/last.ckpt").eval().to(device)
 img_transform = SceneTextDataModule.get_transform(model_recog.hparams.img_size)
 
-model = torch.load('weights/best_model.pth', map_location=torch.device('cpu'))
+model = torch.load('weights/best_model.pth', map_location=torch.device(device))
 model.eval()
 model.float()
 
@@ -21,18 +23,17 @@ SHAPE_X = 384
 SHAPE_Y = 384
 preprocessing_fn = smp.encoders.get_preprocessing_fn('resnet50')
 
-transform_compose = albu.Compose([
-            albu.Lambda(image=preprocessing_fn), albu.Resize(SHAPE_X, SHAPE_Y)
-          ])
 
-transform_tensor = transforms.ToTensor()
 
 def prediction(image_path):
     t_start = process_time()
     image = cv2.imread(image_path)
     image_original = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
-
+    transform_compose = albu.Compose([
+            albu.Lambda(image=preprocessing_fn), albu.Resize(SHAPE_X, SHAPE_Y)
+          ])
     image_result = transform_compose(image=image_original)["image"]
+    transform_tensor = transforms.ToTensor()
     tensor = transform_tensor(image_result)
     tensor = torch.unsqueeze(tensor, 0)
     output = model.predict(tensor.float())
@@ -41,7 +42,7 @@ def prediction(image_path):
     
     image = cv2.cvtColor(result, cv2.COLOR_BGR2RGB)
     im_pil = Image.fromarray(image)
-    image = img_transform(im_pil).unsqueeze(0).to("cpu")
+    image = img_transform(im_pil).unsqueeze(0).to(device)
     
     p = model_recog(image).softmax(-1)
     pred, p = model_recog.tokenizer.decode(p)