simulate.py

"""
人脸识别调度模拟器
模拟五种划分方案，考虑网络延迟和进程占用
"""

import grpc
import base64
import pickle
import argparse
import time
import os
import sys
from typing import Optional, Dict, Any, List, Tuple
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
# 添加中文字体支持
matplotlib.rcParams['font.sans-serif'] = ['SimHei', 'Microsoft YaHei', 'DejaVu Sans']
matplotlib.rcParams['axes.unicode_minus'] = False

# 添加项目路径和rpc路径
sys.path.append(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(os.path.join(os.path.dirname(os.path.abspath(__file__)), 'rpc'))

# 使用绝对导入
import face_recognition_pb2
import face_recognition_pb2_grpc
import config
from dynamic_scheduler.module_wrapper import wrap_detect, wrap_align, wrap_extract, wrap_match
from dynamic_scheduler.partition_scheme import find_optimal_partition
from dynamic_scheduler.time_predictor import predict_all_modules_time

class FaceRecognitionSimulator:
    """人脸识别调度模拟器"""

    def __init__(self, server_address: str = "localhost:50051", 
                 process_usage: float = 0.0, network_delay: float = 0.0):
        """初始化模拟器"""
        self.channel = grpc.insecure_channel(server_address)
        self.stub = face_recognition_pb2_grpc.FaceRecognitionServiceStub(self.channel)
        self.process_usage = process_usage
        self.network_delay = network_delay
        print(f"连接到服务器: {server_address}")
        print(f"模拟参数 - 进程占用: {process_usage}, 网络延迟: {network_delay}ms")

    def __del__(self):
        """关闭连接"""
        if hasattr(self, 'channel'):
            self.channel.close()

    def encode_image(self, image_path: str) -> tuple[bytes, str]:
        """编码图像为base64字符串"""
        with open(image_path, 'rb') as f:
            image_data = f.read()

        ext = os.path.splitext(image_path)[1][1:].upper()
        if ext not in ['JPEG', 'JPG', 'PNG', 'BMP']:
            ext = 'JPEG'

        return base64.b64encode(image_data), ext

    def create_options(self, **kwargs) -> face_recognition_pb2.ProcessOptions:
        """创建处理选项"""
        return face_recognition_pb2.ProcessOptions(
            device=kwargs.get('device', 'auto'),
            confidence_threshold=kwargs.get('confidence_threshold', 0.8),
            max_faces=kwargs.get('max_faces', 5),
            align_output_size=kwargs.get('align_output_size', config.ALIGN_OUTPUT_SIZE),
            pretrained_source=kwargs.get('pretrained_source', config.PRETRAINED_SOURCE),
            batch_size=kwargs.get('batch_size', 32),
            database_path=kwargs.get('database_path', config.DB_PATH),
            top_k=kwargs.get('top_k', 5),
            return_intermediate_results=kwargs.get('return_intermediate_results', True)
        )

    def pipeline_scheme_0(self, image_path: str, **options) -> Dict[str, Any]:
        """方案0: 全云端执行 (detect->align->extract->match)"""
        print(f"\n=== 方案0: 全云端执行 ===")
        
        # 编码图像
        image_b64, image_format = self.encode_image(image_path)

        # 构建请求
        request = face_recognition_pb2.FaceRecognitionRequest(
            stage=face_recognition_pb2.DETECT,
            image_data=image_b64,
            image_format=image_format,
            options=self.create_options(**options)
        )

        # 发送请求（使用网络延迟包装）
        wrapped_stub = self._wrap_grpc_call(self.stub.ProcessFaceRecognition, "network")
        start_time = time.time()
        response = wrapped_stub(request)
        end_time = time.time()

        result = self._parse_response(response, start_time, end_time)
        # 方案0：T_total = T_cloud^0(d), 使用实际墙上时间（已包含网络延迟模拟）
        result['client_time_ms'] = (end_time - start_time) * 1000
        
        return result

    def pipeline_scheme_1(self, image_path: str, **options) -> Dict[str, Any]:
        """方案1: 端侧detect，云端align->extract->match"""
        print(f"\n=== 方案1: 端侧detect，云端align->extract->match ===")
        
        # 端侧：执行detect（使用进程占用包装）
        wrapped_detect = wrap_detect(self.process_usage, self.network_delay, "process")
        detection_result = wrapped_detect(config.DETECTION_DEVICE, image_path)

        # 打印本地模块运行时间
        print(f"端侧检测时间: {detection_result['time_ms']:.2f}ms")

        if detection_result.get('boxes') is None or len(detection_result['boxes']) == 0:
            return {'status': 'no_faces', 'message': '未检测到人脸'}

        # 准备云端请求
        detection_data = {
            'image': base64.b64encode(pickle.dumps(detection_result['image'])).decode(),
            'boxes': detection_result.get('boxes'),
            'probs': detection_result.get('probs'),
            'landmarks': detection_result.get('landmarks'),
            'time_ms': detection_result['time_ms']
        }

        # 构建云端请求
        request = face_recognition_pb2.FaceRecognitionRequest(
            stage=face_recognition_pb2.ALIGN,
            detection_results=base64.b64encode(pickle.dumps(detection_data)),
            options=self.create_options(**options)
        )

        # 发送到云端（使用网络延迟包装）
        wrapped_stub = self._wrap_grpc_call(self.stub.ProcessFaceRecognition, "network")
        start_time = time.time()
        response = wrapped_stub(request)
        end_time = time.time()

        result = self._parse_response(response, start_time, end_time, faces_detected=len(detection_result['boxes']))
        
        # 添加本地模块时间到结果中
        if 'module_times' not in result:
            result['module_times'] = {}
        result['module_times']['端侧检测'] = detection_result['time_ms']
        
        # 方案1：T_total = T_local^1(u) + T_cloud^1(d)
        # 累加实际执行时间：端侧detect + 云端执行时间（包含网络延迟）
        cloud_time = (end_time - start_time) * 1000
        result['client_time_ms'] = detection_result['time_ms'] + cloud_time
        
        return result

    def pipeline_scheme_2(self, image_path: str, **options) -> Dict[str, Any]:
        """方案2: 端侧detect->align，云端extract->match"""
        print(f"\n=== 方案2: 端侧detect->align，云端extract->match ===")
        
        # 端侧：执行detect和align（使用进程占用包装）
        wrapped_detect = wrap_detect(self.process_usage, self.network_delay, "process")
        wrapped_align = wrap_align(self.process_usage, self.network_delay, "process")
        
        detection_result = wrapped_detect(config.DETECTION_DEVICE, image_path)
        alignment_result = wrapped_align(
            detection_result['image'],
            detection_result.get('boxes'),
            detection_result.get('landmarks'),
            output_size=config.ALIGN_OUTPUT_SIZE
        )
        
        # 打印本地模块运行时间
        print(f"端侧检测时间: {detection_result['time_ms']:.2f}ms")
        print(f"端侧对齐时间: {alignment_result['time_ms']:.2f}ms")

        if len(alignment_result['aligned']) == 0:
            return {'status': 'no_faces', 'message': '未检测到人脸'}

        # 准备云端请求
        alignment_data = {
            'aligned': alignment_result['aligned'],
            'time_ms': alignment_result['time_ms']
        }

        # 构建云端请求
        request = face_recognition_pb2.FaceRecognitionRequest(
            stage=face_recognition_pb2.EXTRACT,
            alignment_results=base64.b64encode(pickle.dumps(alignment_data)),
            options=self.create_options(**options)
        )

        # 发送到云端（使用网络延迟包装）
        wrapped_stub = self._wrap_grpc_call(self.stub.ProcessFaceRecognition, "network")
        start_time = time.time()
        response = wrapped_stub(request)
        end_time = time.time()

        result = self._parse_response(response, start_time, end_time, faces_detected=len(alignment_result['aligned']))
        
        # 添加本地模块时间到结果中
        if 'module_times' not in result:
            result['module_times'] = {}
        result['module_times']['端侧检测'] = detection_result['time_ms']
        result['module_times']['端侧对齐'] = alignment_result['time_ms']
        
        # 方案2：T_total = T_local^1(u) + T_local^2(u) + T_cloud^2(d)
        # 累加实际执行时间：端侧detect+align + 云端执行时间（包含网络延迟）
        cloud_time = (end_time - start_time) * 1000
        result['client_time_ms'] = detection_result['time_ms'] + alignment_result['time_ms'] + cloud_time
        
        return result

    def pipeline_scheme_3(self, image_path: str, **options) -> Dict[str, Any]:
        """方案3: 端侧detect->align->extract，云端match"""
        print(f"\n=== 方案3: 端侧detect->align->extract，云端match ===")
        
        # 端侧：执行detect、align和extract（使用进程占用包装）
        wrapped_detect = wrap_detect(self.process_usage, self.network_delay, "process")
        wrapped_align = wrap_align(self.process_usage, self.network_delay, "process")
        wrapped_extract = wrap_extract(self.process_usage, self.network_delay, "process")
        
        detection_result = wrapped_detect(config.DETECTION_DEVICE, image_path)
        alignment_result = wrapped_align(
            detection_result['image'],
            detection_result.get('boxes'),
            detection_result.get('landmarks'),
            output_size=config.ALIGN_OUTPUT_SIZE
        )

        if len(alignment_result['aligned']) == 0:
            return {'status': 'no_faces', 'message': '未检测到人脸'}

        feature_result = wrapped_extract(
            config.DETECTION_DEVICE,
            config.PRETRAINED_SOURCE,
            alignment_result['aligned']
        )

        # 打印本地模块运行时间
        print(f"端侧检测时间: {detection_result['time_ms']:.2f}ms")
        print(f"端侧对齐时间: {alignment_result['time_ms']:.2f}ms")
        print(f"端侧特征提取时间: {feature_result['time_ms']:.2f}ms")

        # 准备云端请求
        feature_data = {
            'embeddings': feature_result['embeddings'],
            'time_ms': feature_result['time_ms']
        }

        # 构建云端请求
        request = face_recognition_pb2.FaceRecognitionRequest(
            stage=face_recognition_pb2.MATCH,
            feature_results=base64.b64encode(pickle.dumps(feature_data)),
            options=self.create_options(**options)
        )

        # 发送到云端（使用网络延迟包装）
        wrapped_stub = self._wrap_grpc_call(self.stub.ProcessFaceRecognition, "network")
        start_time = time.time()
        response = wrapped_stub(request)
        end_time = time.time()

        result = self._parse_response(response, start_time, end_time, faces_detected=len(alignment_result['aligned']))
        
        # 添加本地模块时间到结果中
        if 'module_times' not in result:
            result['module_times'] = {}
        result['module_times']['端侧检测'] = detection_result['time_ms']
        result['module_times']['端侧对齐'] = alignment_result['time_ms']
        result['module_times']['端侧特征提取'] = feature_result['time_ms']
        
        # 方案3：T_total = T_local^1(u) + T_local^2(u) + T_local^3(u) + T_cloud^3(d)
        # 累加实际执行时间：端侧detect+align+extract + 云端执行时间（包含网络延迟）
        cloud_time = (end_time - start_time) * 1000
        result['client_time_ms'] = detection_result['time_ms'] + alignment_result['time_ms'] + feature_result['time_ms'] + cloud_time
        
        return result

    def pipeline_scheme_4(self, image_path: str, **options) -> Dict[str, Any]:
        """方案4: 全端侧执行 (detect->align->extract->match)"""
        print(f"\n=== 方案4: 全端侧执行 ===")
        
        # 端侧：执行所有模块（使用进程占用包装）
        wrapped_detect = wrap_detect(self.process_usage, self.network_delay, "process")
        wrapped_align = wrap_align(self.process_usage, self.network_delay, "process")
        wrapped_extract = wrap_extract(self.process_usage, self.network_delay, "process")
        wrapped_match = wrap_match(self.process_usage, self.network_delay, "process")
        
        # 执行检测
        detection_result = wrapped_detect(config.DETECTION_DEVICE, image_path)
        if detection_result.get('boxes') is None or len(detection_result['boxes']) == 0:
            return {'status': 'no_faces', 'message': '未检测到人脸'}

        # 执行对齐
        alignment_result = wrapped_align(
            detection_result['image'],
            detection_result.get('boxes'),
            detection_result.get('landmarks'),
            output_size=config.ALIGN_OUTPUT_SIZE
        )

        if len(alignment_result['aligned']) == 0:
            return {'status': 'no_faces', 'message': '未检测到人脸'}

        # 执行特征提取
        feature_result = wrapped_extract(
            config.DETECTION_DEVICE,
            config.PRETRAINED_SOURCE,
            alignment_result['aligned']
        )

        # 执行匹配
        from modules.matcher import load_db
        db = load_db(config.DB_PATH)
        match_result = wrapped_match(db, feature_result['embeddings'], options.get('top_k', 5))

        # 打印本地模块运行时间
        print(f"端侧检测时间: {detection_result['time_ms']:.2f}ms")
        print(f"端侧对齐时间: {alignment_result['time_ms']:.2f}ms")
        print(f"端侧特征提取时间: {feature_result['time_ms']:.2f}ms")
        print(f"端侧匹配时间: {match_result['time_ms']:.2f}ms")

        # 方案4：T_total = Σ(i=1 to 4) T_local^i(u)
        # 累加实际执行时间：所有端侧模块
        total_time = detection_result['time_ms'] + alignment_result['time_ms'] + feature_result['time_ms'] + match_result['time_ms']

        return {
            'status': 'success',
            'message': '全端侧执行完成',
            'total_time_ms': total_time,
            'client_time_ms': total_time,
            'faces_detected': len(alignment_result['aligned']),
            'matches': match_result.get('matches', []),
            'module_times': {
                '端侧检测': detection_result['time_ms'],
                '端侧对齐': alignment_result['time_ms'],
                '端侧特征提取': feature_result['time_ms'],
                '端侧匹配': match_result['time_ms']
            }
        }

    def test_all_schemes(self, image_path: str, **options) -> List[Tuple[int, float, str]]:
        """
        测试所有划分方案并返回用时数据
        
        参数:
            image_path: 测试图像路径
            **options: 处理选项
            
        返回:
            List[Tuple[int, float, str]]: 每个方案的(方案编号, 用时(ms), 方案描述)
        """
        print(f"\n=== 开始测试所有划分方案 ===")
        print(f"图像: {image_path}")
        print(f"进程占用: {self.process_usage}, 网络延迟: {self.network_delay}ms")
        
        schemes = [
            (0, self.pipeline_scheme_0, "全云端执行"),
            (1, self.pipeline_scheme_1, "端侧detect，云端align->extract->match"),
            (2, self.pipeline_scheme_2, "端侧detect->align，云端extract->match"),
            (3, self.pipeline_scheme_3, "端侧detect->align->extract，云端match"),
            (4, self.pipeline_scheme_4, "全端侧执行")
        ]
        
        results = []
        
        for scheme_id, scheme_func, description in schemes:
            try:
                print(f"\n--- 测试方案 {scheme_id}: {description} ---")
                # 直接调用方案函数，不再在外层计时
                result = scheme_func(image_path, **options)
                
                if result.get('status') == 'success' or result.get('status') == face_recognition_pb2.FaceRecognitionResponse.SUCCESS:
                    # 使用方案函数内部计算的client_time_ms
                    execution_time = result['client_time_ms']
                    results.append((scheme_id, execution_time, description))
                    print(f"方案 {scheme_id} 执行时间: {execution_time:.2f}ms")
                else:
                    print(f"方案 {scheme_id} 执行失败: {result.get('message', '未知错误')}")
                    results.append((scheme_id, float('inf'), description))
                    
            except Exception as e:
                print(f"方案 {scheme_id} 执行异常: {str(e)}")
                results.append((scheme_id, float('inf'), description))
        
        return results

    def _wrap_grpc_call(self, grpc_func, delay_type: str):
        """包装gRPC调用，添加网络延迟"""
        from dynamic_scheduler.module_wrapper import wrap_module
        return wrap_module(grpc_func, self.process_usage, self.network_delay, delay_type)

    def _parse_response(self, response, start_time, end_time, faces_detected=None):
        """解析响应并记录模块运行时间"""
        result = {
            'status': response.status,
            'message': response.message,
            'total_time_ms': response.total_time_ms,
            'client_time_ms': (end_time - start_time) * 1000,
            'faces_detected': faces_detected or 0,
            'matches': [],
            'module_times': {}  # 新增：记录模块运行时间
        }

        # 记录云端模块运行时间
        if response.detection_results:
            result['module_times']['云端检测'] = response.detection_results.time_ms
        if response.alignment_results:
            result['module_times']['云端对齐'] = response.alignment_results.time_ms
        if response.feature_results:
            result['module_times']['云端特征提取'] = response.feature_results.time_ms
        if response.match_results:
            result['module_times']['云端匹配'] = response.match_results.time_ms

        if response.match_results:
            for i, match_item in enumerate(response.match_results.matches):
                face_matches = []
                for match_result in match_item.items:
                    face_matches.append({
                        'name': match_result.name,
                        'similarity': match_result.similarity
                    })
                result['matches'].append(face_matches)

        return result

    def _print_result(self, result: dict):
        """打印测试结果，包括模块运行时间"""
        print(f"状态: {result['status']}")
        print(f"消息: {result['message']}")
        print(f"处理时间: {result['total_time_ms']:.2f}ms")
        print(f"客户端总时间: {result['client_time_ms']:.2f}ms")
        print(f"检测到人脸数: {result['faces_detected']}")

        # 打印模块运行时间
        if result.get('module_times'):
            print("模块运行时间:")
            for module_name, time_ms in result['module_times'].items():
                print(f"  {module_name}: {time_ms:.2f}ms")

        if result['matches']:
            print("匹配结果:")
            for i, face_matches in enumerate(result['matches']):
                print(f"  人脸 {i+1}:")
                for match in face_matches[:3]:
                    print(f"    {match['name']}: {match['similarity']:.4f}")

    def find_optimal_scheme(self):
        """寻找最优划分方案"""
        print(f"\n=== 寻找最优划分方案 ===")
        print(f"进程占用: {self.process_usage}, 网络延迟: {self.network_delay}ms")
        
        # 直接调用find_optimal_partition函数，传递process_usage和network_delay参数
        optimal_scheme = find_optimal_partition(self.process_usage, self.network_delay)
        
        print(f"最优划分方案: 方案{optimal_scheme}")
        print(f"方案描述: {self._get_scheme_description(optimal_scheme)}")
        
        return optimal_scheme


    def _get_scheme_description(self, scheme: int) -> str:
        """获取方案描述"""
        descriptions = {
            0: "全云端执行 (detect->align->extract->match)",
            1: "端侧detect，云端align->extract->match",
            2: "端侧detect->align，云端extract->match",
            3: "端侧detect->align->extract，云端match",
            4: "全端侧执行 (detect->align->extract->match)"
        }
        return descriptions.get(scheme, "未知方案")

def plot_scheme_times(results: List[Tuple[int, float, str]], process_usage: float, network_delay: float):
    """绘制各划分方案的用时图表并保存到outputs目录"""
    # 过滤掉执行失败的方案
    valid_results = [(scheme_id, time, desc) for scheme_id, time, desc in results if time != float('inf')]
    
    if not valid_results:
        print("没有有效的执行结果可绘制")
        return
    
    schemes = [f"方案{scheme_id}" for scheme_id, _, _ in valid_results]
    times = [time for _, time, _ in valid_results]
    descriptions = [desc for _, _, desc in valid_results]
    
    # 创建图表
    plt.figure(figsize=(14, 8))
    
    # 使用新的配色方案
    cloud_color = '#2878B5'  # 云端执行颜色
    local_colors = ['#9AC9DB', '#F8AC8C', '#C82423', '#FF8884']  # 端侧模块颜色
    # 根据方案类型设置颜色
    colors = []
    for scheme_id, _, _ in valid_results:
        if scheme_id == 0:  # 全云端
            colors.append(cloud_color)
        elif scheme_id == 1:  # 端侧detect，云端其他
            colors.append(local_colors[0])  # detect颜色
        elif scheme_id == 2:  # 端侧detect->align，云端其他
            colors.append(local_colors[1])  # align颜色
        elif scheme_id == 3:  # 端侧detect->align->extract，云端match
            colors.append(local_colors[2])  # extract颜色
        else:  # 方案4：全端侧
            colors.append(local_colors[3])  # match颜色
    
    # 绘制柱状图
    bars = plt.bar(schemes, times, color=colors)
    
    # 添加数值标签
    for bar, time in zip(bars, times):
        plt.text(bar.get_x() + bar.get_width()/2, bar.get_height() + max(times)*0.01, 
                f'{time:.1f}ms', ha='center', va='bottom', fontsize=10)
    
    # 设置图表属性
    plt.title(f'各划分方案执行时间对比\n(进程占用: {process_usage}, 网络延迟: {network_delay}ms)', fontsize=14, fontweight='bold')
    plt.xlabel('划分方案', fontsize=12)
    plt.ylabel('执行时间 (毫秒)', fontsize=12)
    plt.xticks(rotation=45, ha='right')
    plt.grid(axis='y', alpha=0.3)
    
   
    # 添加方案描述说明（右上角）
    legend_text = "\n".join([f"方案{i}: {desc}" for i, desc in enumerate(descriptions)])
    plt.figtext(0.98, 0.98, legend_text, fontsize=9, bbox=dict(boxstyle="round,pad=0.3", facecolor="lightgray"), 
                ha='right', va='top')
    
    # 自动调整布局
    plt.tight_layout()
    
    # 创建outputs目录
    os.makedirs('outputs', exist_ok=True)
    
    # 生成文件名：月-日_时-分-秒.png
    import datetime
    current_time = datetime.datetime.now()
    filename = current_time.strftime("%m-%d_%H-%M-%S.png")
    filepath = os.path.join('outputs', filename)
    
    # 保存图片
    plt.savefig(filepath, dpi=300, bbox_inches='tight')
    print(f"图表已保存到: {filepath}")
    
    # 关闭图表以释放内存
    plt.close()

def main():
    parser = argparse.ArgumentParser(description='人脸识别调度模拟器')
    parser.add_argument('image_path', help='测试图像路径')
    parser.add_argument('--server', default='localhost:6006', help='服务器地址')
    parser.add_argument('--scheme', type=int, choices=[0,1,2,3,4], help='指定划分方案')
    parser.add_argument('--process-usage', type=float, default=0.0, help='模拟进程占用率 (0-1)')
    parser.add_argument('--network-delay', type=float, default=0.0, help='模拟网络延迟 (ms)')
    parser.add_argument('--find-optimal', action='store_true', help='寻找最优划分方案')
    parser.add_argument('--test-all', action='store_true', help='测试所有划分方案并绘制图表')
    parser.add_argument('--device', default='cuda', help='计算设备')
    parser.add_argument('--top_k', type=int, default=5, help='匹配返回的top-k结果数')

    args = parser.parse_args()

    if not os.path.exists(args.image_path):
        print(f"错误: 图像文件不存在: {args.image_path}")
        return

    # 创建模拟器
    simulator = FaceRecognitionSimulator(args.server, args.process_usage, args.network_delay)

    # 准备选项
    options = {
        'device': args.device,
        'top_k': args.top_k
    }

    try:
        if args.test_all:
            # 测试所有方案并绘制图表
            results = simulator.test_all_schemes(args.image_path, **options)
            
            # 打印详细结果
            print(f"\n=== 所有方案执行结果汇总 ===")
            for scheme_id, execution_time, description in results:
                if execution_time != float('inf'):
                    print(f"方案{scheme_id}: {execution_time:.2f}ms - {description}")
                else:
                    print(f"方案{scheme_id}: 执行失败 - {description}")
            
            # 绘制图表
            plot_scheme_times(results, args.process_usage, args.network_delay)
            
            # 同时显示最优方案
            optimal_scheme = simulator.find_optimal_scheme()
            print(f"\n预测最优方案: 方案{optimal_scheme}")
            
        elif args.find_optimal:
            # 寻找最优方案
            optimal_scheme = simulator.find_optimal_scheme()
            print(f"\n执行最优方案 {optimal_scheme}...")
            args.scheme = optimal_scheme

            # 根据指定方案执行
            if args.scheme == 0:
                simulator.pipeline_scheme_0(args.image_path, **options)
            elif args.scheme == 1:
                simulator.pipeline_scheme_1(args.image_path, **options)
            elif args.scheme == 2:
                simulator.pipeline_scheme_2(args.image_path, **options)
            elif args.scheme == 3:
                simulator.pipeline_scheme_3(args.image_path, **options)
            elif args.scheme == 4:
                simulator.pipeline_scheme_4(args.image_path, **options)
            else:
                print("请指定划分方案 (0-4) 或使用 --find-optimal 寻找最优方案")

        else:
            # 根据指定方案执行
            if args.scheme == 0:
                simulator.pipeline_scheme_0(args.image_path, **options)
            elif args.scheme == 1:
                simulator.pipeline_scheme_1(args.image_path, **options)
            elif args.scheme == 2:
                simulator.pipeline_scheme_2(args.image_path, **options)
            elif args.scheme == 3:
                simulator.pipeline_scheme_3(args.image_path, **options)
            elif args.scheme == 4:
                simulator.pipeline_scheme_4(args.image_path, **options)
            else:
                print("请指定划分方案 (0-4) 或使用 --find-optimal 寻找最优方案")

    except grpc.RpcError as e:
        print(f"RPC错误: {e.code()} - {e.details()}")
    except Exception as e:
        print(f"模拟器错误: {str(e)}")

if __name__ == '__main__':
    main()