diff --git a/selfdrive/modeld/modeld.py b/selfdrive/modeld/modeld.py index 9805506432..15b7947d66 100755 --- a/selfdrive/modeld/modeld.py +++ b/selfdrive/modeld/modeld.py @@ -38,8 +38,10 @@ SEND_RAW_PRED = os.getenv('SEND_RAW_PRED') VISION_PKL_PATH = Path(__file__).parent / 'models/driving_vision_tinygrad.pkl' POLICY_PKL_PATH = Path(__file__).parent / 'models/driving_policy_tinygrad.pkl' +MISC_PKL_PATH = Path(__file__).parent / 'models/driving_misc_tinygrad.pkl' VISION_METADATA_PATH = Path(__file__).parent / 'models/driving_vision_metadata.pkl' POLICY_METADATA_PATH = Path(__file__).parent / 'models/driving_policy_metadata.pkl' +MISC_METADATA_PATH = Path(__file__).parent / 'models/driving_misc_metadata.pkl' LAT_SMOOTH_SECONDS = 0.3 LONG_SMOOTH_SECONDS = 0.3 @@ -113,11 +115,19 @@ class ModelState: self.policy_output_slices = policy_metadata['output_slices'] policy_output_size = policy_metadata['output_shapes']['outputs'][1] + with open(MISC_METADATA_PATH, 'rb') as f: + misc_metadata = pickle.load(f) + self.misc_input_shapes = misc_metadata['input_shapes'] + self.misc_output_slices = misc_metadata['output_slices'] + misc_output_size = misc_metadata['output_shapes']['outputs'][1] + # img buffers are managed in openCL transform code self.vision_inputs: dict[str, Tensor] = {} self.vision_output = np.zeros(vision_output_size, dtype=np.float32) self.policy_inputs = {k: Tensor(v, device='NPY').realize() for k,v in self.numpy_inputs.items()} self.policy_output = np.zeros(policy_output_size, dtype=np.float32) + self.misc_inputs = {k: Tensor(v, device='NPY').realize() for k,v in self.numpy_inputs.items() if k not in {"prev_desired_curv", "lateral_control_params"}} + self.misc_output = np.zeros(misc_output_size, dtype=np.float32) self.parser = Parser() with open(VISION_PKL_PATH, "rb") as f: @@ -126,6 +136,9 @@ class ModelState: with open(POLICY_PKL_PATH, "rb") as f: self.policy_run = pickle.load(f) + with open(MISC_PKL_PATH, "rb") as f: + self.misc_run = pickle.load(f) + def slice_outputs(self, model_outputs: np.ndarray, output_slices: dict[str, slice]) -> dict[str, np.ndarray]: parsed_model_outputs = {k: model_outputs[np.newaxis, v] for k,v in output_slices.items()} return parsed_model_outputs @@ -169,12 +182,15 @@ class ModelState: self.policy_output = self.policy_run(**self.policy_inputs).numpy().flatten() policy_outputs_dict = self.parser.parse_policy_outputs(self.slice_outputs(self.policy_output, self.policy_output_slices)) + self.misc_output = self.misc_run(**self.misc_inputs).numpy().flatten() + misc_outputs_dict = self.parser.parse_misc_outputs(self.slice_outputs(self.misc_output, self.misc_output_slices)) + # TODO model only uses last value now self.full_prev_desired_curv[0,:-1] = self.full_prev_desired_curv[0,1:] self.full_prev_desired_curv[0,-1,:] = policy_outputs_dict['desired_curvature'][0, :] self.numpy_inputs['prev_desired_curv'][:] = 0*self.full_prev_desired_curv[0, self.temporal_idxs] - combined_outputs_dict = {**vision_outputs_dict, **policy_outputs_dict} + combined_outputs_dict = {**vision_outputs_dict, **policy_outputs_dict, **misc_outputs_dict} if SEND_RAW_PRED: combined_outputs_dict['raw_pred'] = np.concatenate([self.vision_output.copy(), self.policy_output.copy()])