图片审校区域标记工具

为了方便准备审校数据集，开发标记工具，将存在审校修改的图片区域单独保存，提供给后续流程来产生数据集。

from nicegui import ui
import os
import glob
import time
from PIL import Image

# 自动管理图片路径和保存目录
INPUT_DIR = 'static/input'
DONE_DIR = 'static/done'
CROPS_DIR = 'static/crops'
GRAY_IMAGE = 'static/gray.png'

os.makedirs(INPUT_DIR, exist_ok=True)
os.makedirs(DONE_DIR, exist_ok=True)
os.makedirs(CROPS_DIR, exist_ok=True)

DISPLAY_WIDTH = 1240
DISPLAY_HEIGHT = 1790
SELECT_BOX_HEIGHT = 300

select_box_y = 0
img = None
select_box = None

def ensure_gray_image():
    if not os.path.exists(GRAY_IMAGE):
        img_gray = Image.new('RGB', (DISPLAY_WIDTH, DISPLAY_HEIGHT), (200, 200, 200))
        img_gray.save(GRAY_IMAGE)

def get_input_images():
    return sorted(glob.glob(os.path.join(INPUT_DIR, '*.jpg')) + glob.glob(os.path.join(INPUT_DIR, '*.png')))

def get_first_image():
    images = get_input_images()
    if images:
        return images[0]
    ensure_gray_image()
    return GRAY_IMAGE

current_image_path = get_first_image()

def update_select_box(event):
    global select_box_y
    if event is None:
        y = select_box_y
    else:
        offset_y = getattr(event.args, 'offsetY', None)
        if offset_y is None:
            offset_y = event.args.get('offsetY', 0) if isinstance(event.args, dict) else 0
        y = max(0, min(offset_y, DISPLAY_HEIGHT - SELECT_BOX_HEIGHT))
    select_box_y = y
    select_box.style(f'top: {select_box_y}px; left: 0px; width: {DISPLAY_WIDTH}px; height: {SELECT_BOX_HEIGHT}px; position: absolute; border: 2px solid #00f; background: rgba(0,0,255,0.1); pointer-events: none;')

def crop_and_save(event):
    global select_box_y, current_image_path
    img_path = current_image_path
    if not os.path.exists(img_path):
        ui.notify('图片不存在')
        return
    try:
        img_obj = Image.open(img_path)
        img_width, img_height = img_obj.size
        scale = img_height / DISPLAY_HEIGHT
        crop_left = 0
        crop_upper = int(select_box_y * scale)
        crop_right = img_width
        crop_lower = int(min((select_box_y + SELECT_BOX_HEIGHT) * scale, img_height))
        crop_box = (crop_left, crop_upper, crop_right, crop_lower)
        crop_img = img_obj.crop(crop_box)
        if crop_img.mode == "RGBA":
            crop_img = crop_img.convert("RGB")
        ts = int(time.time() * 1000)
        save_path = os.path.join(CROPS_DIR, f'{ts}.jpg')
        crop_img.save(save_path)
        ui.notify(f'已保存: {save_path}')
    except Exception as e:
        ui.notify(f'裁剪失败: {e}')

def handle_mouse_click(event):
    global SELECT_BOX_HEIGHT
    button = event.args.get('button', 0)
    if button == 0:  # left
    #     SELECT_BOX_HEIGHT = max(10, SELECT_BOX_HEIGHT - 10)
    #     update_select_box(event)
    # elif button == 2:  # right
    #     SELECT_BOX_HEIGHT += 10
    #     update_select_box(event)
    # elif button == 1:  # middle
        crop_and_save(event)

def done_and_next():
    global current_image_path, img
    images = get_input_images()
    if current_image_path in images:
        basename = os.path.basename(current_image_path)
        # 检查done目录是否已存在同名文件，若存在则加时间戳
        done_path = os.path.join(DONE_DIR, basename)
        if os.path.exists(done_path):
            name, ext = os.path.splitext(basename)
            ts = int(time.time() * 1000)
            done_path = os.path.join(DONE_DIR, f"{name}_{ts}{ext}")
        os.rename(current_image_path, done_path)
        ui.notify(f'已移到: {done_path}')
    images = get_input_images()
    if images:
        current_image_path = images[0]
    else:
        ensure_gray_image()
        current_image_path = GRAY_IMAGE
    img.set_source(current_image_path)
    
def adjust_height(delta):
    global SELECT_BOX_HEIGHT
    if delta < 0:
        SELECT_BOX_HEIGHT = max(10, SELECT_BOX_HEIGHT + delta)
    else:
        SELECT_BOX_HEIGHT += delta
    update_select_box(None)

@ui.page('/')
def main_page():
    global img, select_box, current_image_path
    with ui.row().classes('w-full h-full justify-center items-center'):
        with ui.column().classes('w-full h-full items-center'):
            ui.label('审校区域标记工具').classes('text-2xl q-mb-md')
            ui.label('选择有审校修改的区域进行标记, 一页内容标记完成后点击“完成并下一张”按钮, 当图片显示为灰色时表示已经处理完')
            with ui.element('div').style(f'position: relative; width: {DISPLAY_WIDTH}px; height: {DISPLAY_HEIGHT}px;'):
                img = ui.image(current_image_path).style(f'width: {DISPLAY_WIDTH}px; height: {DISPLAY_HEIGHT}px; display: block; margin-left: auto; margin-right: auto;')
                select_box = ui.element('div').style(f'position: absolute; top: 0px; left: 0px; width: {DISPLAY_WIDTH}px; height: {SELECT_BOX_HEIGHT}px; border: 2px solid #00f; background: rgba(0,0,255,0.1); pointer-events: none;')
                ui.button('-', on_click=lambda: adjust_height(-10)).style('position: fixed; top: 10px; right: 10px; width: 30px; height: 30px;')
                ui.button('+', on_click=lambda: adjust_height(10)).style('position: fixed; top: 10px; right: 50px; width: 30px; height: 30px;')
                img.on('mousemove', update_select_box)
                img.on('mousemove', update_select_box)
                if current_image_path != GRAY_IMAGE:
                    img.on('mousedown', handle_mouse_click)
            # 添加一个与图片等宽、高度为100的按钮，点击后执行done_and_next
            ui.button('完成并下一张', on_click=done_and_next).style(f'width: {DISPLAY_WIDTH}px; height: 100px; font-size: 2rem; margin-top: 10px; margin-left: auto; margin-right: auto; display: block;')

ui.run(title='图片裁剪工具', native=False, port=8080)