diff --git a/sunnypilot/modeld/modeld.py b/sunnypilot/modeld/modeld.py index 7339776e75..91d9f84d60 100755 --- a/sunnypilot/modeld/modeld.py +++ b/sunnypilot/modeld/modeld.py @@ -1,11 +1,9 @@ #!/usr/bin/env python3 import os import time -import pickle import numpy as np import cereal.messaging as messaging from cereal import car, log -from pathlib import Path from setproctitle import setproctitle from cereal.messaging import PubMaster, SubMaster from msgq.visionipc import VisionIpcClient, VisionStreamType, VisionBuf @@ -24,15 +22,11 @@ from openpilot.sunnypilot.modeld.fill_model_msg import fill_model_msg, fill_pose from openpilot.sunnypilot.modeld.constants import ModelConstants from openpilot.sunnypilot.modeld.models.commonmodel_pyx import DrivingModelFrame, CLContext -from openpilot.sunnypilot.modeld.runners.run_helpers import METADATA_PATH, prepare_inputs, parse_runner_inputs +from openpilot.sunnypilot.modeld.runners.run_helpers import load_model, load_metadata, prepare_inputs, parse_runner_inputs PROCESS_NAME = "sunnypilot.modeld.modeld" SEND_RAW_PRED = os.getenv('SEND_RAW_PRED') -MODEL_PATHS = { - ModelRunner.THNEED: Path(__file__).parent / 'models/supercombo.thneed', - ModelRunner.ONNX: Path(__file__).parent / 'models/supercombo.onnx'} - class FrameMeta: frame_id: int = 0 @@ -61,15 +55,15 @@ class ModelState: # img buffers are managed in openCL transform code self.inputs = prepare_inputs() - with open(METADATA_PATH, 'rb') as f: - model_metadata = pickle.load(f) + model_paths = load_model() + model_metadata = load_metadata() self.output_slices = model_metadata['output_slices'] net_output_size = model_metadata['output_shapes']['outputs'][1] self.output = np.zeros(net_output_size, dtype=np.float32) self.parser = Parser() - self.model = ModelRunner(MODEL_PATHS, self.output, Runtime.GPU, False, context) + self.model = ModelRunner(model_paths, self.output, Runtime.GPU, False, context) self.model.addInput("input_imgs", None) self.model.addInput("big_input_imgs", None) for k,v in self.inputs.items(): diff --git a/sunnypilot/modeld/runners/run_helpers.py b/sunnypilot/modeld/runners/run_helpers.py index bb8e8bb114..37503c8bfc 100644 --- a/sunnypilot/modeld/runners/run_helpers.py +++ b/sunnypilot/modeld/runners/run_helpers.py @@ -3,15 +3,44 @@ # This file is part of sunnypilot and is licensed under the MIT License. # See the LICENSE.md file in the root directory for more details. +import os import pickle import numpy as np from pathlib import Path +from cereal import custom +from openpilot.sunnypilot.modeld.runners import ModelRunner +from openpilot.sunnypilot.models.helpers import get_active_bundle +from openpilot.system.hardware import PC +from openpilot.system.hardware.hw import Paths +USE_ONNX = int(os.getenv('USE_ONNX', str(int(PC)))) + +CUSTOM_MODEL_PATH = Paths.model_root() METADATA_PATH = Path(__file__).parent / 'models/supercombo_metadata.pkl' +ModelManager = custom.ModelManagerSP + + +def load_model(): + if USE_ONNX: + model_paths = {ModelRunner.ONNX: Path(__file__).parent / 'models/supercombo.onnx'} + elif bundle := get_active_bundle(): + drive_model = next(model for model in bundle.models if model.type == ModelManager.Type.drive) + model_paths = {ModelRunner.THNEED: f"{CUSTOM_MODEL_PATH}/{drive_model.fileName}"} + else: + model_paths = {ModelRunner.THNEED: Path(__file__).parent / 'models/supercombo.thneed'} + + return model_paths + def load_metadata(): - with open(METADATA_PATH, 'rb') as f: + if bundle := get_active_bundle(): + drive_model = next(model for model in bundle.models if model.type == ModelManager.Type.metadata) + metadata_path = f"{CUSTOM_MODEL_PATH}/{drive_model.fileName}" + else: + metadata_path = METADATA_PATH + + with open(metadata_path, 'rb') as f: metadata = pickle.load(f) return metadata