StitchImageServer/app/api/stitch.py

import os
import shutil
import uuid
import zipfile
import logging
from pathlib import Path

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


@router.post("/", response_class=FileResponse, summary="单个拼图接口")
async def stitch_single_puzzle(
        background_tasks: BackgroundTasks,
        zip_file: UploadFile = File(..., description="包含一个文件夹的ZIP压缩包，文件夹内有24张小图。"),
        # --- 通用参数 ---
        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="[模板匹配] 是否启用光照补偿")
):
    """
    上传一个包含24张小图的文件夹的ZIP压缩包，接口会将其拼接成一张大图并返回。

    - **zip_file**: 必须是.zip格式，内部应仅包含一个文件夹，该文件夹内含所有待拼接的.jpg图片。
    - **返回**: 拼接成功后，返回拼接好的图片文件，文件名与ZIP包内的文件夹名相同。
    """
    request_id = str(uuid.uuid4())
    session_dir = TEMP_DIR / request_id
    session_dir.mkdir()
    background_tasks.add_task(cleanup_temp_folder, session_dir)

    zip_path = session_dir / zip_file.filename
    with open(zip_path, "wb") as buffer:
        shutil.copyfileobj(zip_file.file, buffer)

    extracted_dir = session_dir / "extracted"
    extracted_dir.mkdir()

    try:
        with zipfile.ZipFile(zip_path, 'r') as zf:
            zf.extractall(extracted_dir)
    except zipfile.BadZipFile:
        raise HTTPException(status_code=400, detail="上传的文件不是有效的ZIP格式。")

    image_dir = extracted_dir
    output_dir = session_dir / "output"

    # 根据选择的方法调用不同的拼接函数
    stitched_image_path = None
    if method == StitchingMethod.KEY_POINT:
        stitched_image_path = stitcher_keypoint.stitch_img(
            IMAGE_DIR=image_dir, OUTPUT_DIR=output_dir, NUM_COLS=num_cols, NUM_ROWS=num_rows,
            ESTIMATE_OVERLAP_HORIZONTAL_PIXELS=overlap_h, ESTIMATE_OVERLAP_VERTICAL_PIXELS=overlap_v,
            BLEND_TYPE=kp_blend_type.value, FeatureDetector=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=num_cols, NUM_ROWS=num_rows,
            ESTIMATE_OVERLAP_HORIZONTAL_PIXELS=overlap_h, ESTIMATE_OVERLAP_VERTICAL_PIXELS=overlap_v,
            BLEND_TYPE=tm_blend_type.value, LIGHT_COMPENSATION=tm_light_compensation,
            DEBUG_MODE=False
        )

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

    # 使用原始文件夹名命名输出图片
    final_filename = f"{image_dir.name}.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'
    )


@router.post("/batch", response_class=FileResponse, summary="批量拼图接口")
async def stitch_batch_puzzles(
        background_tasks: BackgroundTasks,
        zip_file: UploadFile = File(..., description="包含多个拼图文件夹的ZIP压缩包。"),
        # 参数与单个拼图接口相同
        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压缩包，接口会处理所有文件夹，并将结果打包成一个新的ZIP返回。

    - **zip_file**: 必须是.zip格式，内部可以有多个文件夹，每个文件夹都包含待拼接的图片。
    - **返回**: 一个ZIP压缩包，里面是所有拼接好的图片，每张图片以其对应的原文件夹名命名。
    """
    request_id = str(uuid.uuid4())
    session_dir = TEMP_DIR / request_id
    session_dir.mkdir()
    background_tasks.add_task(cleanup_temp_folder, session_dir)

    zip_path = session_dir / zip_file.filename
    with open(zip_path, "wb") as buffer:
        shutil.copyfileobj(zip_file.file, buffer)

    extracted_dir = session_dir / "extracted"
    extracted_dir.mkdir()
    try:
        with zipfile.ZipFile(zip_path, 'r') as zf:
            zf.extractall(extracted_dir)
    except zipfile.BadZipFile:
        raise HTTPException(status_code=400, detail="上传的文件不是有效的ZIP格式。")

    puzzle_folders = [d for d in extracted_dir.iterdir() if d.is_dir()]
    if not puzzle_folders:
        raise HTTPException(status_code=400, detail="ZIP包中未找到任何拼图文件夹。")

    batch_output_dir = session_dir / "batch_output"
    batch_output_dir.mkdir()

    processed_count = 0
    failed_folders = []

    for image_dir in puzzle_folders:
        logging.info(f"--- 开始处理批量任务中的文件夹: {image_dir.name} ---")
        # 为每个子任务创建一个独立的输出目录
        single_output_dir = session_dir / "single_output"
        if single_output_dir.exists():
            shutil.rmtree(single_output_dir)  # 清理上一次循环的输出

        stitched_image_path = None
        try:
            if method == StitchingMethod.KEY_POINT:
                stitched_image_path = stitcher_keypoint.stitch_img(
                    IMAGE_DIR=image_dir, OUTPUT_DIR=single_output_dir,
                    NUM_COLS=num_cols, NUM_ROWS=num_rows,
                    ESTIMATE_OVERLAP_HORIZONTAL_PIXELS=overlap_h,
                    ESTIMATE_OVERLAP_VERTICAL_PIXELS=overlap_v,
                    BLEND_TYPE=kp_blend_type.value, FeatureDetector=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=single_output_dir,
                    NUM_COLS=num_cols, NUM_ROWS=num_rows,
                    ESTIMATE_OVERLAP_HORIZONTAL_PIXELS=overlap_h,
                    ESTIMATE_OVERLAP_VERTICAL_PIXELS=overlap_v,
                    BLEND_TYPE=tm_blend_type.value, LIGHT_COMPENSATION=tm_light_compensation,
                    DEBUG_MODE=False
                )

            if stitched_image_path and stitched_image_path.exists():
                # 将成功的结果移到最终的批量输出目录
                target_path = batch_output_dir / f"{image_dir.name}.jpg"
                shutil.move(str(stitched_image_path), str(target_path))
                processed_count += 1
            else:
                logging.error(f"文件夹 {image_dir.name} 拼接失败。")
                failed_folders.append(image_dir.name)
        except Exception as e:
            logging.error(f"处理文件夹 {image_dir.name} 时发生严重错误: {e}")
            failed_folders.append(image_dir.name)

    if processed_count == 0:
        detail_msg = f"所有 {len(puzzle_folders)} 个文件夹都拼接失败。失败列表: {failed_folders}"
        raise HTTPException(status_code=500, detail=detail_msg)

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

    if failed_folders:
        logging.warning(f"批量任务完成，但有 {len(failed_folders)} 个文件夹失败: {failed_folders}")
        # 可以在响应头中添加自定义信息来通知客户端部分失败
        # response.headers["X-Failed-Folders"] = ",".join(failed_folders)

    return FileResponse(
        path=output_zip_path,
        filename="stitched_results.zip",
        media_type='application/zip'
    )