1+ import argparse
2+ import numpy as np
3+ from pathlib import Path
4+ from probeinterface import Probe , read_probeinterface
5+ from probeinterface .plotting import plot_probe
6+
7+ import matplotlib .pyplot as plt
8+
9+ def load_probegroup_from_json (manufacturer , model_name ):
10+ """Load probegroup from JSON file based on manufacturer and model name."""
11+ json_file = Path (__file__ ).parent / manufacturer / model_name / f"{ model_name }
12+
13+ # Create probe from JSON data
14+ probegroup = read_probeinterface (json_file )
15+ return probegroup
16+
17+ def plot_and_save_probegroup (
18+ probegroup ,
19+ manufacturer ,
20+ model_name ,
21+ with_contact_id = True ,
22+ title_fontsize = 10 ,
23+ label_fontsize = 5 ,
24+ contact_id_fontsize = 5 ,
25+ save_figure = False ,
26+ zoom_on_tip = False ,
27+ dpi = 300 ,
28+ figsize = (15 , 15 )
29+ ):
30+ """Plot probegroup and optionally save to file."""
31+ probe = probegroup .probes [0 ]
32+ if zoom_on_tip :
33+ ncols = 2
34+ else :
35+ ncols = 1
36+ fig , axs = plt .subplots (figsize = figsize , ncols = ncols )
37+ if zoom_on_tip :
38+ ax_full , ax_tip = axs
39+ else :
40+ ax_full = axs
41+ plot_probe (probe , with_contact_id = with_contact_id , ax = ax_full )
42+ for txt in ax_full .texts :
43+ txt .set_fontsize (contact_id_fontsize )
44+
45+ title = f"{ manufacturer } { model_name }
46+ if probe .description :
47+ title += f"\n { probe .description }
48+ title += f"\n { probe .get_contact_count ()} { probe .get_shank_count ()}
49+ ax_full .set_ylabel ("Y (µm)" , fontsize = label_fontsize )
50+ ax_full .tick_params (axis = 'both' , which = 'major' , labelsize = label_fontsize )
51+
52+ if zoom_on_tip :
53+ ax_full .spines [['top' , 'right' ]].set_visible (False )
54+ ax_full .set_title ("" )
55+ ax_full .set_xlabel ("" )
56+ ax_tip .spines [['top' , 'right' ]].set_visible (False )
57+ plot_probe (probe , with_contact_id = with_contact_id , ax = ax_tip )
58+ shank_tip_y = np .min (probe .probe_planar_contour [:, 1 ])
59+ shank_left = np .min (probe .probe_planar_contour [:, 0 ])
60+ shank_right = np .max (probe .probe_planar_contour [:, 0 ])
61+ shank_width = shank_right - shank_left
62+ ax_tip .set_xlim (shank_left - 0.5 * shank_width , shank_right + 0.5 * shank_width )
63+ ax_tip .set_ylim (shank_tip_y - 20 , 200 )
64+ ax_tip .set_title ("" )
65+ fig .suptitle (title , fontsize = title_fontsize )
66+ else :
67+ ax_full .set_title (title , fontsize = title_fontsize )
68+ ax_full .set_xlabel ("X (µm)" , fontsize = label_fontsize )
69+
70+ if save_figure :
71+ save_dir = Path (__file__ ).parent / manufacturer / model_name
72+ save_dir .mkdir (parents = True , exist_ok = True )
73+ output_file = save_dir / f"{ model_name }
74+ fig .savefig (output_file , dpi = dpi )
75+ print (f"Image saved to: { output_file } )
76+ else :
77+ plt .show ()
78+
79+ plt .close (fig )
80+
81+ def main ():
82+ parser = argparse .ArgumentParser (description = "Generate probe images from JSON files" )
83+ parser .add_argument ("manufacturer" , help = "Probe manufacturer name" )
84+ parser .add_argument ("--model-name" , help = "Probe model name" )
85+ parser .add_argument ("--with-ids" , "-ids" , action = "store_true" , help = "Display contact IDs" )
86+ parser .add_argument ("--save-figure" , "-s" , action = "store_true" , help = "Save the figure to a file" )
87+ parser .add_argument ("--zoom-on-tip" , "-z" , action = "store_true" , help = "Zoom in on the probe tip" )
88+ parser .add_argument ("--dpi" , type = int , default = 600 , help = "DPI for saved figure" )
89+ parser .add_argument ("--title-fontsize" , "-t" , type = int , default = 15 , help = "Font size for the title" )
90+ parser .add_argument ("--label-fontsize" , "-l" , type = int , default = 10 , help = "Font size for the labels" )
91+ parser .add_argument ("--contact-id-fontsize" , "-c" , type = int , default = 2 , help = "Font size for contact IDs" )
92+
93+ args = parser .parse_args ()
94+
95+ try :
96+ if args .model_name is None :
97+ # find model names
98+ manufacturer_dir = Path (__file__ ).parent / args .manufacturer
99+ if not manufacturer_dir .exists () or not manufacturer_dir .is_dir ():
100+ raise FileNotFoundError (f"Manufacturer directory not found: { manufacturer_dir } )
101+ model_names = [d .name for d in manufacturer_dir .iterdir () if d .is_dir ()]
102+ else :
103+ model_names = [args .model_name ]
104+ for model_name in model_names :
105+ # Load probegroup from JSON
106+ probegroup = load_probegroup_from_json (args .manufacturer , model_name )
107+
108+ # Plot and save
109+ plot_and_save_probegroup (
110+ probegroup = probegroup ,
111+ manufacturer = args .manufacturer ,
112+ model_name = model_name ,
113+ save_figure = args .save_figure ,
114+ zoom_on_tip = args .zoom_on_tip ,
115+ with_contact_id = args .with_ids ,
116+ title_fontsize = args .title_fontsize ,
117+ label_fontsize = args .label_fontsize ,
118+ contact_id_fontsize = args .contact_id_fontsize ,
119+ dpi = args .dpi
120+ )
121+
122+ except FileNotFoundError as e :
123+ print (f"Error: { e } )
124+ except Exception as e :
125+ print (f"Unexpected error: { e } )
126+
127+ if __name__ == "__main__" :
128+ main ()
0 commit comments