CheckCardBoxAndDefectServer/Test/model_test01.py

import os
from pathlib import Path
from app.core.config import settings
import json
import cv2

from app.utils.card_inference.fry_bisenetv2_predictor_V04_250819 import FryBisenetV2Predictor

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


def predict_single_image(config_params: dict,
                         img_path: str,
                         output_dir: str,
                         only_json=False):
    # 配置参数
    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,
    )

    # 单张图片预测
    print("=== 单张图片预测 ===")
    now_img_path = img_path
    answer_json_dir_str = output_dir

    if not only_json:
        result = predictor.predict_single_image(
            img_path=now_img_path,
            save_visualization=True,
            save_json=True,
            answer_json_dir_str=answer_json_dir_str
        )
    else:
        img_bgr = cv2.imread(now_img_path)
        result = predictor.predict_from_image(img_bgr)

    return result


if __name__ == '__main__':
    big_img_path = r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\250805_pokemon_0001.jpg"

    config = settings.CARD_MODELS_CONFIG
    # predict_single_image(config['pokemon_front_inner_box'],
    #                      img_path=big_img_path,
    #                      output_dir=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\inner")
    #
    # predict_single_image(config['pokemon_back_inner_box'],
    #                      img_path=r"C:\Users\wow38\Downloads\_250829_1656_最新模型汇总\_250829_1814_宝可梦背面内框模型\pth_and_images\images\Pokémon_back_0805_0001.jpg",
    #                      output_dir=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\back_inner")
    #
    #
    # predict_single_image(config['outer_box'],
    #                      big_img_path,
    #                      output_dir=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\outer")

    # predict_single_image(config['pokemon_back_corner_defect'],
    #                      img_path=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\img_2.png",
    #                      output_dir=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\back_corner")

    # predict_single_image(config['pokemon_front_face_no_reflect_defect'],
    #                      img_path=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\250805_pokemon_0001_grid_r3_c4.jpg",
    #                      output_dir=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\face_no_reflect")

    result = predict_single_image(config['pokemon_front_corner_no_reflect_defect'],
                                  img_path=r"C:\Users\wow38\Downloads\_250829_1656_最新模型汇总\_250829_1817_宝可梦非闪光卡正面边角模型\pth_and_images\images\250730_pokemon_0031_bottom_grid_r0_c5.jpg",
                                  output_dir=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\corner_no_reflect")

    # result = predict_single_image(config['pokemon_front_face_no_reflect_defect'],
    #                      img_path=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\250805_pokemon_0001_grid_r3_c4.jpg",
    #                      output_dir=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\face_no_reflect",
    #                      only_json=True)

    # print(result)