captchaboy/dfff4444

1# from transformers import AutoModel2import argparse3import logging4import os5import glob6import tqdm7import torch, re8import PIL9import cv210import numpy as np11import torch.nn.functional as F12from torchvision import transforms13from utils import Config, Logger, CharsetMapper14import gradio as gr 15#dfgdfg16import gdown17gdown.download(id='16PF_b4dURVkBt4OT7E-a-vq-SRxi0uDl', output='lol.pth')18gdown.download(id='19rGjfo73P25O_keQv30snfe3IHrK0uV2', output='config.yaml')19 20# gdown.download(id='1qyNV80qmYHx_r4KsG3_8PXQ6ff1a1dov', output='modules.zip')21 22# gdown.download(id='1UMZ7i8SpfuNw0N2JvVY8euaNx9gu3x6N', output='configs.zip')23 24# gdown.download(id='1yHD7_4DD_keUwGs2nenAYDaQ2CNEA5IU', output='data.zip')25# os.system('unzip data.zip && unzip configs.zip && unzip modules.zip')26 27 28def get_model(config):29    import importlib30    names = config.model_name.split('.')31    module_name, class_name = '.'.join(names[:-1]), names[-1]32    cls = getattr(importlib.import_module(module_name), class_name)33    model = cls(config)34    logging.info(model)35    model = model.eval()36    return model37 38 39def load(model, file, device=None, strict=True):40    if device is None: device = 'cpu'41    elif isinstance(device, int): device = torch.device('cuda', device)42    assert os.path.isfile(file)43    state = torch.load(file, map_location=device)44    if set(state.keys()) == {'model', 'opt'}:45        state = state['model']46    model.load_state_dict(state, strict=strict)47    return model48 49config = Config('config.yaml')50config.model_vision_checkpoint = None51model = get_model(config)52model = load(model, 'lol.pth')53 54 55def postprocess(output, charset, model_eval):56    def _get_output(last_output, model_eval):57        if isinstance(last_output, (tuple, list)): 58            for res in last_output:59                if res['name'] == model_eval: output = res60        else: output = last_output61        return output62 63    def _decode(logit):64        """ Greed decode """65        out = F.softmax(logit, dim=2)66        pt_text, pt_scores, pt_lengths = [], [], []67        for o in out:68            text = charset.get_text(o.argmax(dim=1), padding=False, trim=False)69            text = text.split(charset.null_char)[0]  # end at end-token70            pt_text.append(text)71            pt_scores.append(o.max(dim=1)[0])72            pt_lengths.append(min(len(text) + 1, charset.max_length))  # one for end-token73        return pt_text, pt_scores, pt_lengths74 75    output = _get_output(output, model_eval)76    logits, pt_lengths = output['logits'], output['pt_lengths']77    pt_text, pt_scores, pt_lengths_ = _decode(logits)78    79    return pt_text, pt_scores, pt_lengths_80 81def preprocess(img, width, height):82    img = cv2.resize(np.array(img), (width, height))83    img = transforms.ToTensor()(img).unsqueeze(0)84    mean = torch.tensor([0.485, 0.456, 0.406])85    std  = torch.tensor([0.229, 0.224, 0.225])86    return (img-mean[...,None,None]) / std[...,None,None]87 88def process_image(image):89    charset = CharsetMapper(filename=config.dataset_charset_path, max_length=config.dataset_max_length + 1)90 91    img = image.convert('RGB')92    img = preprocess(img, config.dataset_image_width, config.dataset_image_height)93    res = model(img)94    return postprocess(res, charset, 'alignment')[0][0]95 96iface = gr.Interface(fn=process_image, 97                     inputs=gr.inputs.Image(type="pil"), 98                     outputs=gr.outputs.Textbox(),99                     title="8kun kek",100                     description="Making Jim Watkins sheete because he is a techlet pedo",101                    #  article=article,102                    #  examples=glob.glob('figs/test/*.png')103                    )104iface.launch(debug=True)