CheckCardBoxAndDefectServer/Test/model_test01.py

import os
from pathlib import Path
from app.utils.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):
    # 配置参数
    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

    result = predictor.predict_single_image(
        img_path=now_img_path,
        save_visualization=True,
        save_json=True,
        answer_json_dir_str=answer_json_dir_str
    )


if __name__ == '__main__':
    config = {
        "pokemon_outer_box": {
            "pth_path": "Model/pokemon_outer_box.pth",
            "class_dict": {1: 'outer_box'},
            "img_size": {'width': 1280, 'height': 1280},
            "confidence": 0.5,
            "input_channels": 3,
        },
        "pokemon_inner_box": {
            "pth_path": "Model/pokemon_inner_box.pth",
            "class_dict": {1: 'inner_box'},
            "img_size": {'width': 1280, 'height': 1280},
            "confidence": 0.5,
            "input_channels": 3,
        },
        "pokemon_back_corner_defect": {
            "pth_path": "Model/pokemon_back_corner_defect.pth",
            "class_dict": {
                1: 'wear', 2: 'wear_and_impact', 3: 'impact',
                4: 'damaged', 5: 'wear_and_stain',
            },
            "img_size": {'width': 512, 'height': 512},
            "confidence": 0.5,
            "input_channels": 3,
        },
        "pokemon_front_corner_reflect_defect": {
            "pth_path": "Model/pokemon_front_corner_reflect_defect.pth",
            "class_dict": {"1": "impact", "2": "wear_and_impact", "3": "wear"},
            "img_size": {'width': 512, 'height': 512},
            "confidence": 0.5,
            "input_channels": 3,
        },
        "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,
        },
        "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,
        },
    }
    # predict_single_image(config['pokemon_inner_box'],
    #                      img_path=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\img.png",
    #                      output_dir=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\inner")

    # predict_single_image(config['pokemon_outer_box'],
    #                      img_path=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\img.png",
    #                      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_r5_c2.jpg",
    #                      output_dir=r"C:\Code\ML\Project\CheckCardBoxAndDefectServer\temp\face_no_reflect")