xfmr-zem 0.2.2__py3-none-any.whl → 0.2.5__py3-none-any.whl

xfmr_zem/servers/ocr/deepdoc_vietocr/utils/file_utils.py

@@ -0,0 +1,246 @@
+#
+#  Copyright 2024 The InfiniFlow Authors. All Rights Reserved.
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+#
+import base64
+import json
+import os
+import re
+import sys
+import threading
+from io import BytesIO
+
+import pdfplumber
+from PIL import Image
+from cachetools import LRUCache, cached
+from ruamel.yaml import YAML
+
+from enum import Enum
+# from .db import FileType
+# from .constants import IMG_BASE64_PREFIX
+IMG_BASE64_PREFIX = 'data:image/png;base64,'
+
+class FileType(Enum):
+    PDF = 'pdf'
+    DOC = 'doc'
+    VISUAL = 'visual'
+    AURAL = 'aural'
+    VIRTUAL = 'virtual'
+    FOLDER = 'folder'
+    OTHER = "other"
+
+PROJECT_BASE = os.getenv("RAG_PROJECT_BASE") or os.getenv("RAG_DEPLOY_BASE")
+RAG_BASE = os.getenv("RAG_BASE")
+
+LOCK_KEY_pdfplumber = "global_shared_lock_pdfplumber"
+if LOCK_KEY_pdfplumber not in sys.modules:
+    sys.modules[LOCK_KEY_pdfplumber] = threading.Lock()
+
+
+def get_project_base_directory(*args):
+    base_dir = os.path.abspath(os.path.join(os.path.dirname(os.path.realpath(__file__)), os.pardir))
+    if args:
+        return os.path.join(base_dir, *args)
+    return base_dir
+
+def get_rag_directory(*args):
+    global RAG_BASE
+    if RAG_BASE is None:
+        RAG_BASE = os.path.abspath(
+            os.path.join(
+                os.path.dirname(os.path.realpath(__file__)),
+                os.pardir,
+                os.pardir,
+                os.pardir,
+            )
+        )
+    if args:
+        return os.path.join(RAG_BASE, *args)
+    return RAG_BASE
+
+
+def get_rag_python_directory(*args):
+    return get_rag_directory("python", *args)
+
+
+def get_home_cache_dir():
+    dir = os.path.join(os.path.expanduser('~'), ".ragflow")
+    try:
+        os.mkdir(dir)
+    except OSError:
+        pass
+    return dir
+
+
+@cached(cache=LRUCache(maxsize=10))
+def load_json_conf(conf_path):
+    if os.path.isabs(conf_path):
+        json_conf_path = conf_path
+    else:
+        json_conf_path = os.path.join(get_project_base_directory(), conf_path)
+    try:
+        with open(json_conf_path) as f:
+            return json.load(f)
+    except BaseException:
+        raise EnvironmentError(
+            "loading json file config from '{}' failed!".format(json_conf_path)
+        )
+
+
+def dump_json_conf(config_data, conf_path):
+    if os.path.isabs(conf_path):
+        json_conf_path = conf_path
+    else:
+        json_conf_path = os.path.join(get_project_base_directory(), conf_path)
+    try:
+        with open(json_conf_path, "w") as f:
+            json.dump(config_data, f, indent=4)
+    except BaseException:
+        raise EnvironmentError(
+            "loading json file config from '{}' failed!".format(json_conf_path)
+        )
+
+
+def load_json_conf_real_time(conf_path):
+    if os.path.isabs(conf_path):
+        json_conf_path = conf_path
+    else:
+        json_conf_path = os.path.join(get_project_base_directory(), conf_path)
+    try:
+        with open(json_conf_path) as f:
+            return json.load(f)
+    except BaseException:
+        raise EnvironmentError(
+            "loading json file config from '{}' failed!".format(json_conf_path)
+        )
+
+
+def load_yaml_conf(conf_path):
+    if not os.path.isabs(conf_path):
+        conf_path = os.path.join(get_project_base_directory(), conf_path)
+    try:
+        with open(conf_path) as f:
+            yaml = YAML(typ='safe', pure=True)
+            return yaml.load(f)
+    except Exception as e:
+        raise EnvironmentError(
+            "loading yaml file config from {} failed:".format(conf_path), e
+        )
+
+
+def rewrite_yaml_conf(conf_path, config):
+    if not os.path.isabs(conf_path):
+        conf_path = os.path.join(get_project_base_directory(), conf_path)
+    try:
+        with open(conf_path, "w") as f:
+            yaml = YAML(typ="safe")
+            yaml.dump(config, f)
+    except Exception as e:
+        raise EnvironmentError(
+            "rewrite yaml file config {} failed:".format(conf_path), e
+        )
+
+
+def rewrite_json_file(filepath, json_data):
+    with open(filepath, "w", encoding='utf-8') as f:
+        json.dump(json_data, f, indent=4, separators=(",", ": "))
+    f.close()
+
+
+def filename_type(filename):
+    filename = filename.lower()
+    if re.match(r".*\.pdf$", filename):
+        return FileType.PDF.value
+
+    if re.match(
+             r".*\.(eml|doc|docx|ppt|pptx|yml|xml|htm|json|csv|txt|ini|xls|xlsx|wps|rtf|hlp|pages|numbers|key|md|py|js|java|c|cpp|h|php|go|ts|sh|cs|kt|html|sql)$", filename):
+        return FileType.DOC.value
+
+    if re.match(
+            r".*\.(wav|flac|ape|alac|wavpack|wv|mp3|aac|ogg|vorbis|opus|mp3)$", filename):
+        return FileType.AURAL.value
+
+    if re.match(r".*\.(jpg|jpeg|png|tif|gif|pcx|tga|exif|fpx|svg|psd|cdr|pcd|dxf|ufo|eps|ai|raw|WMF|webp|avif|apng|icon|ico|mpg|mpeg|avi|rm|rmvb|mov|wmv|asf|dat|asx|wvx|mpe|mpa|mp4)$", filename):
+        return FileType.VISUAL.value
+
+    return FileType.OTHER.value
+
+def thumbnail_img(filename, blob):
+    """
+    MySQL LongText max length is 65535
+    """
+    filename = filename.lower()
+    if re.match(r".*\.pdf$", filename):
+        with sys.modules[LOCK_KEY_pdfplumber]:
+            pdf = pdfplumber.open(BytesIO(blob))
+            buffered = BytesIO()
+            resolution = 32
+            img = None
+            for _ in range(10):
+                # https://github.com/jsvine/pdfplumber?tab=readme-ov-file#creating-a-pageimage-with-to_image
+                pdf.pages[0].to_image(resolution=resolution).annotated.save(buffered, format="png")
+                img = buffered.getvalue()
+                if len(img) >= 64000 and resolution >= 2:
+                    resolution = resolution / 2
+                    buffered = BytesIO()
+                else:
+                    break
+        pdf.close()
+        return img
+
+    elif re.match(r".*\.(jpg|jpeg|png|tif|gif|icon|ico|webp)$", filename):
+        image = Image.open(BytesIO(blob))
+        image.thumbnail((30, 30))
+        buffered = BytesIO()
+        image.save(buffered, format="png")
+        return buffered.getvalue()
+
+    elif re.match(r".*\.(ppt|pptx)$", filename):
+        import aspose.slides as slides
+        import aspose.pydrawing as drawing
+        try:
+            with slides.Presentation(BytesIO(blob)) as presentation:
+                buffered = BytesIO()
+                scale = 0.03
+                img = None
+                for _ in range(10):
+                    # https://reference.aspose.com/slides/python-net/aspose.slides/slide/get_thumbnail/#float-float
+                    presentation.slides[0].get_thumbnail(scale, scale).save(
+                        buffered, drawing.imaging.ImageFormat.png)
+                    img = buffered.getvalue()
+                    if len(img) >= 64000:
+                        scale = scale / 2.0
+                        buffered = BytesIO()
+                    else:
+                        break
+                return img
+        except Exception:
+            pass
+    return None
+
+
+def thumbnail(filename, blob):
+    img = thumbnail_img(filename, blob)
+    if img is not None:
+        return IMG_BASE64_PREFIX + \
+            base64.b64encode(img).decode("utf-8")
+    else:
+        return ''
+
+
+def traversal_files(base):
+    for root, ds, fs in os.walk(base):
+        for f in fs:
+            fullname = os.path.join(root, f)
+            yield fullname

xfmr_zem/servers/ocr/deepdoc_vietocr/vietocr/config/base.yml

@@ -0,0 +1,58 @@
+# change to list chars of your dataset or use default vietnamese chars
+vocab: 'aAàÀảẢãÃáÁạẠăĂằẰẳẲẵẴắẮặẶâÂầẦẩẨẫẪấẤậẬbBcCdDđĐeEèÈẻẺẽẼéÉẹẸêÊềỀểỂễỄếẾệỆfFgGhHiIìÌỉỈĩĨíÍịỊjJkKlLmMnNoOòÒỏỎõÕóÓọỌôÔồỒổỔỗỖốỐộỘơƠờỜởỞỡỠớỚợỢpPqQrRsStTuUùÙủỦũŨúÚụỤưƯừỪửỬữỮứỨựỰvVwWxXyYỳỲỷỶỹỸýÝỵỴzZ0123456789!"#$%&''()*+,-./:;<=>?@[\]^_`{|}~ '
+
+# cpu, cuda, cuda:0
+device: cuda:0
+
+seq_modeling: transformer
+transformer:  
+    d_model: 256
+    nhead: 8
+    num_encoder_layers: 6
+    num_decoder_layers: 6
+    dim_feedforward: 2048
+    max_seq_length: 1024
+    pos_dropout: 0.1
+    trans_dropout: 0.1
+
+optimizer:
+    max_lr: 0.0003 
+    pct_start: 0.1
+
+trainer:
+    batch_size: 32
+    print_every: 200
+    valid_every: 4000
+    iters: 100000
+    # where to save our model for prediction
+    export: ./weights/transformerocr.pth
+    checkpoint: ./checkpoint/transformerocr_checkpoint.pth
+    log: ./train.log
+    # null to disable compuate accuracy, or change to number of sample to enable validiation while training
+    metrics: null
+
+dataset:    
+    # name of your dataset
+    name: data
+    # path to annotation and image
+    data_root: ./img/
+    train_annotation: annotation_train.txt
+    valid_annotation: annotation_val_small.txt
+    # resize image to 32 height, larger height will increase accuracy
+    image_height: 32
+    image_min_width: 32
+    image_max_width: 512
+
+dataloader:
+    num_workers: 3
+    pin_memory: True
+
+aug:
+    image_aug: true
+    masked_language_model: true
+
+predictor:
+    # disable or enable beamsearch while prediction, use beamsearch will be slower
+    beamsearch: False
+
+quiet: False 

xfmr_zem/servers/ocr/deepdoc_vietocr/vietocr/config/vgg-seq2seq.yml

@@ -0,0 +1,38 @@
+pretrain: 
+    id_or_url: 1nTKlEog9YFK74kPyX0qLwCWi60_YHHk4
+    md5: efcabaa6d3adfca8e52bda2fd7d2ee04
+    cached: /tmp/tranformerorc.pth
+
+# url or local path
+weights: https://drive.google.com/uc?id=1nTKlEog9YFK74kPyX0qLwCWi60_YHHk4
+
+backbone: vgg19_bn
+cnn:
+    # pooling stride size
+    ss:
+        - [2, 2]
+        - [2, 2]
+        - [2, 1]
+        - [2, 1]
+        - [1, 1]         
+    # pooling kernel size 
+    ks:
+        - [2, 2]
+        - [2, 2]
+        - [2, 1]
+        - [2, 1]
+        - [1, 1]
+    # dim of ouput feature map
+    hidden: 256
+
+seq_modeling: seq2seq
+transformer:
+    encoder_hidden: 256
+    decoder_hidden: 256
+    img_channel: 256
+    decoder_embedded: 256
+    dropout: 0.1
+
+optimizer:
+    max_lr: 0.001
+    pct_start: 0.1

xfmr_zem/servers/ocr/deepdoc_vietocr/vietocr/model/backbone/cnn.py

@@ -0,0 +1,25 @@
+import torch
+from torch import nn
+
+from . import vgg
+
+class CNN(nn.Module):
+    def __init__(self, backbone, **kwargs):
+        super(CNN, self).__init__()
+
+        if backbone == 'vgg11_bn':
+            self.model = vgg.vgg11_bn(**kwargs)
+        elif backbone == 'vgg19_bn':
+            self.model = vgg.vgg19_bn(**kwargs)
+
+    def forward(self, x):
+        return self.model(x)
+
+    def freeze(self):
+        for name, param in self.model.features.named_parameters():
+            if name != 'last_conv_1x1':
+                param.requires_grad = False
+
+    def unfreeze(self):
+        for param in self.model.features.parameters():
+            param.requires_grad = True

xfmr_zem/servers/ocr/deepdoc_vietocr/vietocr/model/backbone/vgg.py

@@ -0,0 +1,51 @@
+import torch
+from torch import nn
+from torchvision import models
+from einops import rearrange
+from torchvision.models._utils import IntermediateLayerGetter
+
+
+class Vgg(nn.Module):
+    def __init__(self, name, ss, ks, hidden, pretrained=True, dropout=0.5):
+        super(Vgg, self).__init__()
+
+        weights = "DEFAULT" if pretrained else None
+        if name == 'vgg11_bn':
+            cnn = models.vgg11_bn(weights=weights)
+        elif name == 'vgg19_bn':
+            cnn = models.vgg19_bn(weights=weights)
+
+        pool_idx = 0
+
+        for i, layer in enumerate(cnn.features):
+            if isinstance(layer, torch.nn.MaxPool2d):
+                cnn.features[i] = torch.nn.AvgPool2d(kernel_size=ks[pool_idx], stride=ss[pool_idx], padding=0)
+                pool_idx += 1
+
+        self.features = cnn.features
+        self.dropout = nn.Dropout(dropout)
+        self.last_conv_1x1 = nn.Conv2d(512, hidden, 1)
+
+    def forward(self, x):
+        """
+        Shape:
+            - x: (N, C, H, W)
+            - output: (W, N, C)
+        """
+
+        conv = self.features(x)
+        conv = self.dropout(conv)
+        conv = self.last_conv_1x1(conv)
+
+#        conv = rearrange(conv, 'b d h w -> b d (w h)')
+        conv = conv.permute(0, 1, 3, 2)
+        conv = conv.flatten(2)
+        conv = conv.permute(2, 0, 1)
+
+        return conv
+
+def vgg11_bn(ss, ks, hidden, pretrained=True, dropout=0.5):
+    return Vgg('vgg11_bn', ss, ks, hidden, pretrained, dropout)
+
+def vgg19_bn(ss, ks, hidden, pretrained=True, dropout=0.5):
+    return Vgg('vgg19_bn', ss, ks, hidden, pretrained, dropout)

xfmr_zem/servers/ocr/deepdoc_vietocr/vietocr/model/seqmodel/seq2seq.py

@@ -0,0 +1,175 @@
+import torch
+import torch.nn as nn
+import torch.optim as optim
+import torch.nn.functional as F
+
+class Encoder(nn.Module):
+    def __init__(self, emb_dim, enc_hid_dim, dec_hid_dim, dropout):
+        super().__init__()
+                
+        self.rnn = nn.GRU(emb_dim, enc_hid_dim, bidirectional = True)
+        self.fc = nn.Linear(enc_hid_dim * 2, dec_hid_dim)
+        self.dropout = nn.Dropout(dropout)
+        
+    def forward(self, src):
+        """
+        src: src_len x batch_size x img_channel
+        outputs: src_len x batch_size x hid_dim 
+        hidden: batch_size x hid_dim
+        """
+
+        embedded = self.dropout(src)
+        
+        outputs, hidden = self.rnn(embedded)
+                                 
+        hidden = torch.tanh(self.fc(torch.cat((hidden[-2,:,:], hidden[-1,:,:]), dim = 1)))
+        
+        return outputs, hidden
+
+class Attention(nn.Module):
+    def __init__(self, enc_hid_dim, dec_hid_dim):
+        super().__init__()
+        
+        self.attn = nn.Linear((enc_hid_dim * 2) + dec_hid_dim, dec_hid_dim)
+        self.v = nn.Linear(dec_hid_dim, 1, bias = False)
+        
+    def forward(self, hidden, encoder_outputs):
+        """
+        hidden: batch_size x hid_dim
+        encoder_outputs: src_len x batch_size x hid_dim,
+        outputs: batch_size x src_len
+        """
+        
+        batch_size = encoder_outputs.shape[1]
+        src_len = encoder_outputs.shape[0]
+        
+        hidden = hidden.unsqueeze(1).repeat(1, src_len, 1)
+  
+        encoder_outputs = encoder_outputs.permute(1, 0, 2)
+        
+        energy = torch.tanh(self.attn(torch.cat((hidden, encoder_outputs), dim = 2))) 
+        
+        attention = self.v(energy).squeeze(2)
+        
+        return F.softmax(attention, dim = 1)
+
+class Decoder(nn.Module):
+    def __init__(self, output_dim, emb_dim, enc_hid_dim, dec_hid_dim, dropout, attention):
+        super().__init__()
+
+        self.output_dim = output_dim
+        self.attention = attention
+        
+        self.embedding = nn.Embedding(output_dim, emb_dim)
+        self.rnn = nn.GRU((enc_hid_dim * 2) + emb_dim, dec_hid_dim)
+        self.fc_out = nn.Linear((enc_hid_dim * 2) + dec_hid_dim + emb_dim, output_dim)
+        self.dropout = nn.Dropout(dropout)
+        
+    def forward(self, input, hidden, encoder_outputs):
+        """
+        inputs: batch_size
+        hidden: batch_size x hid_dim
+        encoder_outputs: src_len x batch_size x hid_dim
+        """
+             
+        input = input.unsqueeze(0)
+        
+        embedded = self.dropout(self.embedding(input))
+        
+        a = self.attention(hidden, encoder_outputs)
+                
+        a = a.unsqueeze(1)
+        
+        encoder_outputs = encoder_outputs.permute(1, 0, 2)
+        
+        weighted = torch.bmm(a, encoder_outputs)
+        
+        weighted = weighted.permute(1, 0, 2)
+        
+        rnn_input = torch.cat((embedded, weighted), dim = 2)
+        
+        output, hidden = self.rnn(rnn_input, hidden.unsqueeze(0))
+        
+        assert (output == hidden).all()
+        
+        embedded = embedded.squeeze(0)
+        output = output.squeeze(0)
+        weighted = weighted.squeeze(0)
+        
+        prediction = self.fc_out(torch.cat((output, weighted, embedded), dim = 1))
+        
+        return prediction, hidden.squeeze(0), a.squeeze(1)
+
+class Seq2Seq(nn.Module):
+    def __init__(self, vocab_size, encoder_hidden, decoder_hidden, img_channel, decoder_embedded, dropout=0.1):
+        super().__init__()
+        
+        attn = Attention(encoder_hidden, decoder_hidden)
+        
+        self.encoder = Encoder(img_channel, encoder_hidden, decoder_hidden, dropout)
+        self.decoder = Decoder(vocab_size, decoder_embedded, encoder_hidden, decoder_hidden, dropout, attn)
+        
+    def forward_encoder(self, src):       
+        """
+        src: timestep x batch_size x channel
+        hidden: batch_size x hid_dim
+        encoder_outputs: src_len x batch_size x hid_dim
+        """
+
+        encoder_outputs, hidden = self.encoder(src)
+
+        return (hidden, encoder_outputs)
+
+    def forward_decoder(self, tgt, memory):
+        """
+        tgt: timestep x batch_size 
+        hidden: batch_size x hid_dim
+        encouder: src_len x batch_size x hid_dim
+        output: batch_size x 1 x vocab_size
+        """
+        
+        tgt = tgt[-1]
+        hidden, encoder_outputs = memory
+        output, hidden, _ = self.decoder(tgt, hidden, encoder_outputs)
+        output = output.unsqueeze(1)
+        
+        return output, (hidden, encoder_outputs)
+
+    def forward(self, src, trg):
+        """
+        src: time_step x batch_size
+        trg: time_step x batch_size
+        outputs: batch_size x time_step x vocab_size
+        """
+
+        batch_size = src.shape[1]
+        trg_len = trg.shape[0]
+        trg_vocab_size = self.decoder.output_dim
+        device = src.device
+
+        outputs = torch.zeros(trg_len, batch_size, trg_vocab_size).to(device)
+        encoder_outputs, hidden = self.encoder(src)
+                
+        for t in range(trg_len):
+            input = trg[t] 
+            output, hidden, _ = self.decoder(input, hidden, encoder_outputs)
+            
+            outputs[t] = output
+            
+        outputs = outputs.transpose(0, 1).contiguous()
+
+        return outputs
+
+    def expand_memory(self, memory, beam_size):
+        hidden, encoder_outputs = memory
+        hidden = hidden.repeat(beam_size, 1)
+        encoder_outputs = encoder_outputs.repeat(1, beam_size, 1)
+
+        return (hidden, encoder_outputs)
+    
+    def get_memory(self, memory, i):
+        hidden, encoder_outputs = memory
+        hidden = hidden[[i]]
+        encoder_outputs = encoder_outputs[:, [i],:]
+
+        return (hidden, encoder_outputs)

xfmr_zem/servers/ocr/deepdoc_vietocr/vietocr/model/transformerocr.py

@@ -0,0 +1,29 @@
+from .backbone.cnn import CNN
+from .seqmodel.seq2seq import Seq2Seq
+from torch import nn
+
+class VietOCR(nn.Module):
+    def __init__(self, vocab_size,
+                 backbone,
+                 cnn_args,
+                 transformer_args, seq_modeling='transformer'):
+
+        super(VietOCR, self).__init__()
+
+        self.cnn = CNN(backbone, **cnn_args)
+        self.seq_modeling = seq_modeling
+        self.transformer = Seq2Seq(vocab_size, **transformer_args)
+
+
+    def forward(self, img, tgt_input, tgt_key_padding_mask):
+        """
+        Shape:
+            - img: (N, C, H, W)
+            - tgt_input: (T, N)
+            - tgt_key_padding_mask: (N, T)
+            - output: b t v
+        """
+        src = self.cnn(img)
+        outputs = self.transformer(src, tgt_input)
+
+        return outputs

xfmr_zem/servers/ocr/deepdoc_vietocr/vietocr/model/vocab.py

@@ -0,0 +1,36 @@
+class Vocab():
+    def __init__(self, chars):
+        self.pad = 0
+        self.go = 1
+        self.eos = 2
+        self.mask_token = 3
+
+        self.chars = chars
+
+        self.c2i = {c:i+4 for i, c in enumerate(chars)}
+
+        self.i2c = {i+4:c for i, c in enumerate(chars)}
+        
+        self.i2c[0] = '<pad>'
+        self.i2c[1] = '<sos>'
+        self.i2c[2] = '<eos>'
+        self.i2c[3] = '*'
+
+    def encode(self, chars):
+        return [self.go] + [self.c2i[c] for c in chars] + [self.eos]
+    
+    def decode(self, ids):
+        first = 1 if self.go in ids else 0
+        last = ids.index(self.eos) if self.eos in ids else None
+        sent = ''.join([self.i2c[i] for i in ids[first:last]])
+        return sent
+    
+    def __len__(self):
+        return len(self.c2i) + 4
+    
+    def batch_decode(self, arr):
+        texts = [self.decode(ids) for ids in arr]
+        return texts
+
+    def __str__(self):
+        return self.chars