StitchImageServer/app/api/stitch_v2.py

import os
import shutil
import uuid
import zipfile
import logging
from pathlib import Path
from typing import List  # 导入 List 类型

from fastapi import APIRouter, UploadFile, File, Form, HTTPException, BackgroundTasks
from fastapi.responses import FileResponse, JSONResponse

# 导入我们的核心逻辑和数据模型
from app.core import stitcher_keypoint, stitcher_template
from app.schemas import StitchingMethod, KeypointFeatureDetector, KeypointBlendType, TemplateBlendType

from utils.utils import cleanup_temp_folder


router = APIRouter(prefix="/stitch", tags=['拼图'])

TEMP_DIR = Path("_temp_work")
TEMP_DIR.mkdir(exist_ok=True)
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')


# --- 内部辅助函数，用于处理单个拼图任务，避免代码重复 ---
def _process_single_puzzle(
    image_dir: Path,
    output_dir: Path,
    method: StitchingMethod,
    params: dict
) -> Path | None:
    """
    处理单个拼图任务的核心逻辑。

    :param image_dir: 包含所有小图的输入目录。
    :param output_dir: 存放拼接结果的输出目录。
    :param method: 拼图方法。
    :param params: 包含所有拼图参数的字典。
    :return: 成功则返回拼接后图片的路径，否则返回 None。
    """
    output_dir.mkdir(exist_ok=True)
    stitched_image_path = None
    try:
        if method == StitchingMethod.KEY_POINT:
            stitched_image_path = stitcher_keypoint.stitch_img(
                IMAGE_DIR=image_dir, OUTPUT_DIR=output_dir,
                NUM_COLS=params["num_cols"], NUM_ROWS=params["num_rows"],
                ESTIMATE_OVERLAP_HORIZONTAL_PIXELS=params["overlap_h"],
                ESTIMATE_OVERLAP_VERTICAL_PIXELS=params["overlap_v"],
                BLEND_TYPE=params["kp_blend_type"].value,
                FeatureDetector=params["kp_feature_detector"].value,
                DEBUG_MODE=False
            )
        elif method == StitchingMethod.TEMPLATE_MATCH:
            stitched_image_path = stitcher_template.stitch_img(
                IMAGE_DIR=image_dir, OUTPUT_DIR=output_dir,
                NUM_COLS=params["num_cols"], NUM_ROWS=params["num_rows"],
                ESTIMATE_OVERLAP_HORIZONTAL_PIXELS=params["overlap_h"],
                ESTIMATE_OVERLAP_VERTICAL_PIXELS=params["overlap_v"],
                BLEND_TYPE=params["tm_blend_type"].value,
                LIGHT_COMPENSATION=params["tm_light_compensation"],
                DEBUG_MODE=False
            )
    except Exception as e:
        logging.error(f"处理文件夹 {image_dir.name} 时发生错误: {e}")
        return None

    return stitched_image_path


# --- 新的单个拼图接口 (直接上传图片文件夹) ---
@router.post("/from-folder", response_class=FileResponse, summary="单个拼图接口 (从文件夹上传)")
async def stitch_single_from_folder(
        background_tasks: BackgroundTasks,
        files: List[UploadFile] = File(..., description="一个文件夹中的所有待拼接图片。"),
        output_filename_base: str = Form(..., description="输出图片的基础名称（不含扩展名），例如 'puzzle_A'。"),
        # --- 通用参数 ---
        method: StitchingMethod = Form(StitchingMethod.TEMPLATE_MATCH, description="选择拼图方法"),
        num_cols: int = Form(4, description="拼图的列数"),
        num_rows: int = Form(6, description="拼图的行数"),
        overlap_h: int = Form(405, description="预估的水平重叠像素"),
        overlap_v: int = Form(440, description="预估的垂直重叠像素"),
        # --- 点匹配法 (key_point) 特定参数 ---
        kp_blend_type: KeypointBlendType = Form(KeypointBlendType.COMBINE, description="[点匹配] 融合模式"),
        kp_feature_detector: KeypointFeatureDetector = Form(KeypointFeatureDetector.SIFT,
                                                            description="[点匹配] 特征检测器"),
        # --- 模板匹配法 (template_match) 特定参数 ---
        tm_blend_type: TemplateBlendType = Form(TemplateBlendType.HALF_IMPORTANCE_ADD_WEIGHT,
                                                description="[模板匹配] 融合模式"),
        tm_light_compensation: bool = Form(True, description="[模板匹配] 是否启用光照补偿")
):
    """
    上传一个文件夹内的所有图片进行拼接，直接返回拼接好的单张大图。

    - **files**: 选择一个文件夹中的所有图片进行上传。
    - **output_filename_base**: 为你的拼图任务命名，这个名字将作为返回图片的文件名。
    - **返回**: 拼接成功后，返回拼接好的图片文件。
    """
    if not files:
        raise HTTPException(status_code=400, detail="没有上传任何文件。")

    request_id = str(uuid.uuid4())
    session_dir = TEMP_DIR / request_id
    session_dir.mkdir()
    background_tasks.add_task(cleanup_temp_folder, session_dir)

    # 创建用于存放上传图片的临时目录
    image_dir = session_dir / "images"
    image_dir.mkdir()

    # 保存所有上传的文件
    for upload_file in files:
        file_path = image_dir / upload_file.filename
        with open(file_path, "wb") as buffer:
            shutil.copyfileobj(upload_file.file, buffer)

    output_dir = session_dir / "output"

    # 将所有参数打包到一个字典中，方便传递
    params = locals()

    # 调用核心处理函数
    stitched_image_path = _process_single_puzzle(image_dir, output_dir, method, params)

    if not stitched_image_path or not stitched_image_path.exists():
        raise HTTPException(status_code=500, detail=f"图片拼接失败，请检查服务器日志（请求ID: {request_id}）。")

    # 使用用户提供的基础名称命名输出图片
    final_filename = f"{output_filename_base}.jpg"
    final_filepath = stitched_image_path.rename(stitched_image_path.parent / final_filename)

    return FileResponse(
        path=final_filepath,
        filename=final_filename,
        media_type='image/jpeg'
    )


# --- 新的批量拼图接口 (上传一个拼图文件夹，返回ZIP) ---
@router.post("/batch/from-folder", response_class=FileResponse, summary="批量拼图接口 (单文件夹上传, ZIP返回)")
async def stitch_batch_from_folder(
        background_tasks: BackgroundTasks,
        files: List[UploadFile] = File(..., description="一个文件夹中的所有待拼接图片。"),
        output_filename_base: str = Form(..., description="输出图片的基础名称（不含扩展名），例如 'puzzle_A'。"),
        # --- 参数与单个拼图接口相同 ---
        method: StitchingMethod = Form(StitchingMethod.TEMPLATE_MATCH, description="选择拼图方法"),
        num_cols: int = Form(4, description="拼图的列数"),
        num_rows: int = Form(6, description="拼图的行数"),
        overlap_h: int = Form(405, description="预估的水平重叠像素"),
        overlap_v: int = Form(440, description="预估的垂直重叠像素"),
        kp_blend_type: KeypointBlendType = Form(KeypointBlendType.COMBINE, description="[点匹配] 融合模式"),
        kp_feature_detector: KeypointFeatureDetector = Form(KeypointFeatureDetector.SIFT,
                                                            description="[点匹配] 特征检测器"),
        tm_blend_type: TemplateBlendType = Form(TemplateBlendType.HALF_IMPORTANCE_ADD_WEIGHT,
                                                description="[模板匹配] 融合模式"),
        tm_light_compensation: bool = Form(True, description="[模板匹配] 是否启用光照补偿"),
):
    """
    上传一个文件夹内的所有图片进行拼接，将结果打包成一个ZIP压缩文件返回。

    - **files**: 选择一个文件夹中的所有图片进行上传。
    - **output_filename_base**: 为你的拼图任务命名，这个名字将作为ZIP包内图片的文件名。
    - **返回**: 一个ZIP压缩包，里面包含拼接好的单张图片。
    """
    if not files:
        raise HTTPException(status_code=400, detail="没有上传任何文件。")

    request_id = str(uuid.uuid4())
    session_dir = TEMP_DIR / request_id
    session_dir.mkdir()
    background_tasks.add_task(cleanup_temp_folder, session_dir)

    image_dir = session_dir / "images"
    image_dir.mkdir()

    for upload_file in files:
        file_path = image_dir / upload_file.filename
        with open(file_path, "wb") as buffer:
            shutil.copyfileobj(upload_file.file, buffer)

    single_output_dir = session_dir / "single_output"
    params = locals()

    stitched_image_path = _process_single_puzzle(image_dir, single_output_dir, method, params)

    if not stitched_image_path or not stitched_image_path.exists():
        raise HTTPException(status_code=500, detail=f"图片拼接失败，请检查服务器日志（请求ID: {request_id}）。")

    # --- 将单个结果打包成ZIP ---
    batch_output_dir = session_dir / "batch_output"
    batch_output_dir.mkdir()

    # 将成功的结果移到最终的批量输出目录
    target_path = batch_output_dir / f"{output_filename_base}.jpg"
    shutil.move(str(stitched_image_path), str(target_path))

    # 将最终结果打包成ZIP
    output_zip_path = session_dir / "stitched_result.zip"
    shutil.make_archive(str(output_zip_path.with_suffix('')), 'zip', batch_output_dir)

    return FileResponse(
        path=output_zip_path,
        filename=f"{output_filename_base}_stitched.zip",
        media_type='application/zip'
    )