import argparse import logging import multiprocessing import os import pickle import time from functools import partial import h5py import numpy as np import pandas as pd import tensorflow as tf from data_reader import DataReader_mseed_array, DataReader_pred from postprocess import ( extract_amplitude, extract_picks, save_picks, save_picks_json, save_prob_h5, ) from tqdm import tqdm from visulization import plot_waveform from model import ModelConfig, UNet tf.compat.v1.disable_eager_execution() tf.compat.v1.logging.set_verbosity(tf.compat.v1.logging.ERROR) def read_args(): parser = argparse.ArgumentParser() parser.add_argument("--batch_size", default=20, type=int, help="batch size") parser.add_argument("--model_dir", help="Checkpoint directory (default: None)") parser.add_argument("--data_dir", default="", help="Input file directory") parser.add_argument("--data_list", default="", help="Input csv file") parser.add_argument("--hdf5_file", default="", help="Input hdf5 file") parser.add_argument("--hdf5_group", default="data", help="data group name in hdf5 file") parser.add_argument("--result_dir", default="results", help="Output directory") parser.add_argument("--result_fname", default="picks", help="Output file") parser.add_argument("--min_p_prob", default=0.3, type=float, help="Probability threshold for P pick") parser.add_argument("--min_s_prob", default=0.3, type=float, help="Probability threshold for S pick") parser.add_argument("--mpd", default=50, type=float, help="Minimum peak distance") parser.add_argument("--amplitude", action="store_true", help="if return amplitude value") parser.add_argument("--format", default="numpy", help="input format") parser.add_argument("--s3_url", default="localhost:9000", help="s3 url") parser.add_argument("--stations", default="", help="seismic station info") parser.add_argument("--plot_figure", action="store_true", help="If plot figure for test") parser.add_argument("--save_prob", action="store_true", help="If save result for test") parser.add_argument("--pre_sec", default=1, type=float, help="Window length before pick") parser.add_argument("--post_sec", default=4, type=float, help="Window length after pick") parser.add_argument("--highpass_filter", default=0.0, type=float, help="Highpass filter") parser.add_argument("--response_xml", default=None, type=str, help="response xml file") parser.add_argument("--sampling_rate", default=100, type=float, help="sampling rate") args = parser.parse_args() return args def pred_fn(args, data_reader, figure_dir=None, prob_dir=None, log_dir=None): current_time = time.strftime("%y%m%d-%H%M%S") if log_dir is None: log_dir = os.path.join(args.log_dir, "pred", current_time) if not os.path.exists(log_dir): os.makedirs(log_dir) if (args.plot_figure == True) and (figure_dir is None): figure_dir = os.path.join(log_dir, "figures") if not os.path.exists(figure_dir): os.makedirs(figure_dir) if (args.save_prob == True) and (prob_dir is None): prob_dir = os.path.join(log_dir, "probs") if not os.path.exists(prob_dir): os.makedirs(prob_dir) if args.save_prob: h5 = h5py.File(os.path.join(args.result_dir, "result.h5"), "w", libver="latest") prob_h5 = h5.create_group("/prob") logging.info("Pred log: %s" % log_dir) logging.info("Dataset size: {}".format(data_reader.num_data)) with tf.compat.v1.name_scope("Input_Batch"): if args.format == "mseed_array": batch_size = 1 else: batch_size = args.batch_size dataset = data_reader.dataset(batch_size) batch = tf.compat.v1.data.make_one_shot_iterator(dataset).get_next() config = ModelConfig(X_shape=data_reader.X_shape) with open(os.path.join(log_dir, "config.log"), "w") as fp: fp.write("\n".join("%s: %s" % item for item in vars(config).items())) model = UNet(config=config, input_batch=batch, mode="pred") # model = UNet(config=config, mode="pred") sess_config = tf.compat.v1.ConfigProto() sess_config.gpu_options.allow_growth = True # sess_config.log_device_placement = False with tf.compat.v1.Session(config=sess_config) as sess: saver = tf.compat.v1.train.Saver(tf.compat.v1.global_variables(), max_to_keep=5) init = tf.compat.v1.global_variables_initializer() sess.run(init) latest_check_point = tf.train.latest_checkpoint(args.model_dir) logging.info(f"restoring model {latest_check_point}") saver.restore(sess, latest_check_point) picks = [] amps = [] if args.amplitude else None if args.plot_figure: multiprocessing.set_start_method("spawn") pool = multiprocessing.Pool(multiprocessing.cpu_count()) for _ in tqdm(range(0, data_reader.num_data, batch_size), desc="Pred"): if args.amplitude: pred_batch, X_batch, amp_batch, fname_batch, t0_batch, station_batch = sess.run( [model.preds, batch[0], batch[1], batch[2], batch[3], batch[4]], feed_dict={model.drop_rate: 0, model.is_training: False}, ) # X_batch, amp_batch, fname_batch, t0_batch = sess.run([batch[0], batch[1], batch[2], batch[3]]) else: pred_batch, X_batch, fname_batch, t0_batch, station_batch = sess.run( [model.preds, batch[0], batch[1], batch[2], batch[3]], feed_dict={model.drop_rate: 0, model.is_training: False}, ) # X_batch, fname_batch, t0_batch = sess.run([model.preds, batch[0], batch[1], batch[2]]) # pred_batch = [] # for i in range(0, len(X_batch), 1): # pred_batch.append(sess.run(model.preds, feed_dict={model.X: X_batch[i:i+1], model.drop_rate: 0, model.is_training: False})) # pred_batch = np.vstack(pred_batch) waveforms = None if args.amplitude: waveforms = amp_batch picks_ = extract_picks( preds=pred_batch, file_names=fname_batch, station_ids=station_batch, begin_times=t0_batch, config=args, waveforms=waveforms, use_amplitude=args.amplitude, dt=1.0 / args.sampling_rate, ) picks.extend(picks_) ## save pick per file if len(fname_batch) == 1: df = pd.DataFrame(picks_) df = df[df["phase_index"] > 10] if not os.path.exists(os.path.join(args.result_dir, "picks")): os.makedirs(os.path.join(args.result_dir, "picks")) df = df[ [ "station_id", "begin_time", "phase_index", "phase_time", "phase_score", "phase_type", "phase_amplitude", "dt", ] ] df.to_csv( os.path.join( args.result_dir, "picks", fname_batch[0].decode().split("/")[-1].rstrip(".mseed") + ".csv" ), index=False, ) if args.plot_figure: if not (isinstance(fname_batch, np.ndarray) or isinstance(fname_batch, list)): fname_batch = [fname_batch.decode().rstrip(".mseed") + "_" + x.decode() for x in station_batch] else: fname_batch = [x.decode() for x in fname_batch] pool.starmap( partial( plot_waveform, figure_dir=figure_dir, ), # zip(X_batch, pred_batch, [x.decode() for x in fname_batch]), zip(X_batch, pred_batch, fname_batch), ) if args.save_prob: # save_prob(pred_batch, fname_batch, prob_dir=prob_dir) if not (isinstance(fname_batch, np.ndarray) or isinstance(fname_batch, list)): fname_batch = [fname_batch.decode().rstrip(".mseed") + "_" + x.decode() for x in station_batch] else: fname_batch = [x.decode() for x in fname_batch] save_prob_h5(pred_batch, fname_batch, prob_h5) if len(picks) > 0: # save_picks(picks, args.result_dir, amps=amps, fname=args.result_fname+".csv") # save_picks_json(picks, args.result_dir, dt=data_reader.dt, amps=amps, fname=args.result_fname+".json") df = pd.DataFrame(picks) # df["fname"] = df["file_name"] # df["id"] = df["station_id"] # df["timestamp"] = df["phase_time"] # df["prob"] = df["phase_prob"] # df["type"] = df["phase_type"] base_columns = [ "station_id", "begin_time", "phase_index", "phase_time", "phase_score", "phase_type", "file_name", ] if args.amplitude: base_columns.append("phase_amplitude") base_columns.append("phase_amp") df["phase_amp"] = df["phase_amplitude"] df = df[base_columns] df.to_csv(os.path.join(args.result_dir, args.result_fname + ".csv"), index=False) print( f"Done with {len(df[df['phase_type'] == 'P'])} P-picks and {len(df[df['phase_type'] == 'S'])} S-picks" ) else: print(f"Done with 0 P-picks and 0 S-picks") return 0 def main(args): logging.basicConfig(format="%(asctime)s %(message)s", level=logging.INFO) with tf.compat.v1.name_scope("create_inputs"): if args.format == "mseed_array": data_reader = DataReader_mseed_array( data_dir=args.data_dir, data_list=args.data_list, stations=args.stations, amplitude=args.amplitude, highpass_filter=args.highpass_filter, ) else: data_reader = DataReader_pred( format=args.format, data_dir=args.data_dir, data_list=args.data_list, hdf5_file=args.hdf5_file, hdf5_group=args.hdf5_group, amplitude=args.amplitude, highpass_filter=args.highpass_filter, response_xml=args.response_xml, sampling_rate=args.sampling_rate, ) pred_fn(args, data_reader, log_dir=args.result_dir) return if __name__ == "__main__": args = read_args() main(args)