Upload 7 files

.gitattributes

@@ -1,36 +1,35 @@
-*.7z filter=lfs diff=lfs merge=lfs -text
-*.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
-*.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
-*.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
-*.h5 filter=lfs diff=lfs merge=lfs -text
-*.joblib filter=lfs diff=lfs merge=lfs -text
-*.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
-*.model filter=lfs diff=lfs merge=lfs -text
-*.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
-*.onnx filter=lfs diff=lfs merge=lfs -text
-*.ot filter=lfs diff=lfs merge=lfs -text
-*.parquet filter=lfs diff=lfs merge=lfs -text
-*.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
-*.pt filter=lfs diff=lfs merge=lfs -text
-*.pth filter=lfs diff=lfs merge=lfs -text
-*.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
-saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
-*.tflite filter=lfs diff=lfs merge=lfs -text
-*.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
-*.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
-*.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
-MTO[[:space:]]Font/MTO[[:space:]]Getting[[:space:]]Angry.ttf filter=lfs diff=lfs merge=lfs -text
+*.7z filter=lfs diff=lfs merge=lfs -text
+*.arrow filter=lfs diff=lfs merge=lfs -text
+*.bin filter=lfs diff=lfs merge=lfs -text
+*.bz2 filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.ftz filter=lfs diff=lfs merge=lfs -text
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.h5 filter=lfs diff=lfs merge=lfs -text
+*.joblib filter=lfs diff=lfs merge=lfs -text
+*.lfs.* filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.model filter=lfs diff=lfs merge=lfs -text
+*.msgpack filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.onnx filter=lfs diff=lfs merge=lfs -text
+*.ot filter=lfs diff=lfs merge=lfs -text
+*.parquet filter=lfs diff=lfs merge=lfs -text
+*.pb filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.pt filter=lfs diff=lfs merge=lfs -text
+*.pth filter=lfs diff=lfs merge=lfs -text
+*.rar filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+saved_model/**/* filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.tflite filter=lfs diff=lfs merge=lfs -text
+*.tgz filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
+*.xz filter=lfs diff=lfs merge=lfs -text
+*.zip filter=lfs diff=lfs merge=lfs -text
+*.zst filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

README.md

@@ -1,13 +1,13 @@
----
-title: TheEditor
-emoji: 📉
-colorFrom: pink
-colorTo: yellow
-sdk: gradio
-sdk_version: 5.3.0
-app_file: app.py
-pinned: false
-license: apache-2.0
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+---
+title: TheEditor
+emoji: 📉
+colorFrom: pink
+colorTo: yellow
+sdk: gradio
+sdk_version: 5.3.0
+app_file: app.py
+pinned: false
+license: apache-2.0
+---
+
+Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

TransUnet.py

@@ -0,0 +1,320 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from positional_encodings.torch_encodings import PositionalEncoding2D
+
+
+class LayerNorm2D(nn.Module):
+    def __init__(self, embed_dim):
+        super().__init__()
+        self.layer_norm = nn.LayerNorm(embed_dim)
+
+    def forward(self, x):
+        x = x.permute(0, 2, 3, 1)
+        x = self.layer_norm(x)
+        x = x.permute(0, 3, 1, 2)
+        return x
+    
+class Image_Adaptor(nn.Module):
+    def __init__(self, in_channels, adp_channels, dropout=0.1):
+        super().__init__()
+        
+        self.adaptor = nn.Sequential(
+            nn.Conv2d(in_channels, adp_channels // 4, kernel_size=4, padding='same'),
+            LayerNorm2D(adp_channels // 4),
+            nn.GELU(),
+            nn.Conv2d(adp_channels // 4, adp_channels // 4, kernel_size=2, padding='same'),
+            LayerNorm2D(adp_channels // 4),
+            nn.GELU(),
+            nn.Conv2d(adp_channels // 4, adp_channels, kernel_size=2, padding='same')
+        )
+        self.dropout = nn.Dropout(dropout)
+        
+    def forward(self, images):
+        """
+        input:  [N, in_channels, H, W]
+        output: [N, apd_channels, H, W]
+        """
+        adapt_imgs = self.adaptor(images)
+        return self.dropout(adapt_imgs)
+    
+class Positional_Encoding(nn.Module):
+    def __init__(self, adp_channels):
+        super().__init__()
+        self.pe = PositionalEncoding2D(adp_channels)
+        
+    def forward(self, adapt_imgs):
+        """
+        input:  [N, apd_channels, H, W]
+        output: [N, apd_channels, H, W]
+        """
+        x = adapt_imgs.permute(0, -2, -1, -3)
+        encode = self.pe(x)
+        encode = encode.permute(0, -1, -3, -2)
+        return encode
+    
+class GeGLU(nn.Module):
+    def __init__(self, emb_channels, ffn_size):
+        super().__init__()
+        self.wi_0 = nn.Linear(emb_channels, ffn_size, bias=False)
+        self.wi_1 = nn.Linear(emb_channels, ffn_size, bias=False)
+        self.act = nn.GELU()
+
+    def forward(self, x):
+        x_gelu = self.act(self.wi_0(x))
+        x_linear = self.wi_1(x)
+        x = x_gelu * x_linear
+        return x
+
+class Feed_Forward(nn.Module):
+    def __init__(self, in_channels, ffw_channels, dropout=0.1):
+        super().__init__()
+        
+        self.ln1 = GeGLU(in_channels, ffw_channels)
+        self.dropout = nn.Dropout(dropout)
+        self.ln2 = GeGLU(ffw_channels, in_channels)
+        
+    def forward(self, x):
+        '''
+        input:  [N, H, W, channels]
+        output: [N, H, W, channels]
+        '''
+        x = self.ln1(x)
+        x = self.dropout(x)
+        x = self.ln2(x)
+        return x
+    
+class MultiHeadAttention(nn.Module):
+    def __init__(self, channels, num_attn_heads, dropout=0.1):
+        super().__init__()
+        
+        self.head_size = num_attn_heads
+        self.channels = channels
+        self.attn_size = channels // num_attn_heads
+        self.scale = self.attn_size ** -0.5
+        assert num_attn_heads * self.attn_size == channels, "Input channels of attention must divisible by number of attention head!"
+        
+        self.lq = nn.Linear(channels, self.head_size*self.attn_size, bias=False)
+        self.lk = nn.Linear(channels, self.head_size*self.attn_size, bias=False)
+        self.lv = nn.Linear(channels, self.head_size*self.attn_size, bias=False)
+        self.lout = nn.Linear(self.head_size*self.attn_size, channels, bias=False)
+        self.dropout = nn.Dropout(dropout)
+        
+    def forward(self, q, k, v):
+        '''
+        input:  [N, H, W, channels] cho cả 3 cái q, k, v
+        output: [N, H, W, channels]
+        '''
+        bz, H, W, C = q.shape
+        
+        # Duỗi ảnh ra trước        
+        q = q.view(bz, -1, C) # [N, H*W, C]
+        k = k.view(bz, -1, C) # [N, H*W, C]
+        v = v.view(bz, -1, C) # [N, H*W, C]
+        
+        q = self.lq(q).view(bz, -1, self.head_size, self.attn_size) # [N, H*W, hz, az]
+        k = self.lk(k).view(bz, -1, self.head_size, self.attn_size) # [N, H*W, hz, az]
+        v = self.lv(v).view(bz, -1, self.head_size, self.attn_size) # [N, H*W, hz, az]
+        
+        q = q.transpose(1, 2)                   # [N, hz, H*W, az]
+        k = k.transpose(1, 2).transpose(-1, -2) # [N, hz, az, H*W]
+        v = v.transpose(1, 2)                   # [N, hz, H*W, az]
+        
+        q *= self.scale
+        
+        x = torch.matmul(q, k) # [N, hz, H*W, H*W]
+        x = torch.softmax(x, dim=-1)
+        x = self.dropout(x)
+        x = x.matmul(v) # [N, hz, H*W, az]
+        
+        x = x.transpose(1, 2).contiguous() # [N, H*W, hz, az]
+        x = x.view(bz, -1, C) # [N, H*W, C]
+        x = x.view(bz, H, W, C) # [N, H, W, C]
+        
+        x = self.lout(x) # [N, H, W, C]
+        
+        return x
+    
+class Transformer_Encoder_Layer(nn.Module):
+    def __init__(self, channels, num_attn_heads, ffw_channels, dropout=0.1):
+        super().__init__()
+        
+        self.attn_norm = nn.LayerNorm(channels)
+        self.attn_layer = MultiHeadAttention(channels, num_attn_heads, dropout)
+        self.attn_dropout = nn.Dropout(dropout)
+        
+        self.ffw_norm = nn.LayerNorm(channels)
+        self.ffw_layer = Feed_Forward(channels, ffw_channels, dropout)
+        self.ffw_dropout = nn.Dropout(dropout)
+        
+    def forward(self, adp_pos_imgs):
+        """
+        input:  [N, H, W, channels]
+        output: [N, H, W, channels]
+        """
+        _x = adp_pos_imgs
+        x = self.attn_norm(adp_pos_imgs)
+        x = self.attn_layer(x, x, x)
+        x = self.attn_dropout(x)
+        x = x + _x
+        
+        _x = x
+        x = self.ffw_norm(x)
+        x = self.ffw_layer(x)
+        x = self.ffw_dropout(x)
+        x = x + _x
+        return x
+    
+class Transformer_Encoder(nn.Module):
+    def __init__(self, in_channels, out_channels, num_layers, num_attn_heads, ffw_channels, dropout=0.1):
+        super().__init__()
+        
+        self.encoder_layers = nn.ModuleList([
+            Transformer_Encoder_Layer(in_channels, num_attn_heads, ffw_channels, dropout) for _ in range(num_layers)
+        ])
+        self.linear = nn.Linear(in_channels, out_channels)
+        self.last_norm = LayerNorm2D(out_channels)
+        self.dropout = nn.Dropout(dropout)
+        
+    def forward(self, adp_pos_imgs):
+        """
+        input:  [N, in_channels, H, W]
+        output: [N, out_channels, H, W]
+        """
+        x = adp_pos_imgs.permute(0, -2, -1, -3) # [N, H, W, in_channels]
+        
+        for layer in self.encoder_layers:
+            x = layer(x)
+            
+        x = self.linear(x) # [N, H, W, out_channels]
+        x = x.permute(0, -1, -3, -2)
+        x = self.last_norm(x)
+        out = self.dropout(x)
+        return out
+    
+class Double_Conv(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        
+        self.double_conv = nn.Sequential(
+            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True),
+            nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
+            nn.BatchNorm2d(out_channels),
+            nn.ReLU(inplace=True)
+        )
+        
+    def forward(self, X):
+        """
+        input:  [N, in_channels, H, W]
+        output: [N, out_channels, H//2, W//2]
+        """
+        return self.double_conv(X)
+    
+class Down(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+        
+        self.down = nn.Sequential(
+            nn.MaxPool2d(2),
+            Double_Conv(in_channels, out_channels)
+        )
+        
+    def forward(self, X):
+        """
+        input:  [N, in_channels, H, W]
+        output: [N, out_channels, H//2, W//2]
+        """
+        return self.down(X)
+    
+class Up(nn.Module):
+    def __init__(self, in_channels, out_channels):
+        super().__init__()
+    
+        self.up = nn.ConvTranspose2d(in_channels, in_channels//2, kernel_size=2, stride=2)
+        self.conv = Double_Conv(in_channels, out_channels)
+    
+    def forward(self, X1, X2):
+        """
+        input:  X1 : [N, in_channels,      H // 2, W // 2]
+                X2 : [N, in_channels // 2, H,      W]
+        output: X  : [N, out_channels,     H,      W]
+        """
+        X1 = self.up(X1)
+        
+        diffY = X2.shape[-2] - X1.shape[-2]
+        diffX = X2.shape[-1] - X1.shape[-1]
+        
+        pad_top = diffY // 2
+        pad_bottom = diffY - pad_top
+        pad_left = diffX // 2
+        pad_right = diffX - pad_left
+        
+        X1 = F.pad(X1, (pad_left, pad_right, pad_top, pad_bottom))
+        
+        X = torch.cat((X2, X1), dim=-3)
+        return self.conv(X)
+    
+class Out_Conv(nn.Module):
+    def __init__(self, adp_channels, out_channels):
+        super().__init__()
+        
+        self.out_conv = nn.Conv2d(adp_channels, out_channels, kernel_size=1)
+        
+    def forward(self, X):
+        return self.out_conv(X)
+    
+class Trans_UNet(nn.Module):
+    def __init__(self, 
+                in_channels, 
+                adp_channels, 
+                out_channels, 
+                trans_num_layers=5, 
+                trans_num_attn_heads=8, 
+                trans_ffw_channels=1024,
+                dropout=0.1):
+        super().__init__()
+        
+        self.img_adaptor = Image_Adaptor(in_channels, adp_channels, dropout)
+        self.pos_encoding = Positional_Encoding(adp_channels)
+        
+        self.down1 = Down(adp_channels * 1, adp_channels * 2)
+        self.down2 = Down(adp_channels * 2, adp_channels * 4)
+        self.down3 = Down(adp_channels * 4, adp_channels * 8)
+        self.down4 = Down(adp_channels * 8, adp_channels * 16)
+        self.down5 = Down(adp_channels * 16, adp_channels * 32)
+        
+        self.trans_encoder = Transformer_Encoder(adp_channels * 32, adp_channels * 32, trans_num_layers, trans_num_attn_heads, trans_ffw_channels, dropout)
+        
+        self.up5 = Up(adp_channels * 32, adp_channels * 16)
+        self.up4 = Up(adp_channels * 16, adp_channels * 8)
+        self.up3 = Up(adp_channels * 8, adp_channels * 4)
+        self.up2 = Up(adp_channels * 4, adp_channels * 2)
+        self.up1 = Up(adp_channels * 2, adp_channels * 1)
+        
+        self.out_conv = Out_Conv(adp_channels, out_channels)
+        self.sigmoid = nn.Sigmoid()
+        
+    def forward(self, images):
+        adp_imgs = self.img_adaptor(images)
+        pos_enc  = self.pos_encoding(adp_imgs)
+        adp_imgs += pos_enc        
+        
+        d1 = self.down1(adp_imgs)
+        d2 = self.down2(d1)
+        d3 = self.down3(d2)
+        d4 = self.down4(d3)
+        d5 = self.down5(d4)
+        
+        x = self.trans_encoder(d5)
+        
+        u5 = self.up5(x, d4)
+        u4 = self.up4(u5, d3)
+        u3 = self.up3(u4, d2)
+        u2 = self.up2(u3, d1)
+        u1 = self.up1(u2, adp_imgs)
+        
+        x = self.out_conv(u1)
+        out = self.sigmoid(x)
+        return out

TransUnet_Config.py

@@ -0,0 +1,7 @@
+in_channels = 3
+adp_channels = 32
+out_channels = 1
+trans_num_layers = 5
+trans_num_attn_heads = 8
+trans_ffw_channels = 512
+dropout = 0.1

app.py

@@ -0,0 +1,428 @@
+import cv2
+import gradio as gr
+import os
+from edit_func import *
+from TransUnet import Trans_UNet
+import TransUnet_Config as config2
+from huggingface_hub import hf_hub_download
+from googletrans import Translator
+import random
+import torch.nn as nn
+import spaces
+
+@spaces.GPU
+class DTM(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        self.detect_text_model = Trans_UNet(
+            config2.in_channels, config2.adp_channels, config2.out_channels, 
+            config2.trans_num_layers, config2.trans_num_attn_heads, config2.trans_ffw_channels,
+            config2.dropout
+        ).to(self.device)
+        self.repo_name = 'SS3M/detect-text-model'
+        files = ['detect-text-v3-0.pt', 'detect-text-v3-1.pt', 
+                      'detect-text-v3-2.pt', 'detect-text-v3-3.pt', 
+                      'detect-text-v3-4.pt', 'detect-text-v3-5.pt', 
+                      'detect-text-v3-6.pt', 'detect-text-v3-7.pt']
+        self.files = []
+        for file in files:
+            self.files.append(hf_hub_download(repo_id=self.repo_name, filename=file))
+    
+    def forward(self, X):
+        X = X.to(self.device)
+        N, C, H, W = X.shape
+        result = torch.zeros((N, 1, H, W))
+        for file in self.files:
+            model_path = file
+            best_model_state = torch.load(
+                model_path, 
+                weights_only=True, 
+                map_location=self.device
+            )
+            self.detect_text_model.load_state_dict(best_model_state)
+            result += self.detect_text_model(X)
+        result /= len(self.files)
+        return result
+
+@spaces.GPU
+class DWBM(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.device = 'cuda' if torch.cuda.is_available() else 'cpu'
+        self.detect_wordball_model = Trans_UNet(
+            config2.in_channels, config2.adp_channels, config2.out_channels, 
+            config2.trans_num_layers, config2.trans_num_attn_heads, config2.trans_ffw_channels,
+            config2.dropout
+        ).to(self.device)
+        self.repo_name = 'SS3M/detect-wordball-model'
+        files = ['detect-text-v3-0.pt', 'detect-text-v3-1.pt', 
+                      'detect-text-v3-2.pt', 'detect-text-v3-3.pt', 
+                      'detect-text-v3-4.pt', 'detect-text-v3-5.pt', 
+                      'detect-text-v3-6.pt', 'detect-text-v3-7.pt']
+        self.files = []
+        for file in files:
+            self.files.append(hf_hub_download(repo_id=self.repo_name, filename=file))
+    
+    def forward(self, X):
+        X = X.to(self.device)
+        N, C, H, W = X.shape
+        result = torch.zeros((N, 1, H, W))
+        for file in self.files:
+            model_path = file
+            best_model_state = torch.load(
+                model_path, 
+                weights_only=True, 
+                map_location=self.device
+            )
+            self.detect_wordball_model.load_state_dict(best_model_state)
+            result += self.detect_wordball_model(X)
+        result /= len(self.files)
+        return result   
+
+detect_text_model = DTM()
+detect_wordball_model = DWBM()
+
+translator = Translator()
+
+def down1(src_img):
+    src_img = cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR)
+    text_msk = create_text_mask(src_img, detect_text_model)
+    wordball_msk = create_wordball_mask(src_img, detect_wordball_model)
+
+    text_positions, areas = get_text_positions(text_msk, text_value=0)
+    rgbs = []
+    for _ in range(len(areas)):
+        rgbs.append((random.randint(0, 255), random.randint(0, 255), random.randint(0, 255)))
+
+    idx = '; '.join(str(i) for i in range(len(areas)))
+    text_positions = '; '.join([', '.join(str(i) for i in pos) for pos in text_positions])
+    areas = '; '.join(str(i) for i in areas)
+    rgbs = '; '.join([', '.join(str(i) for i in rgb) for rgb in rgbs])
+    src_img = cv2.cvtColor(src_img, cv2.COLOR_BGR2RGB)
+    return text_msk*255, wordball_msk*255, idx, text_positions, areas, rgbs, 'Xong'
+
+def idx_txt_change(src_img, idx_txt, pos_txt, rgb_txt):
+    try:
+        src_img2 = cv2.cvtColor(src_img, cv2.COLOR_RGB2BGR)
+        text_positions = pos_txt.split('; ')
+        for idx in range(len(text_positions)):
+            text_positions[idx] = (int(i) for i in text_positions[idx].split(', '))
+        rgbs = rgb_txt.split('; ')
+        for idx in range(len(rgbs)):
+            rgbs[idx] = (int(i) for i in rgbs[idx].split(', '))
+        idxes = [int(idx) for idx in idx_txt.split('; ')]
+
+        for idx, ((min_x, min_y, max_x, max_y), (r, g, b)) in enumerate(zip(text_positions, rgbs)):
+            if idx in idxes:
+                cv2.rectangle(src_img2, (min_x, min_y), (max_x, max_y), (b, g, r), thickness=4)
+        src_img2 = cv2.cvtColor(src_img2, cv2.COLOR_BGR2RGB)
+        return src_img2
+    except:
+        return src_img
+
+def scale_area_change(min_area, max_area, area_txt):
+    areas = [int(area) for area in area_txt.split('; ')]
+    idxes = []
+    for idx, area in enumerate(areas):
+        if min_area <= area <= max_area:
+            idxes.append(idx)
+    idxes = '; '.join(str(i) for i in idxes)
+    return idxes
+
+def position_block_change(X, Y, W, H, ID, pos_txt_value):
+    text_positions = pos_txt_value.split('; ')
+    for idx in range(len(text_positions)):
+        text_positions[idx] = (int(i) for i in text_positions[idx].split(', '))
+    
+    text_positions2 = []
+    for idx, (min_x, min_y, max_x, max_y) in enumerate(text_positions):
+        if idx == ID:
+            text_positions2.append((X, Y, X+W, Y+H))
+        else:
+            text_positions2.append((min_x, min_y, max_x, max_y))
+    text_positions2 = '; '.join([', '.join(str(i) for i in pos) for pos in text_positions2])
+    return text_positions2
+
+def ID_block_change(ID_value, checkbox_value, ID_txt_value):
+    ID_txt_value = [int(i) for i in ID_txt_value.split('; ')]
+    if checkbox_value and ID_value not in ID_txt_value:
+        ID_txt_value.append(ID_value)
+    if not checkbox_value and ID_value in ID_txt_value:
+        ID_txt_value.remove(ID_value)
+    ID_txt_value = sorted(ID_txt_value)
+    ID_txt_value = '; '.join([str(i) for i in ID_txt_value])
+    return ID_txt_value
+
+def down2(src_img_value, txt_mask_value, wordball_mask_value, idx_txt_value, pos_txt_value):
+    src_img_value = cv2.cvtColor(src_img_value, cv2.COLOR_RGB2BGR)
+    text_positions = pos_txt_value.split('; ')
+    for idx in range(len(text_positions)):
+        text_positions[idx] = (int(i) for i in text_positions[idx].split(', '))
+    idxes = [int(i) for i in idx_txt_value.split('; ')]
+    
+    for idx, (min_x, min_y, max_x, max_y) in enumerate(text_positions):
+        if idx not in idxes:
+            txt_mask_value[min_y:max_y+1, min_x:max_x+1] = 255
+    txt_mask_value = txt_mask_value[:, :, 0].astype(np.uint8)
+    non_text_src_img = clear_text(src_img_value, txt_mask_value, wordball_mask_value, text_value=0, non_text_value=255, r=5)
+
+    list_texts = get_list_texts(src_img_value, [tuple(map(int, pos.split(', '))) for idx, pos in enumerate(pos_txt_value.split('; ')) if idx in idxes])
+    list_translated_texts = translate(list_texts, translator)
+    list_fonts = '; '.join(['MTO Astro City.ttf' for _ in range(len(list_translated_texts))])
+    list_sizes = '; '.join(['20' for _ in range(len(list_translated_texts))])
+    list_strokes = '; '.join(['3' for _ in range(len(list_translated_texts))])
+    list_pads = '; '.join(['5' for _ in range(len(list_translated_texts))])
+    list_translated_texts = '; '.join(list_translated_texts)
+    switch = str(random.random())
+
+    return non_text_src_img, list_translated_texts, list_fonts, list_sizes, list_strokes, list_pads, switch, 'Xong'
+
+def text_info_change(non_txt_img_value, translated_txt_value, pos_txt_value, idx_txt_value, font_txt_value, size_txt_value, stroke_txt_value, pad_txt_value):
+    non_txt_img_value = non_txt_img_value.copy()
+    idxes = [int(i) for i in idx_txt_value.split('; ')]
+
+    translated_text_src_img = insert_text(non_txt_img_value, 
+                                          translated_txt_value.split('; '), 
+                                          [tuple(map(int, pos.split(', '))) for idx, pos in enumerate(pos_txt_value.split('; ')) if idx in idxes], 
+                                          font=font_txt_value.split('; '), 
+                                          font_size=[int(i) for i in size_txt_value.split('; ')], 
+                                          pad=[int(i) for i in pad_txt_value.split('; ')], 
+                                          stroke=[int(i) for i in stroke_txt_value.split('; ')])
+    return translated_text_src_img
+
+def value2_change(value, ID2_value, txt_value):
+    txt_value = txt_value.split('; ')
+    
+    txt_value2 = []
+    for idx, text in enumerate(txt_value):
+        if idx == ID2_value:
+            txt_value2.append(str(value))
+        else:
+            txt_value2.append(str(text))
+    txt_value2 = '; '.join(txt_value2)
+    return txt_value2
+
+# Tạo giao diện Gradio
+with gr.Blocks() as demo:
+    # Cấu trúc
+    src_img = gr.Image(type="numpy", label="Upload Image")
+
+    down_bttn_1 = gr.Button("↓", elem_classes="arrow-button")
+
+    with gr.Row():
+        txt_mask = gr.Image(type="numpy", label="Upload Image", visible=True)
+        wordball_mask = gr.Image(type="numpy", label="Upload Image", visible=True)
+    complete = gr.Textbox()
+    with gr.Row():
+        idx_txt = gr.Textbox(label='ID', interactive=False, visible=False)
+        pos_txt = gr.Textbox(label='Pos', interactive=False, visible=False)
+        area_txt = gr.Textbox(label='Area', interactive=False, visible=False)
+        rgb_txt = gr.Textbox(label='rgb', interactive=False, visible=False)
+    with gr.Row():
+        boxed_txt_img = gr.Image(type="numpy", label="Upload Image")
+        with gr.Column() as down_1_column:            
+            @gr.render(inputs=[pos_txt, rgb_txt], triggers=[rgb_txt.change])
+            def create_box(pos_txt_value, rgb_txt_value):
+                text_positions = pos_txt_value.split('; ')
+                for idx in range(len(text_positions)):
+                    text_positions[idx] = (int(i) for i in text_positions[idx].split(', '))
+                rgbs = rgb_txt_value.split('; ')
+                for idx in range(len(rgbs)):
+                    rgbs[idx] = (int(i) for i in rgbs[idx].split(', '))
+
+                elements = []
+                for idx, (min_x, min_y, max_x, max_y) in enumerate(text_positions):
+                    with gr.Group() as box:
+                        r, g, b = rgbs[idx]
+                        with gr.Row():
+                            gr.Markdown(
+                                f"""
+                                <div style="margin-left: 20px; display: flex; align-items: center;">
+                                    <div style="width: 10px; height: 10px; background-color: rgb({r}, {g}, {b}); margin-right: 5px;"></div>
+                                    <span style="font-size: 20px;">Textbox {idx+1}</span>
+                                </div>
+                                """
+                            )
+                            checkbox = gr.Checkbox(value=True, label='', min_width=50, interactive=True)
+                        with gr.Row():
+                            X = gr.Number(label="X", value=min_x, interactive=True)
+                            Y = gr.Number(label="Y", value=min_y, interactive=True)
+                            W = gr.Number(label="W", value=max_x-min_x, interactive=True)
+                            H = gr.Number(label="H", value=max_y-min_y, interactive=True)
+                            ID = gr.Number(label="ID", value=idx, interactive=True, visible=False)
+                            elements.append((X, Y, W, H, ID))
+                            
+                            checkbox.change(
+                                fn=ID_block_change,
+                                inputs=[ID, checkbox, idx_txt],
+                                outputs=idx_txt,
+                                show_progress=False
+                            ).then(
+                                fn=idx_txt_change,
+                                inputs=[src_img, idx_txt, pos_txt, rgb_txt],
+                                outputs=boxed_txt_img,
+                            )
+                            X.change(
+                                fn=position_block_change,
+                                inputs=[X, Y, W, H, ID, pos_txt],
+                                outputs=pos_txt,
+                                show_progress=False
+                            ).then(
+                                fn=idx_txt_change,
+                                inputs=[src_img, idx_txt, pos_txt, rgb_txt],
+                                outputs=boxed_txt_img,
+                                show_progress=False
+                            )
+                            Y.change(
+                                fn=position_block_change,
+                                inputs=[X, Y, W, H, ID, pos_txt],
+                                outputs=pos_txt,
+                                show_progress=False
+                            ).then(
+                                fn=idx_txt_change,
+                                inputs=[src_img, idx_txt, pos_txt, rgb_txt],
+                                outputs=boxed_txt_img,
+                                show_progress=False
+                            )
+                            W.change(
+                                fn=position_block_change,
+                                inputs=[X, Y, W, H, ID, pos_txt],
+                                outputs=pos_txt,
+                                show_progress=False
+                            ).then(
+                                fn=idx_txt_change,
+                                inputs=[src_img, idx_txt, pos_txt, rgb_txt],
+                                outputs=boxed_txt_img,
+                                show_progress=False
+                            )
+                            H.change(
+                                fn=position_block_change,
+                                inputs=[X, Y, W, H, ID, pos_txt],
+                                outputs=pos_txt,
+                                show_progress=False
+                            ).then(
+                                fn=idx_txt_change,
+                                inputs=[src_img, idx_txt, pos_txt, rgb_txt],
+                                outputs=boxed_txt_img,
+                                show_progress=False
+                            )
+    down_bttn_2 = gr.Button("↓", elem_classes="arrow-button")
+    
+    non_txt_img = gr.Image(type="numpy", label="Upload Image", visible=False)
+    complete2 = gr.Textbox()
+    with gr.Row():
+        translated_txt = gr.Textbox(label='translated', interactive=False, visible=False)
+        font_txt = gr.Textbox(label='font', interactive=False, visible=False)
+        size_txt = gr.Textbox(label='size', interactive=False, visible=False)
+        stroke_txt = gr.Textbox(label='stroke', interactive=False, visible=False)
+        pad_txt = gr.Textbox(label='pad', interactive=False, visible=False)
+        switch_txt = gr.Textbox(label='switch', value='1', interactive=False, visible=False)
+    with gr.Row():
+        boxed_inserted_non_txt_img = gr.Image(type="numpy", label="Upload Image")
+        with gr.Column():
+            @gr.render(inputs=[translated_txt, font_txt, size_txt, stroke_txt, pad_txt], triggers=[switch_txt.change])
+            def create_box2(translated_txt_value, font_txt_value, size_txt_value, stroke_txt_value, pad_txt_value):
+                translated_txt_value = translated_txt_value.split('; ')
+                font_txt_value = font_txt_value.split('; ')
+                size_txt_value = size_txt_value.split('; ')
+                stroke_txt_value = stroke_txt_value.split('; ')
+                pad_txt_value = pad_txt_value.split('; ')
+
+                elements = []
+                for idx in range(len(font_txt_value)):
+                    with gr.Group():
+                        gr.Markdown(
+                            f"""
+                            <div style="margin-left: 20px; display: flex; align-items: center;">
+                                <div style="width: 10px; height: 10px; background-color: rgb(255, 255, 255); margin-right: 5px;"></div>
+                                <span style="font-size: 20px;">Text box {idx}</span>
+                            </div>
+                            """
+                        )
+                        translated_text_box = gr.Textbox(label="Translate", value=translated_txt_value[idx], interactive=True)
+                        with gr.Row():
+                            font = gr.Dropdown(choices=os.listdir('MTO Font'), label="Phông chữ", value=font_txt_value[idx], interactive=True, scale=7)
+                            size = gr.Number(label="Size", value=int(size_txt_value[idx]), interactive=True, minimum=1)
+                            stroke = gr.Number(label="Stroke", value=int(stroke_txt_value[idx]), interactive=True, minimum=0, maximum=5)
+                            pad = gr.Number(label="Pad", value=int(pad_txt_value[idx]), interactive=True, minimum=1, maximum=10)
+                            ID2 = gr.Number(label="ID", value=int(idx), interactive=True, visible=False)
+
+                        translated_text_box.submit(
+                            fn=value2_change,
+                            inputs=[translated_text_box, ID2, translated_txt],
+                            outputs=translated_txt,
+                            show_progress=False
+                        ).then(
+                            fn=text_info_change,
+                            inputs=[non_txt_img, translated_txt, pos_txt, idx_txt, font_txt, size_txt, stroke_txt, pad_txt],
+                            outputs=boxed_inserted_non_txt_img,
+                        )
+                        font.change(
+                            fn=value2_change,
+                            inputs=[font, ID2, font_txt],
+                            outputs=font_txt,
+                            show_progress=False
+                        ).then(
+                            fn=text_info_change,
+                            inputs=[non_txt_img, translated_txt, pos_txt, idx_txt, font_txt, size_txt, stroke_txt, pad_txt],
+                            outputs=boxed_inserted_non_txt_img,
+                        )
+                        size.change(
+                            fn=value2_change,
+                            inputs=[size, ID2, size_txt],
+                            outputs=size_txt,
+                            show_progress=False
+                        ).then(
+                            fn=text_info_change,
+                            inputs=[non_txt_img, translated_txt, pos_txt, idx_txt, font_txt, size_txt, stroke_txt, pad_txt],
+                            outputs=boxed_inserted_non_txt_img,
+                        )
+                        stroke.change(
+                            fn=value2_change,
+                            inputs=[stroke, ID2, stroke_txt],
+                            outputs=stroke_txt,
+                            show_progress=False
+                        ).then(
+                            fn=text_info_change,
+                            inputs=[non_txt_img, translated_txt, pos_txt, idx_txt, font_txt, size_txt, stroke_txt, pad_txt],
+                            outputs=boxed_inserted_non_txt_img,
+                        )
+                        pad.change(
+                            fn=value2_change,
+                            inputs=[pad, ID2, pad_txt],
+                            outputs=pad_txt,
+                            show_progress=False
+                        ).then(
+                            fn=text_info_change,
+                            inputs=[non_txt_img, translated_txt, pos_txt, idx_txt, font_txt, size_txt, stroke_txt, pad_txt],
+                            outputs=boxed_inserted_non_txt_img,
+                        )
+
+    # Css
+    demo.css = """
+    .arrow-button {
+        font-size: 40px; /* Kích thước font */
+    }
+    .group-elem {
+        height: 70px;
+    }
+    """
+
+    # Điều khiển
+    down_bttn_1.click(
+        fn=down1,
+        inputs=src_img,
+        outputs=[txt_mask, wordball_mask, idx_txt, pos_txt, area_txt, rgb_txt, complete],
+    )
+    down_bttn_2.click(
+        fn=down2,
+        inputs=[src_img, txt_mask, wordball_mask, idx_txt, pos_txt],
+        outputs=[non_txt_img, translated_txt, font_txt, size_txt, stroke_txt, pad_txt, switch_txt, complete2],
+    ).then(
+        fn=text_info_change,
+        inputs=[non_txt_img, translated_txt, pos_txt, idx_txt, font_txt, size_txt, stroke_txt, pad_txt],
+        outputs=boxed_inserted_non_txt_img,
+    )
+
+demo.launch()

edit_func.py

@@ -0,0 +1,311 @@
+import torch
+import cv2
+import pytesseract
+from PIL import Image, ImageDraw, ImageFont
+from collections import deque
+import numpy as np
+import os
+
+# pytesseract.pytesseract.tesseract_cmd = 'Tesseract\\tesseract.exe'
+
+def get_full_img_path(src_dir):
+    """
+    input:  Đường dẫn đền folder chứa ảnh
+    output: Danh sách tên của tất cả các ảnh
+    """
+    list_img_names = []
+    for dirname, _, filenames in os.walk(src_dir):
+        for filename in filenames:
+            path = os.path.join(dirname, filename).replace(src_dir, '')
+            if path[0] == '/':
+                path = path[1:]
+            list_img_names.append(path)
+    return list_img_names
+
+
+def create_text_mask(src_img, detect_text_model, kernel_size=5, iterations=3):
+    """
+    input:  Ảnh gốc, để dưới định dạng là np.array, shape: [H, W, C]
+    output: Mask đánh dấu text trong ảnh gốc, 0 là chữ, 1 là nền; shape: [H, W]
+    """
+    img = torch.from_numpy(src_img).to(torch.uint8).to(detect_text_model.device)
+    imgT = (img / 255).unsqueeze(0).permute(0, -1, -3, -2)
+    
+    detect_text_model.eval()
+    with torch.no_grad():
+        result = detect_text_model(imgT).squeeze()
+    result = (result >= 0.5).detach().cpu().numpy()
+    
+    mask = ((1-result) * 255).astype(np.uint8)
+
+    kernel = np.ones((kernel_size, kernel_size), np.uint8)
+    mask = cv2.erode(mask, kernel, iterations=iterations)
+    mask = cv2.dilate(mask, kernel, iterations=2*iterations)
+    mask = cv2.erode(mask, kernel, iterations=iterations)
+
+    mask = (1 - mask // 255).astype(np.uint8)
+    return mask
+
+
+def create_wordball_mask(src_img, detect_wordball_model, kernel_size=5, iterations=3):
+    """
+    input:  Ảnh gốc, để dưới định dạng là np.array, shape: [H, W, C]
+    output: Mask đánh dấu text trong ảnh gốc, 0 là chữ, 1 là nền; shape: [H, W]
+    """
+    img = torch.from_numpy(src_img).to(torch.uint8).to(detect_wordball_model.device)
+    imgT = (img / 255).unsqueeze(0).permute(0, -1, -3, -2)
+    
+    detect_wordball_model.eval()
+    with torch.no_grad():
+        result = detect_wordball_model(imgT).squeeze()
+    result = (result >= 0.5).detach().cpu().numpy()
+    
+    mask = ((1-result) * 255).astype(np.uint8)
+
+    kernel = np.ones((kernel_size, kernel_size), np.uint8)
+    mask = cv2.erode(mask, kernel, iterations=iterations)
+    mask = cv2.dilate(mask, kernel, iterations=2*iterations)
+    mask = cv2.erode(mask, kernel, iterations=iterations)
+
+    mask = (1 - mask // 255).astype(np.uint8)
+    return mask
+
+
+def clear_text(src_img, text_msk, wordball_msk, text_value=0, non_text_value=1, r=5):
+    """
+    input:  src_img: Ảnh gốc, để dưới định dạng là np.array, shape: [H, W, C]
+            text_msk: Mask đánh dấu text trong ảnh gốc; shape: [H, W]
+            text_value: Giá trị mà trong mặt nạ nó là text
+            non_text_value: Giá trị mà trong mặt nạ nó là nền
+            r: Bán kính để sử dụng cho việc xoá text và vẽ lại phần bị xoá
+    output: Ảnh sau khi xoá text, để dưới định dạng là np.array, shape: [H, W, C]
+    """
+    MAX = max(text_value, non_text_value)
+    MIN = min(text_value, non_text_value)
+    
+    scale_text_value = (text_value - MIN) / (MAX - MIN)
+    scale_non_text_value = (non_text_value - MIN) / (MAX - MIN)
+    
+    text_msk[text_msk==text_value] = scale_text_value
+    text_msk[text_msk==non_text_value] = scale_non_text_value
+    
+    wordball_msk[wordball_msk==text_value] = scale_text_value
+    wordball_msk[wordball_msk==non_text_value] = scale_non_text_value
+    
+    if scale_text_value == 0:
+        text_msk = 1 - text_msk
+        wordball_msk = 1 - wordball_msk
+    text_msk = text_msk * 255
+    
+    remove_txt = cv2.inpaint(src_img, text_msk, r, cv2.INPAINT_TELEA)
+    remove_wordball = remove_txt.copy()
+    remove_wordball[wordball_msk==1] = 255
+    
+    return remove_wordball
+
+
+def dfs(grid, y, x, visited, value):
+    """
+    Thuật toán tìm miền liên thông, xem thêm về đồ thị nếu không biết nó là gì
+    Output: Một HCN bao phủ miền liên thông + Diện tích của miền liên thông
+    """
+    max_y, max_x = y, x
+    min_y, min_x = y+1, x+1
+    area = 0
+    
+    stack = deque([(y, x)])
+    while stack:
+        y, x = stack.pop()
+
+        max_x = max(max_x, x)
+        max_y = max(max_y, y)
+        min_x = min(min_x, x)
+        min_y = min(min_y, y)
+
+        if (y, x) not in visited:
+            visited.add((y, x))
+            area += 1
+            # Kiểm tra các ô liền kề
+            for dx, dy in [(-1, 0), (1, 0), (0, -1), (0, 1), (-1, -1), (-1, 1), (1, -1), (1, 1)]:
+                nx, ny = x + dx, y + dy
+                if 0 <= ny < grid.shape[0] and 0 <= nx < grid.shape[1] and grid[ny, nx] == value and (ny, nx) not in visited:
+                    stack.append((ny, nx))
+
+    return (min_x, min_y, max_x, max_y), area
+
+
+def find_clusters(grid, value):
+    """
+    Thuật toán tìm danh sách các miền liên thông
+    """
+    visited = set()
+    clusters = []
+    areas = []
+    
+    for y in range(grid.shape[0]):
+        for x in range(grid.shape[1]):
+            if grid[y, x] == value and (y, x) not in visited:
+                cluster, area = dfs(grid, y, x, visited, value)
+                clusters.append(cluster)
+                areas.append(area)
+    
+    return clusters, areas
+
+def get_text_positions(text_msk, text_value=0):
+    """
+    input:  text_msk: Mask đánh dấu text trong ảnh gốc; shape: [H, W]
+            text_value: Giá trị mà trong mặt nạ nó là text
+            min_area: Giả trị tối thiểu của vùng có thể có text
+    output: Danh sách các cùng chứa text, định dạng (min_x, min_y, max_x, max_y)
+    """
+    
+    clusters, areas = find_clusters(text_msk, value=text_value)
+    return clusters, areas
+
+def filter_text_positions(clusters, areas, min_area=1200, max_area=10000):
+    clusters = clusters[(areas >= min_area) & (areas <= max_area)]
+    return clusters
+
+
+def get_list_texts(src_img, text_positions, lang='eng'):
+    """
+    input:  src_img: Ảnh gốc, để dưới định dạng là np.array, shape: [H, W, C]
+            text_positions: Danh sách các cùng chứa text, định dạng (min_x, min_y, max_x, max_y)
+            lang: Ngôn ngữ của text
+    output: Danh sách các câu text
+    """
+    list_texts = []
+    for idx, (min_x, min_y, max_x, max_y) in enumerate(text_positions):
+        crop_img = src_img[min_y:max_y+1, min_x:max_x+1]
+        img_rgb = cv2.cvtColor(crop_img, cv2.COLOR_BGR2RGB)
+        img = Image.fromarray(img_rgb)
+        text = pytesseract.image_to_string(img, lang=lang).replace('\n', ' ').strip()
+        while '  ' in text:
+            text = text.replace('  ', ' ')
+        list_texts.append(text)
+    return list_texts
+
+
+def translate(list_texts, translator):
+    translated_texts = []
+    for text in list_texts:
+        if not text:
+            text = 'a'
+        translated_text = translator.translate(text, src='en', dest='vi').text
+        translated_texts.append(translated_text)
+    return translated_texts
+
+
+def add_centered_multiline_text(image, text, box, font_path="arial.ttf", font_size=36, pad=5, text_color=0):
+    # Mở ảnh
+    draw = ImageDraw.Draw(image)
+
+    # Giải nén box (min_x, min_y, max_x, max_y)
+    min_x, min_y, max_x, max_y = box
+
+    # Tạo font
+    font = ImageFont.truetype(font_path, font_size)
+
+    # Chia văn bản thành nhiều dòng nếu cần
+    wrapped_lines = wrap_text(text, font, draw, max_x - min_x)
+
+    # Tính chiều cao của tất cả các dòng cộng lại
+    total_text_height = sum(get_text_height(line, draw, font) for line in wrapped_lines)
+    
+    # Tính toạ độ y bắt đầu để căn giữa theo chiều dọc
+    start_y = min_y + (max_y - min_y - total_text_height) // 2
+
+    # Vẽ từng dòng và căn giữa theo chiều ngang
+    current_y = start_y
+    for line in wrapped_lines:
+        text_width, text_height = get_text_dimensions(line, draw, font)
+        text_x = min_x + (max_x - min_x - text_width) // 2  # Căn giữa theo chiều ngang
+        draw.text((text_x, current_y), line, fill=text_color, font=font)
+        current_y += text_height + pad # Di chuyển y xuống để vẽ dòng tiếp theo
+
+    # Lưu ảnh mới
+    return image
+
+def get_text_dimensions(text, draw, font):
+    """Trả về (width, height) của văn bản."""
+    bbox = draw.textbbox((0, 0), text, font=font)
+    width = bbox[2] - bbox[0]
+    height = bbox[3] - bbox[1]
+    return width, height
+
+def get_text_height(text, draw, font):
+    """Trả về chiều cao của văn bản."""
+    _, _, _, height = draw.textbbox((0, 0), text, font=font)
+    return height
+
+def wrap_text(text, font, draw, max_width):
+    """Chia văn bản thành nhiều dòng dựa trên chiều rộng tối đa."""
+    words = text.split()
+    lines = []
+    current_line = ""
+
+    for word in words:
+        # Thử thêm từ vào dòng hiện tại
+        test_line = f"{current_line} {word}".strip()
+        test_width, _ = get_text_dimensions(test_line, draw, font)
+
+        if test_width <= max_width:
+            current_line = test_line
+        else:
+            # Nếu quá rộng, lưu dòng hiện tại và bắt đầu dòng mới
+            lines.append(current_line)
+            current_line = word
+
+    # Thêm dòng cuối cùng
+    if current_line:
+        lines.append(current_line)
+
+    return lines
+
+def insert_text(non_text_src_img, list_translated_texts, text_positions, font=['MTO Astro City.ttf'], font_size=[20], pad=[5], text_color=0, stroke=[3]):
+    # Copy ảnh không chữ
+    img_bgr = non_text_src_img.copy()
+    
+    # Thêm text vào măt nạ 1
+    for idx, text in enumerate(list_translated_texts):
+        # Tạo mặt nạ trắng
+        mask1 = Image.new("L", img_bgr.shape[:2][::-1], 255)
+        mask2 = Image.new("L", img_bgr.shape[:2][::-1], 255)
+        mask1 = add_centered_multiline_text(mask1, text, text_positions[idx], f'MTO Font/{font[idx]}', font_size[idx], pad=pad[idx], text_color=text_color)
+    
+        # Chuyển ảnh từ PIL sang cv2
+        mask1 = (np.array(mask1) >= 127).astype(np.uint8) * 255
+        mask1 = cv2.cvtColor(mask1, cv2.COLOR_RGB2BGR)
+        
+        if stroke[idx] > 0:
+            mask2 = np.array(mask2).astype(np.uint8)
+            mask2 = cv2.cvtColor(mask2, cv2.COLOR_RGB2BGR)
+            
+            mask2 = mask2 - mask1
+            kernel = np.ones((stroke[idx]+1, stroke[idx]+1), np.uint8)
+            mask2 = cv2.dilate(mask2, kernel, iterations=1)
+            img_bgr[mask2==255] = 255
+        
+        img_bgr[mask1==text_color] = text_color
+    return img_bgr
+
+
+def save_img(path, translated_text_src_img):
+    """
+    input:  path: Đường dẫn đến ảnh gốc ban đầu
+            translated_text_src_img: Ảnh sau khi được dịch
+    output: Ảnh sau dịch được lưu lại, trong tên có thêm "translated-"
+    """
+    dot = path.rfind('.')
+    last_slash = -1
+    if '/' in path:
+        last_slash = path.rfind('/')
+        
+    ext = path[dot:]
+    parent_path = path[:last_slash+1]
+    name = path[last_slash+1:dot]
+    
+    if parent_path and not os.path.exists(parent_path):
+        os.mkdir(parent_path)    
+    cv2.imwrite(f'{parent_path}translated-{name}{ext}', translated_text_src_img)
+    print(f'Image saved at {parent_path}translated-{name}{ext}')

requirements.txt

@@ -0,0 +1,13 @@
+chardet==3.0.4
+googletrans==4.0.0rc1
+h2==3.2.0
+hstspreload==2024.10.1
+opencv-python==4.10.0.84
+pip==24.2
+positional-encodings==6.0.3
+pytesseract==0.3.13
+rfc3986==1.5.0
+setuptools==75.1.0
+spaces==0.30.4
+torch==2.5.0
+wheel==0.44.0