Update demo.py

demo.py

@@ -1,15 +1,29 @@
+# from transformers import AutoModel
 import argparse
 import logging
 import os
 import glob
 import tqdm
-import torch
+import torch, re
 import PIL
 import cv2
 import numpy as np
 import torch.nn.functional as F
 from torchvision import transforms
 from utils import Config, Logger, CharsetMapper
+import gradio as gr 
+#dfgdfg
+import gdown
+gdown.download(id='16PF_b4dURVkBt4OT7E-a-vq-SRxi0uDl', output='lol.pth')
+gdown.download(id='19rGjfo73P25O_keQv30snfe3IHrK0uV2', output='config.yaml')
+
+# gdown.download(id='1qyNV80qmYHx_r4KsG3_8PXQ6ff1a1dov', output='modules.zip')
+
+# gdown.download(id='1UMZ7i8SpfuNw0N2JvVY8euaNx9gu3x6N', output='configs.zip')
+
+# gdown.download(id='1yHD7_4DD_keUwGs2nenAYDaQ2CNEA5IU', output='data.zip')
+# os.system('unzip data.zip && unzip configs.zip && unzip modules.zip')
+
 
 def get_model(config):
     import importlib
@@ -21,12 +35,22 @@ def get_model(config):
     model = model.eval()
     return model
 
-def preprocess(img, width, height):
-    img = cv2.resize(np.array(img), (width, height))
-    img = transforms.ToTensor()(img).unsqueeze(0)
-    mean = torch.tensor([0.485, 0.456, 0.406])
-    std  = torch.tensor([0.229, 0.224, 0.225])
-    return (img-mean[...,None,None]) / std[...,None,None]
+
+def load(model, file, device=None, strict=True):
+    if device is None: device = 'cpu'
+    elif isinstance(device, int): device = torch.device('cuda', device)
+    assert os.path.isfile(file)
+    state = torch.load(file, map_location=device)
+    if set(state.keys()) == {'model', 'opt'}:
+        state = state['model']
+    model.load_state_dict(state, strict=strict)
+    return model
+
+config = Config('config.yaml')
+config.model_vision_checkpoint = None
+model = get_model(config)
+model = load(model, 'lol.pth')
+
 
 def postprocess(output, charset, model_eval):
     def _get_output(last_output, model_eval):
@@ -54,56 +78,27 @@ def postprocess(output, charset, model_eval):
     
     return pt_text, pt_scores, pt_lengths_
 
-def load(model, file, device=None, strict=True):
-    if device is None: device = 'cpu'
-    elif isinstance(device, int): device = torch.device('cuda', device)
-    assert os.path.isfile(file)
-    state = torch.load(file, map_location=device)
-    if set(state.keys()) == {'model', 'opt'}:
-        state = state['model']
-    model.load_state_dict(state, strict=strict)
-    return model
+def preprocess(img, width, height):
+    img = cv2.resize(np.array(img), (width, height))
+    img = transforms.ToTensor()(img).unsqueeze(0)
+    mean = torch.tensor([0.485, 0.456, 0.406])
+    std  = torch.tensor([0.229, 0.224, 0.225])
+    return (img-mean[...,None,None]) / std[...,None,None]
+
+def process_image(image):
+    charset = CharsetMapper(filename=config.dataset_charset_path, max_length=config.dataset_max_length + 1)
+
+    img = image.convert('RGB')
+    img = preprocess(img, config.dataset_image_width, config.dataset_image_height)
+    res = model(img)
+    return postprocess(res, charset, 'alignment')[0][0]
 
-def main():
-    parser = argparse.ArgumentParser()
-    parser.add_argument('--config', type=str, default='configs/train_abinet.yaml',
-                        help='path to config file')
-    parser.add_argument('--input', type=str, default='figs/test')
-    parser.add_argument('--cuda', type=int, default=-1)
-    parser.add_argument('--checkpoint', type=str, default='workdir/train-abinet/best-train-abinet.pth')
-    parser.add_argument('--model_eval', type=str, default='alignment', 
-                        choices=['alignment', 'vision', 'language'])
-    args = parser.parse_args()
-    config = Config(args.config)
-    if args.checkpoint is not None: config.model_checkpoint = args.checkpoint
-    if args.model_eval is not None: config.model_eval = args.model_eval
-    config.global_phase = 'test'
-    config.model_vision_checkpoint, config.model_language_checkpoint = None, None
-    device = 'cpu' if args.cuda < 0 else f'cuda:{args.cuda}'
-
-    Logger.init(config.global_workdir, config.global_name, config.global_phase)
-    Logger.enable_file()
-    logging.info(config)
-
-    logging.info('Construct model.')
-    model = get_model(config).to(device)
-    model = load(model, config.model_checkpoint, device=device)
-    charset = CharsetMapper(filename=config.dataset_charset_path,
-                            max_length=config.dataset_max_length + 1)
-
-    if os.path.isdir(args.input):
-        paths = [os.path.join(args.input, fname) for fname in os.listdir(args.input)]
-    else:
-        paths = glob.glob(os.path.expanduser(args.input))
-        assert paths, "The input path(s) was not found"
-    paths = sorted(paths)
-    for path in tqdm.tqdm(paths):
-        img = PIL.Image.open(path).convert('RGB')
-        img = preprocess(img, config.dataset_image_width, config.dataset_image_height)
-        img = img.to(device)
-        res = model(img)
-        pt_text, _, __ = postprocess(res, charset, config.model_eval)
-        logging.info(f'{path}: {pt_text[0]}')
-
-if __name__ == '__main__':
-    main()
+iface = gr.Interface(fn=process_image, 
+                     inputs=gr.inputs.Image(type="pil"), 
+                     outputs=gr.outputs.Textbox(),
+                     title="8kun kek",
+                     description="Making Jim Watkins sheete because he is a techlet pedo",
+                    #  article=article,
+                    #  examples=glob.glob('figs/test/*.png')
+                    )
+iface.launch(debug=True)