def get_size_training_chunk(task, date, shuffle, snapshot=4999, chunk=0, nc=200, ns=10): # task = 'ibl3' # date = '2020-03-16' # ns = 10 # nc = 200 # %% # shuffle = 0 data_info = DataLoader(task) # snapshot = 4999 # set to snapshot # to overwrite snapshot read # chunk_size = 1000 # distribute among gpus num_chunks = 1 cfg = get_model_config(task, data_info.model_data_dir, scorer=data_info.scorer, date=date) dlc_cfg = get_train_config(cfg, shuffle) trainingsnapshot_name, snapshot, dlc_cfg = load_dlc_snapshot( dlc_cfg, overwrite_snapshot=snapshot) # %% Load training data # %% dataset = train_loader(dlc_cfg) # %% outputdir = dataset.resnet_out_chunk_path / ('ns{}nc{}'.format(ns, nc)) # %% if not outputdir.exists() or (len(os.listdir(str(outputdir))) == 0): raise FileExistsError('\n\nRun training process\n') # store_training_features_v1(dlc_cfg, ns=ns, num_chunks=num_chunks,chunk_size=chunk_size, nc=nc) # %% print('\n Reading stored resnet outputs in {}'.format(outputdir)) # load a single chunk # num_chunks = len(os.listdir(outputdir)) resnet_out_chunk = dataset.load_resnet_outs_chunk(chunk=chunk, ns=ns, nc=nc)[0] # %% # test frame is 0 nt_chunk = resnet_out_chunk.shape[0] return nt_chunk
# %% task = 'ibl1' date = '2031-05-01' shuffle = 1 snapshot = 5000 # %% data_info = DataLoader(task) # construct proj name cfg = get_model_config(task, data_info.model_data_dir, scorer=data_info.scorer, date=date) dlc_cfg = get_train_config(cfg, shuffle) # load training file trainingsnapshot_name, snapshot, dlc_cfg = load_dlc_snapshot( dlc_cfg, overwrite_snapshot=snapshot) # %% resout_from_file = extract_resnet_output(task=task, date=date, shuffle=shuffle, overwrite_snapshot=snapshot) # %% debug_key = '' # Load training data dataset = train_loader(dlc_cfg, debug_key=debug_key) # %% images = dataset.load_train_data()[0] # %%
def store_resnet_output(task, date, shuffle, overwrite_snapshot=None, allow_growth=True, videofile_path=None, resnet_output_dir=None): from deeplabcut.pose_estimation_tensorflow.nnet.net_factory import pose_net """ task = 'ibl1' date = '2020-01-25' shuffle = 1 overwrite_snapshot = 5000 ibl_chunk_path = '/data/libraries/deepgraphpose/etc/lab_ekb/debug_va_semi_pipeline/run_long_video_aqweight/movies' videofile_path = Path(ibl_chunk_path) / 'movie_chunk_00.mp4' resnet_output_dir = videofile_path.parent / videofile_path.stem allow_growth = True """ #%% if isinstance(videofile_path, str): videofile_path = Path(videofile_path) else: pass # assert isinstance(videofile_path, Path) if isinstance(resnet_output_dir, str): resnet_output_dir = Path(resnet_output_dir) else: pass #assert isinstance(resnet_output_dir, Path) #%% data_info = DataLoader(task) cfg = get_model_config(task, data_info.model_data_dir, scorer=data_info.scorer, date=date) #%% dlc_cfg = get_train_config(cfg, shuffle) # dlc_cfg_init = edict(dlc_cfg) #%% trainingsnapshot_name, trainingsnapshot, dlc_cfg = load_dlc_snapshot( dlc_cfg, overwrite_snapshot=overwrite_snapshot) if not overwrite_snapshot == None: assert trainingsnapshot == overwrite_snapshot # dlc_cfg_init = edict(dlc_cfg) #%% Update dlc_cfg files just to init network dlc_cfg["batch_size"] = 1 dlc_cfg["num_outputs"] = cfg.get("num_outputs", 1) dlc_cfg["deterministic"] = True # %% # Load data if videofile_path is None: videofile_path = Path(data_info.videofile_path) #%% cap = cv2.VideoCapture(str(videofile_path)) nframes = int(cap.get(7)) #%% # We want to pass all frames through network nx_in, ny_in = int(cap.get(4)), int(cap.get(3)) frames = np.zeros((nframes, nx_in, ny_in, 3), dtype="ubyte") # this keeps all frames in a batch pbar = tqdm(total=nframes) counter = 0 step = nframes // 3 #max(10, int(nframes / 100)) while cap.isOpened(): if counter % step == 0: pbar.update(step) ret, frame = cap.read() if ret: frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB) frames[counter] = img_as_ubyte(frame) counter = counter + 1 else: print("the end") break pbar.close() # read all frames assert counter == nframes # %% TF.reset_default_graph() inputs = TF.placeholder(tf.float32, shape=[1, nx_in, ny_in, 3]) pn = pose_net(dlc_cfg) # extract freatures using resnet net, end_points = pn.extract_features(inputs) # heads = pn.prediction_layers(net, end_points) # %% # always restore from snapshot do not restore from IBL if trainingsnapshot == 0: variables_to_restore = slim.get_variables_to_restore( include=["resnet_v1"]) else: variables_to_restore = slim.get_variables_to_restore() restorer = TF.train.Saver(variables_to_restore) # Init session config_TF = TF.ConfigProto() config_TF.gpu_options.allow_growth = allow_growth sess = TF.Session(config=config_TF) sess.run(TF.global_variables_initializer()) sess.run(TF.local_variables_initializer()) # %% # Restore the one variable from disk restorer.restore(sess, dlc_cfg.init_weights) #%% if resnet_output_dir is None: resnet_output_dir = Path(dlc_cfg.init_weights).parent print(resnet_output_dir) #%% resnet_outdir = resnet_output_dir / "resnet_output_mat" / ( "{}".format(trainingsnapshot_name)) if not resnet_outdir.exists(): os.makedirs(resnet_outdir) #%% for ii in range(nframes): if ii % 10 == 0: print("iter {}/{}".format(ii, nframes)) ff = frames[ii, :, :, :] ff = np.expand_dims(ff, axis=0) [net_output] = sess.run([net], feed_dict={inputs: ff}) # net_heads = sess.run(heads, feed_dict={inputs: ff}) ss = resnet_outdir / ("resnet_output_{:03d}.mat".format(ii)) sio.savemat(str(ss), {"net_output": net_output}) print("Stored resnet outputs in:\n{}".format(resnet_outdir)) #%% sess.close() return
def get_targets(task, date, shuffle): from deeplabcut.pose_estimation_tensorflow.dataset.factory import ( create as create_dataset, ) #%% #task = 'reach' #date = '2020-02-19' #shuffle = 0 #%% data_info = DataLoader(task) # Load project configuration cfg = get_model_config(task, data_info.model_data_dir, scorer=data_info.scorer, date=date) # Load training configuration for cfg file dlc_cfg = get_train_config(cfg, shuffle=shuffle) # update _, _, dlc_cfg = load_dlc_snapshot(dlc_cfg) # %% dlc_cfg["deterministic"] = True #%% Create dataset dataset = create_dataset(dlc_cfg) nt = len(dataset.data) # number of training frames assert nt >= 1 nj = max([dat_.joints[0].shape[0] for dat_ in dataset.data]) ncolors, nxraw, nyraw = dataset.data[0].im_size #%% stride = dlc_cfg['stride'] #%% #clip = VideoFileClip(str(dlc_cfg.video_path)) #%% # TO DO: make exact should be around 1/8 for resnet 50 nx = nxraw // 4 ny = nyraw // 4 #%% extract_frame_num = lambda x: int(x.split("/")[-1].split(".")[0][3:]) frame_ids = [] # frame_imgs = [] counter = 0 #datas = [] datas = np.zeros((nt, nx, ny, nj)) joinss = [] #%% while counter < (nt): # for ii in range(nt): data = dataset.next_batch() data_keys = list(data.keys()) # inputs = 0 # part_score_targets = 1 # part_score_weights = 2 # locref_targets = 3 # locref_mask = 4 # pairwise_targets = 5 # pairwise_mask = 6 # data_item = 7 data_item = data[data_keys[-1]] joinss += [np.copy(data_item.joints[0])] inputs = data[data_keys[0]].squeeze() print(inputs.shape) part_score_targets = data[data_keys[1]].squeeze() # multi class labels nx, ny = part_score_targets.shape[0], part_score_targets.shape[1] # part_score_weights = data[data_keys[2]].squeeze() # locref_targets = data[data_keys[3]].squeeze() # locref_mask = data[data_keys[4]].squeeze() frame_id = extract_frame_num(data[data_keys[5]].im_path) # print(part_score_targets.max((0, 1))) print(frame_id) if frame_id in frame_ids: continue else: print("Adding frame {}".format(frame_id)) frame_ids.append(frame_id) datas[counter, :nx, :ny, :] = part_score_targets #datas.append(part_score_targets) counter += 1 #%% #datas = np.stack(datas, 0) datas = datas[:, :nx, :ny, :] #nx, ny = datas.shape[1:3] frame_ids = np.asarray(frame_ids) #%% # ignore when tongue is not present? # assert datas.sum((1, 2)).min() >= 1 print(datas.sum((1, 2)).min()) # %% # To find 2D coordinates, we must update target2d_train = np.zeros((nt, nj, 2)) * np.nan # nt x nj x 2 for ntt in range(nt): cjoin = joinss[ntt] # D x nj njtt = cjoin.shape[0] for njj_id in range(njtt): njj = cjoin[njj_id][0] joinss_ntt_njj = cjoin[njj_id][1:] target2d_train[ntt, njj] = np.flip( (joinss_ntt_njj - stride / 2) / stride) # %% # visualize some randomly # ntt = 23 # njj = 0 # center_loc = target2d_train[ntt, njj] # plt.imshow(datas[ntt, :, :, njj]) # plt.plot(center_loc[1], center_loc[0], 'ro') # plt.show(block=True) # %% # target is shared for all snapshots regargless of whathever else nx_out = int(part_score_targets.shape[0]) ny_out = int(part_score_targets.shape[1]) return frame_ids, target2d_train, nx_out, ny_out, datas