Update generate_test_data.py to not use detect function

MouseLand · Sep 5, 2022 · 636629f · 636629f
1 parent 04b09a7
commit 636629f
Showing 1 changed file with 23 additions and 29 deletions.
diff --git a/scripts/generate_test_data.py b/scripts/generate_test_data.py
@@ -32,14 +32,6 @@ def generate_1p1c1500_expected_data(ops):
 		suite2p.run_s2p(ops=test_ops)
 		rename_output_dir('1plane1chan1500')
 
-	# def generate_1p2c_expected_data(ops):
-	#   """
-	#   Generates expected output for test_1plane_2chan_sourcery of test_full_pipeline.py.
-	#   """
-	#   test_ops = FullPipelineTestUtils.initialize_ops_test_1plane_2chan_sourcery(ops.copy())
-	#   suite2p.run_s2p(ops=test_ops)
-	#   rename_output_dir('1plane2chan')
-
 	def generate_2p2c1500_expected_data(ops):
 		"""
 		Generates expected output for test_2plane_2chan_with_batches of test_full_pipeline.py.
@@ -78,17 +70,18 @@ def generate_detection_1plane1chan_test_data(ops):
 			[[Path(ops['data_path'][0]).joinpath('detection/pre_registered.npy')]],
 			(404, 360)
 		)
-		ops, stat = suite2p.detection.detect(ops[0])
-		ops['neuropil_extract'] = True
-		cell_masks, neuropil_masks = masks.create_masks(stat, ops['Ly'], ops['Lx'], ops=ops)
-		output_dict = {
-			'stat': stat,
-			'cell_masks': cell_masks,
-			'neuropil_masks': neuropil_masks
-		}
-		np.save('expected_detect_output_1p1c0.npy', output_dict)
-		# Remove suite2p directory generated by prepare function
-		shutil.rmtree(os.path.join(test_data_dir_path, 'suite2p'))
+		with suite2p.io.BinaryRWFile(Ly = ops[0]['Ly'], Lx = ops[0]['Lx'], filename=ops[0]['reg_file']) as f_reg:
+			ops, stat = suite2p.detection.detection_wrapper(f_reg, ops=ops[0])
+			ops['neuropil_extract'] = True
+			cell_masks, neuropil_masks = masks.create_masks(stat, ops['Ly'], ops['Lx'], ops=ops)
+			output_dict = {
+				'stat': stat,
+				'cell_masks': cell_masks,
+				'neuropil_masks': neuropil_masks
+			}
+			np.save('expected_detect_output_1p1c0.npy', output_dict)
+			# Remove suite2p directory generated by prepare function
+			shutil.rmtree(os.path.join(test_data_dir_path, 'suite2p'))
 
 	def generate_detection_2plane2chan_test_data(ops):
 		"""
@@ -114,16 +107,17 @@ def generate_detection_2plane2chan_test_data(ops):
 		two_plane_ops[1]['meanImg_chan2'] = np.load(detection_dir.joinpath('meanImg_chan2p1.npy'))
 		for i in range(len(two_plane_ops)):
 			op = two_plane_ops[i]
-			# Neuropil_masks are later needed for extraction test data step
-			op['neuropil_extract'] = True
-			op, stat = suite2p.detection.detect(ops=op)
-			cell_masks, neuropil_masks = masks.create_masks(stat, op['Ly'], op['Lx'], ops=op)
-			output_dict = {
-				'stat': stat,
-				'cell_masks': cell_masks,
-				'neuropil_masks': neuropil_masks
-			}
-			np.save('expected_detect_output_%ip%ic%i.npy' % (ops['nchannels'], ops['nplanes'], i), output_dict)
+			with suite2p.io.BinaryRWFile(Ly = op['Ly'], Lx = op['Lx'], filename=op['reg_file']) as f_reg:
+				# Neuropil_masks are later needed for extraction test data step
+				op['neuropil_extract'] = True
+				op, stat = suite2p.detection.detection_wrapper(f_reg, ops=op)
+				cell_masks, neuropil_masks = masks.create_masks(stat, op['Ly'], op['Lx'], ops=op)
+				output_dict = {
+					'stat': stat,
+					'cell_masks': cell_masks,
+					'neuropil_masks': neuropil_masks
+				}
+				np.save('expected_detect_output_%ip%ic%i.npy' % (ops['nchannels'], ops['nplanes'], i), output_dict)
 		# Get rid of registered binary files that were created for detection module in 
 		# DetectionTestUtils.prepare
 		remove_binary_file(test_data_dir_path, 0, '')