diff --git a/common/models.py b/common/models.py index 00054c5..efe21b9 100644 --- a/common/models.py +++ b/common/models.py @@ -72,6 +72,7 @@ class PipelineContext(BaseModel): lazy_read_input: bool = Field(default=True, description='Lazy read input file') write_recording: bool = Field(default=False, description='Write recording') stub_test: bool = Field(default=False, description='Stub test') + stub_test_duration_sec: float = Field(default=300, description='Stub test duration in seconds') job_kwargs: JobKwargs = Field(default=JobKwargs(), description='Job kwargs') recording_context: RecordingContext = Field(description='Recording context') run_preprocessing: bool = Field(default=True, description='Run preprocessing') diff --git a/common/processor_pipeline.py b/common/processor_pipeline.py index 5c4289b..c7f477c 100644 --- a/common/processor_pipeline.py +++ b/common/processor_pipeline.py @@ -44,7 +44,8 @@ def run_pipeline(context: PipelineContext): if context.stub_test: logger.info('Running in stub test mode') - n_frames = int(min(300_000, recording.get_num_frames())) + stub_test_num_frames = context.stub_test_duration_sec * recording.get_sampling_frequency() + n_frames = int(min(stub_test_num_frames, recording.get_num_frames())) recording = recording.frame_slice(start_frame=0, end_frame=n_frames) logger.info(recording) diff --git a/si_kilosort25/spec.json b/si_kilosort25/spec.json index a97e70c..a2f5128 100644 --- a/si_kilosort25/spec.json +++ b/si_kilosort25/spec.json @@ -15,12 +15,14 @@ "description": "Input NWB file" } ], + "inputFolders": [], "outputs": [ { "name": "output", "description": "Output NWB file" } ], + "outputFolders": [], "parameters": [ { "name": "lazy_read_input", @@ -40,6 +42,12 @@ "type": "bool", "default": false }, + { + "name": "stub_test_duration_sec", + "description": "Stub test duration in seconds", + "type": "float", + "default": 300 + }, { "name": "job_kwargs.n_jobs", "description": "Number of jobs, must be a positive number between 0 and 1, or -1 for all processors.",