from uav_seg import UAVSegPredictor

def test_uav_seg(model_path, input_path, output_path):
    predictor = UAVSegPredictor(
        model_path= model_path,
        model_type='deeplabv3plus',
        num_classes=7
    )
    color_map = [
                [255, 0, 0],    # 类别0: 红色;背景
                [0, 255, 0],    # 类别1: 绿色；荒地
                [0, 0, 255],    # 类别2: 蓝色；林地
                [255, 255, 0],  # 类别3: 黄色；农田
                [255, 0, 255],  # 类别4: 品红；水
                [0, 255, 255],  # 类别5: 青色；道路
                [128, 0, 128]   # 类别6: 紫色；建筑用地
            ]
    predictor.predict_folder(input_path, output_path,color_map)

if __name__ == '__main__':
    test_uav_seg('checkpoints/deeplabv3plus_best.pth', r'D:\work\AI_Python\Ai_tottle\map\test_demo', './output')