CheckCardBoxAndDefectServer/Test/切割合并.py

import json
import numpy as np
import cv2
from typing import Dict
from app.utils.CardDefectAggregator import CardDefectAggregator
from pathlib import Path

# 假设你的预测器类在这里，我们为了测试会创建一个MockPredictor
from app.utils.card_inference.fry_bisenetv2_predictor_V04_250819 import FryBisenetV2Predictor

BASE_PATH = Path(__file__).parent.parent.absolute()


def get_predictor(config_params: dict):
    # 配置参数
    model_path = BASE_PATH / config_params["pth_path"]
    real_seg_class_dict = config_params['class_dict']
    imgSize_train_dict = config_params['img_size']
    confidence = config_params['confidence']
    input_channels = config_params['input_channels']

    # 为不同类别设置不同颜色（可选）
    label_colors_dict = {
        'outer_box': (255, 0, 0),
    }

    # 创建预测器
    predictor = FryBisenetV2Predictor(
        pth_path=str(model_path),
        real_seg_class_dict=real_seg_class_dict,
        imgSize_train_dict=imgSize_train_dict,
        confidence=confidence,
        label_colors_dict=label_colors_dict,
        input_channels=input_channels,
    )
    return predictor


def _test_face_big_img():
    large_image_path = r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\250805_pokemon_0001.jpg"

    pokemon_front_face_no_reflect_defect = {
        "pth_path": "Model/pokemon_front_face_no_reflect_defect.pth",
        "class_dict": {"1": "scratch", "2": "pit", "3": "stain"},
        "img_size": {'width': 512, 'height': 512},
        "confidence": 0.5,
        "input_channels": 3,
    }
    predictor = get_predictor(pokemon_front_face_no_reflect_defect)

    # 3. 实例化我们的聚合器，传入预测器
    aggregator = CardDefectAggregator(
        predictor=predictor,
        tile_size=512,
        overlap_ratio=0.1,  # 10% 重叠
    )

    # --- 执行任务 ---
    # 任务1: 对整个卡片表面进行缺陷检测
    aggregator.process_image(
        image_path=large_image_path,
        output_json_path="output/final_face_defects.json",
        mode='face'
    )

    print("\n" + "=" * 50 + "\n")



def _test_corner_big_img():
    large_image_path = r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\250805_pokemon_0001.jpg"

    pokemon_front_corner_no_reflect_defect = {
        "pth_path": "Model/pokemon_front_corner_no_reflect_defect.pth",
        "class_dict": {"1": "wear", "2": "wear_and_impact", "3": "impact", "4": "damaged"},
        "img_size": {'width': 512, 'height': 512},
        "confidence": 0.5,
        "input_channels": 3,
    }
    predictor = get_predictor(pokemon_front_corner_no_reflect_defect)

    # 3. 实例化我们的聚合器，传入预测器
    aggregator = CardDefectAggregator(
        predictor=predictor,
        tile_size=512,
        overlap_ratio=0.1,  # 10% 重叠
    )

    # 任务2: 仅对卡片边缘进行缺陷检测 (使用另一个模型)
    # 假设你有一个专门用于边角的模型
    # aggregator.process_image(
    #     image_path=large_image_path,
    #     output_json_path="output/final_edge_defects.json",
    #     mode='edge'
    # )

    json_data = aggregator.process_image(
        image_path=large_image_path,
        mode='edge'
    )
    print(json_data)


def _test_split_img(split_mode):
    # "edge", "face"
    # --- 输入与输出路径配置 ---
    large_image_path = r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\250805_pokemon_0001.jpg"
    output_dir = r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\split_output"

    # 打印信息
    print(f"输入大图: {large_image_path}")
    print(f"输出目录: {output_dir}")
    print("-" * 50)

    # --- 实例化聚合器 ---
    # 由于我们只进行分割，不需要预测器，所以 predictor=None 即可。
    aggregator = CardDefectAggregator(
        predictor=None,  # 无需预测器
        tile_size=512,
        overlap_ratio=0.1,  # 10% 重叠
    )

    # --- 调用新的分割保存功能 ---
    # mode='face' 表示对整个图片进行网格分割
    saved_tile_paths = aggregator.split_and_save_tiles(
        image_path=large_image_path,
        output_dir=output_dir,
        mode=split_mode
    )

    # 打印一些结果信息
    if saved_tile_paths:
        print(f"\n分割出的第一个文件是: {saved_tile_paths[0]}")
        print(f"分割出的最后一个文件是: {saved_tile_paths[-1]}")

if __name__ == "__main__":
    # _test_face_big_img()
    _test_corner_big_img()
    # _test_split_img(split_mode='edge')