|
|
| from strhub.data.module import SceneTextDataModule |
| from strhub.models.utils import load_from_checkpoint |
| from post import filter_mask |
| import segmentation_models_pytorch as smp |
| import albumentations as albu |
| from torchvision import transforms |
| from PIL import Image |
| import torch |
| import cv2 |
| from time import process_time |
| import imutils |
| device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") |
|
|
| model_recog = load_from_checkpoint("weights/parseq/last.ckpt").eval().to(device).float() |
| img_transform = SceneTextDataModule.get_transform(model_recog.hparams.img_size) |
|
|
| model = torch.load('weights/best_model.pth', map_location=torch.device(device)) |
| model.eval() |
|
|
|
|
| SHAPE_X = 384 |
| SHAPE_Y = 384 |
| preprocessing_fn = smp.encoders.get_preprocessing_fn('resnet50') |
|
|
|
|
|
|
| def prediction(image_path): |
| t_start = process_time() |
| image = cv2.imread(image_path) |
| image_original = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) |
| image_original = imutils.resize(image_original, width=image_original.shape[1]//5) |
| 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().to(device)) |
|
|
| result, img_vis = filter_mask(output, image_original) |
| |
| image = cv2.cvtColor(result, cv2.COLOR_BGR2RGB) |
| im_pil = Image.fromarray(image) |
| image = img_transform(im_pil).float().unsqueeze(0).to(device) |
| |
| p = model_recog(image).softmax(-1) |
| pred, p = model_recog.tokenizer.decode(p) |
| t_stop = process_time() |
| t_process = (t_stop - t_start)/3 |
| print(f'{image_path}: {pred[0]}, {t_process} seconds') |
| |
| return img_vis, pred[0], t_process |
|
|
|
|
| |
|
|
