Spaces:
Runtime error
Runtime error
Commit
•
835dcbb
1
Parent(s):
43a2d3b
Upload demo copy.py
Browse files- demo copy.py +106 -0
demo copy.py
ADDED
@@ -0,0 +1,106 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# from transformers import AutoModel
|
2 |
+
import argparse
|
3 |
+
import logging
|
4 |
+
import os
|
5 |
+
import glob
|
6 |
+
import tqdm
|
7 |
+
import torch, re
|
8 |
+
import PIL
|
9 |
+
import cv2
|
10 |
+
import numpy as np
|
11 |
+
import torch.nn.functional as F
|
12 |
+
from torchvision import transforms
|
13 |
+
from utils import Config, Logger, CharsetMapper
|
14 |
+
import gradio as gr
|
15 |
+
|
16 |
+
import gdown
|
17 |
+
gdown.download(id='16PF_b4dURVkBt4OT7E-a-vq-SRxi0uDl', output='lol.pth')
|
18 |
+
gdown.download(id='19rGjfo73P25O_keQv30snfe3IHrK0uV2', output='config.yaml')
|
19 |
+
|
20 |
+
gdown.download(id='1qyNV80qmYHx_r4KsG3_8PXQ6ff1a1dov', output='modules.zip')
|
21 |
+
os.system('unzip modules.zip')
|
22 |
+
|
23 |
+
gdown.download(id='1UMZ7i8SpfuNw0N2JvVY8euaNx9gu3x6N', output='configs.zip')
|
24 |
+
os.system('unzip configs.zip')
|
25 |
+
|
26 |
+
gdown.download(id='1yHD7_4DD_keUwGs2nenAYDaQ2CNEA5IU', output='data.zip')
|
27 |
+
os.system('unzip data.zip')
|
28 |
+
|
29 |
+
|
30 |
+
def get_model(config):
|
31 |
+
import importlib
|
32 |
+
names = config.model_name.split('.')
|
33 |
+
module_name, class_name = '.'.join(names[:-1]), names[-1]
|
34 |
+
cls = getattr(importlib.import_module(module_name), class_name)
|
35 |
+
model = cls(config)
|
36 |
+
logging.info(model)
|
37 |
+
model = model.eval()
|
38 |
+
return model
|
39 |
+
|
40 |
+
|
41 |
+
def load(model, file, device=None, strict=True):
|
42 |
+
if device is None: device = 'cpu'
|
43 |
+
elif isinstance(device, int): device = torch.device('cuda', device)
|
44 |
+
assert os.path.isfile(file)
|
45 |
+
state = torch.load(file, map_location=device)
|
46 |
+
if set(state.keys()) == {'model', 'opt'}:
|
47 |
+
state = state['model']
|
48 |
+
model.load_state_dict(state, strict=strict)
|
49 |
+
return model
|
50 |
+
|
51 |
+
config = Config('config.yaml')
|
52 |
+
config.model_vision_checkpoint = None
|
53 |
+
model = get_model(config)
|
54 |
+
model = load(model, 'lol.pth')
|
55 |
+
|
56 |
+
|
57 |
+
def postprocess(output, charset, model_eval):
|
58 |
+
def _get_output(last_output, model_eval):
|
59 |
+
if isinstance(last_output, (tuple, list)):
|
60 |
+
for res in last_output:
|
61 |
+
if res['name'] == model_eval: output = res
|
62 |
+
else: output = last_output
|
63 |
+
return output
|
64 |
+
|
65 |
+
def _decode(logit):
|
66 |
+
""" Greed decode """
|
67 |
+
out = F.softmax(logit, dim=2)
|
68 |
+
pt_text, pt_scores, pt_lengths = [], [], []
|
69 |
+
for o in out:
|
70 |
+
text = charset.get_text(o.argmax(dim=1), padding=False, trim=False)
|
71 |
+
text = text.split(charset.null_char)[0] # end at end-token
|
72 |
+
pt_text.append(text)
|
73 |
+
pt_scores.append(o.max(dim=1)[0])
|
74 |
+
pt_lengths.append(min(len(text) + 1, charset.max_length)) # one for end-token
|
75 |
+
return pt_text, pt_scores, pt_lengths
|
76 |
+
|
77 |
+
output = _get_output(output, model_eval)
|
78 |
+
logits, pt_lengths = output['logits'], output['pt_lengths']
|
79 |
+
pt_text, pt_scores, pt_lengths_ = _decode(logits)
|
80 |
+
|
81 |
+
return pt_text, pt_scores, pt_lengths_
|
82 |
+
|
83 |
+
def preprocess(img, width, height):
|
84 |
+
img = cv2.resize(np.array(img), (width, height))
|
85 |
+
img = transforms.ToTensor()(img).unsqueeze(0)
|
86 |
+
mean = torch.tensor([0.485, 0.456, 0.406])
|
87 |
+
std = torch.tensor([0.229, 0.224, 0.225])
|
88 |
+
return (img-mean[...,None,None]) / std[...,None,None]
|
89 |
+
|
90 |
+
def process_image(image):
|
91 |
+
charset = CharsetMapper(filename=config.dataset_charset_path, max_length=config.dataset_max_length + 1)
|
92 |
+
|
93 |
+
img = image.convert('RGB')
|
94 |
+
img = preprocess(img, config.dataset_image_width, config.dataset_image_height)
|
95 |
+
res = model(img)
|
96 |
+
return postprocess(res, charset, 'alignment')[0][0]
|
97 |
+
|
98 |
+
iface = gr.Interface(fn=process_image,
|
99 |
+
inputs=gr.inputs.Image(type="pil"),
|
100 |
+
outputs=gr.outputs.Textbox(),
|
101 |
+
title="8kun kek",
|
102 |
+
description="Making Jim Watkins sheete because he is a techlet pedo",
|
103 |
+
# article=article,
|
104 |
+
# examples=glob.glob('figs/test/*.png')
|
105 |
+
)
|
106 |
+
iface.launch(debug=True)
|