app.py

import gradio as gr
import numpy as np
import math
import torch

from PIL import Image, ImageDraw
from rect_main import docscanner_rec, load_docscanner_model
from data_utils.image_utils import unwarp, mask2point, get_corner, _rotate_90_degrees
from config import Config


config = Config()
cuda = torch.device("cuda" if torch.cuda.is_available() else "cpu")


docscanner = load_docscanner_model(
    cuda, path_l=config.get_rec_model_path, path_m=config.get_seg_model_path
)

# 좌표를 초기화하는 함수
def reset_points(image, state):
    state = []
    return image, state

def cutting_image(image, state):
    min_x = min(point[0] for point in state)
    max_x = max(point[0] for point in state)
    min_y = min(point[1] for point in state)
    max_y = max(point[1] for point in state)
    
    cutted_image = image[min_y:max_y, min_x:max_x]
    state = []
    return cutted_image, cutted_image, state

def rotate_image(image):
    rotated_image = _rotate_90_degrees(image)
    state = []
    return rotated_image, state

def reset_image(image, state):
    out_image, msk_np = docscanner_rec(image, docscanner)
    state = list(get_corner(mask2point(mask=msk_np)))
    
    img = Image.fromarray(image)
    area = image.shape[0]*image.shape[1]
    radius=max(5, round(area**0.5 / 120))
    # 좌표가 최소 3개 이상일 때만 폴리곤 그리기
    draw = ImageDraw.Draw(img)
    for pt in state:
        left_up_point = (pt[0] - radius, pt[1] - radius)
        right_down_point = (pt[0] + radius, pt[1] + radius)
        draw.ellipse([left_up_point, right_down_point], outline="black", fill="red")
    
    center = (sum(p[0] for p in state) / len(state), sum(p[1] for p in state) / len(state))
    # 각도에 따라 점들을 정렬
    sorted_points = sorted(state, key=lambda p: calculate_angle(p, center))
    draw.polygon(sorted_points, outline="red", fill=None, width=round(radius/2))
    
    return img, state

def auto_point_detect(image):
    out_image, msk_np = docscanner_rec(image, docscanner, cuda)
    state = list(get_corner(mask2point(mask=msk_np)))
    
    img = Image.fromarray(image)
    area = image.shape[0]*image.shape[1]
    radius=max(5, round(area**0.5 / 120))
    # 좌표가 최소 3개 이상일 때만 폴리곤 그리기
    draw = ImageDraw.Draw(img)
    for pt in state:
        left_up_point = (pt[0] - radius, pt[1] - radius)
        right_down_point = (pt[0] + radius, pt[1] + radius)
        draw.ellipse([left_up_point, right_down_point], outline="black", fill="red")
    
    center = (sum(p[0] for p in state) / len(state), sum(p[1] for p in state) / len(state))
    # 각도에 따라 점들을 정렬
    sorted_points = sorted(state, key=lambda p: calculate_angle(p, center))
    draw.polygon(sorted_points, outline="red", fill=None, width=round(radius/2))
    
    return img, state


def calculate_angle(point, center):
    return math.atan2(point[1] - center[1], point[0] - center[0])

# 좌표를 받아서 폴리곤을 그리는 함수
def draw_polygon_on_image(image, evt: gr.SelectData, state):
    img = Image.fromarray(image)

    pt = (evt.index[0], evt.index[1])
    state.append(pt)
    # 클릭한 좌표를 저장
    area = image.shape[0]*image.shape[1]
    radius=max(5, round(area**0.5 / 120))
    
    draw = ImageDraw.Draw(img)
    for pt in state:
        left_up_point = (pt[0] - radius, pt[1] - radius)
        right_down_point = (pt[0] + radius, pt[1] + radius)
        draw.ellipse([left_up_point, right_down_point], outline="black", fill="red")
        
    if len(state) == 2:
        draw.line([state[0], state[1]], fill="red", width=round(radius/2))
    if len(state) >= 3:  # 좌표가 최소 3개 이상일 때만 폴리곤 그리기
        center = (sum(p[0] for p in state) / len(state), sum(p[1] for p in state) / len(state))
        # 각도에 따라 점들을 정렬
        sorted_points = sorted(state, key=lambda p: calculate_angle(p, center))
        draw.polygon(sorted_points, outline="red", fill=None, width=round(radius/2))

    return img, state

def sort_corners(corners):
    # 각 좌표를 (x, y) 형태로 받는다고 가정합니다.
    # corners = [(x1, y1), (x2, y2), (x3, y3), (x4, y4)]
    
    if len(corners) != 4:
        raise ValueError("Input should contain exactly four coordinates.")
    
    # 좌표를 y 기준으로 정렬하여 각 좌표를 결정합니다.
    sorted_by_y = sorted(corners, key=lambda p: p[1])  # y 기준으로 정렬
    lt, rt = sorted(sorted_by_y[:2], key=lambda p: p[0])
    lb, rb = sorted(sorted_by_y[2:], key=lambda p: p[0])
    
    return lt, rt, rb, lb

def convert(image, state):
    h,w = image.shape[:2]
    if len(state) < 4:
        out_image, msk_np = docscanner_rec(image, docscanner, cuda)
        out_image = out_image[:,:,::-1]
    elif len(state) ==4:
        state = list(sort_corners(state))
        src = np.array(state).astype(np.float32)
        dst = np.float32([
            (0, 0),
            (w - 1, 0),
            (w - 1, h - 1),
            (0, h - 1)
        ])
        out_image, M = unwarp(image, src, dst)
    return out_image
css = """
.image-container {
    padding: 20px;
    background-color: #f0f0f0;
}
"""
# Gradio Blocks 컨텍스트에서 인터페이스 구성
with gr.Blocks(css=css) as demo:
    state = gr.State([])
    with gr.Row():
        with gr.Column():
            text = gr.Textbox("입력 이미지(코너를 클릭하세요)", show_label=False)
            image_input = gr.Image(show_label=False, interactive=True, elem_classes="image-container")
            clear_button = gr.Button("Clear Points")
            cutting_button = gr.Button("Cutting Image(need more than 2 points)")
            rotating_button = gr.Button("Rotate Image(clock wise 90 degree)")
            auto_button = gr.Button("Auto Points detection")
            convert_button = gr.Button("Convert Image")
        with gr.Column():
            text = gr.Textbox("변환될 영역", show_label=False)
            image_output = gr.Image(show_label=False)
            # state_display = gr.Textbox(label="Current State")
            # coordinates_text = gr.Textbox(label="Coordinates", placeholder="Enter coordinates (x, y) for each point")
            # update_coords_button = gr.Button("Update Coordinates")
        with gr.Column():
            text = gr.Textbox("결과 이미지", show_label=False)
            result_image = gr.Image(show_label=False, format="png")

    # # 이미지 위에서 클릭 이벤트 처리
    image_input.select(draw_polygon_on_image, inputs=[image_input,state], outputs=[image_output,state])
    
    
    
    # 좌표 초기화 버튼 클릭 시 좌표 리셋
    clear_button.click(fn=reset_points, inputs=[image_input,state], outputs=[image_output,state])
    # 이미지 자르기 편집
    cutting_button.click(fn=cutting_image, inputs=[image_input,state], outputs=[image_input, image_output, state])
    # 이미지 회전
    rotating_button.click(fn=rotate_image, inputs=[image_input], outputs=[image_input, state])
    # 자동 검출 버튼
    auto_button.click(fn=auto_point_detect, inputs=image_input, outputs=[image_output,state])
    # 변환 버튼
    convert_button.click(fn=convert, inputs=[image_input,state], outputs=result_image)

demo.launch(share=True)