diff --git a/ipi_extract.py b/ipi_extract.py index c1882549df5ce2893832b2704b9959bc7cd55f02..079a050473b1c2ab38e485f80d34d2c838738fbf 100644 --- a/ipi_extract.py +++ b/ipi_extract.py @@ -707,10 +707,22 @@ def main(args): outpath = args.input.rsplit('.', 1)[0] + '.pred.h5' else: outpath = args.out - if os.path.isfile(outpath) and not args.erase: - print(f'Out file {outpath} already exist and erase option isn\'t set.') + if os.path.isfile(outpath): + if not (args.erase or args.resume): + print(f'Out file {outpath} already exist and erase option isn\'t set.') + return 1 + elif args.resume: + print(f'Out file {outpath} does not already exist and resume option is set.') return 1 + ref_dict = init(args.input, args.channel) + if args.resume: + df = pd.read_hdf(outpath) + ref_dict['callback'].df = df.to_dict(orient='index') + ref_dict['callback'].offset = np.tile(np.array(list(ref_dict['callback'].df.keys()))[:, np.newaxis], (1,2)) + ref_dict['callback'].offset[:, 1] = ref_dict['callback'].song_resample[(ref_dict['callback'].offset[:, 1] * FSSR).astype(int)] + ref_dict['callback'].scat.set_offsets(ref_dict['callback'].offset) + ref_dict['callback'].scat.set_color(np.array(len(ref_dict['callback'].offset)*[[0.7, 0.2, 0.5, 1]])) plt.draw() plt.pause(0.2) ref_dict['fig'].set_constrained_layout(False) @@ -732,8 +744,11 @@ if __name__ == '__main__': parser.add_argument("input", type=str, help="Input file") parser.add_argument("--out", type=str, default='', help="Output file. Default to the input_path'.pred.h5'") parser.add_argument("--channel", type=int, default=0, help="Sound channel to be analysed. Indices start from 0.") - parser.add_argument("--erase", action='store_true', help="If out file exist and this option is not given," + group = parser.add_mutually_exclusive_group() + group.add_argument("--erase", action='store_true', help="If out file exist and this option is not given," " the computation will be halted") + group.add_argument("--resume", action='store_true', help="If out file exist and this option is given," + " the previous annotation file will be loaded") try: args = parser.parse_args() except ValueError as e: @@ -756,6 +771,10 @@ if __name__ == '__main__': self.out = input("What is the out file path? (Leave empty for default)") self.channel = int(input("Which channel do you want to use starting from 0? ")) self.erase = ask("Do you want to erase the out file if it already exist?") + if not self.erase: + self.resume = ask("Do you want to resume the out file if it already exist?") + else: + self.resume = False args = VirtualArgParse()