| 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
|
|
|
|
|
| with gr.Blocks() as demo:
|
|
|
| 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,
|
| )
|
|
|
|
|
| demo.css = """
|
| .arrow-button {
|
| font-size: 40px; /* Kích thước font */
|
| }
|
| .group-elem {
|
| height: 70px;
|
| }
|
| """
|
|
|
|
|
| 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()
|
|
|
