edit_func.py

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}')