def _benchmark_paf_graphs( config, inference_cfg, data, paf_inds, greedy=False, add_discarded=True, identity_only=False, calibration_file="", oks_sigma=0.1, ): n_multi = len(auxfun_multianimal.extractindividualsandbodyparts(config)[2]) data_ = {"metadata": data.pop("metadata")} for k, v in data.items(): data_[k] = v["prediction"] ass = Assembler( data_, max_n_individuals=inference_cfg["topktoretain"], n_multibodyparts=n_multi, greedy=greedy, pcutoff=inference_cfg.get("pcutoff", 0.1), min_affinity=inference_cfg.get("pafthreshold", 0.1), add_discarded=add_discarded, identity_only=identity_only, ) if calibration_file: ass.calibrate(calibration_file) params = ass.metadata image_paths = params["imnames"] bodyparts = params["joint_names"] idx = (data[image_paths[0]]["groundtruth"][2].unstack("coords").reindex( bodyparts, level="bodyparts").index) individuals = idx.get_level_values("individuals").unique() n_individuals = len(individuals) map_ = dict(zip(individuals, range(n_individuals))) # Form ground truth beforehand ground_truth = [] for i, imname in enumerate(image_paths): temp = data[imname]["groundtruth"][2] ground_truth.append(temp.to_numpy().reshape((-1, 2))) ground_truth = np.stack(ground_truth) temp = np.ones((*ground_truth.shape[:2], 3)) temp[..., :2] = ground_truth temp = temp.reshape((temp.shape[0], n_individuals, -1, 3)) ass_true_dict = _parse_ground_truth_data(temp) ids = np.vectorize(map_.get)( idx.get_level_values("individuals").to_numpy()) ground_truth = np.insert(ground_truth, 2, ids, axis=2) # Assemble animals on the full set of detections paf_inds = sorted(paf_inds, key=len) paf_graph = ass.graph n_graphs = len(paf_inds) all_scores = [] all_metrics = [] for j, paf in enumerate(paf_inds, start=1): print(f"Graph {j}|{n_graphs}") graph = [paf_graph[i] for i in paf] ass.paf_inds = paf ass.graph = graph ass.assemble() oks = evaluate_assembly(ass.assemblies, ass_true_dict, oks_sigma) all_metrics.append(oks) scores = np.full((len(image_paths), 2), np.nan) for i, imname in enumerate(tqdm(image_paths)): gt = ground_truth[i] gt = gt[~np.isnan(gt).any(axis=1)] if len(np.unique( gt[:, 2])) < 2: # Only consider frames with 2+ animals continue # Count the number of unassembled bodyparts n_dets = len(gt) animals = ass.assemblies.get(i) if animals is None: if n_dets: scores[i, 0] = 1 else: animals = [ np.c_[animal.data, np.ones(animal.data.shape[0]) * n] for n, animal in enumerate(animals) ] hyp = np.concatenate(animals) hyp = hyp[~np.isnan(hyp).any(axis=1)] scores[i, 0] = (n_dets - hyp.shape[0]) / n_dets neighbors = _find_closest_neighbors(gt[:, :2], hyp[:, :2]) valid = neighbors != -1 id_gt = gt[valid, 2] id_hyp = hyp[neighbors[valid], -1] mat = contingency_matrix(id_gt, id_hyp) purity = mat.max(axis=0).sum() / mat.sum() scores[i, 1] = purity all_scores.append((scores, paf)) dfs = [] for score, inds in all_scores: df = pd.DataFrame(score, columns=["miss", "purity"]) df["ngraph"] = len(inds) dfs.append(df) big_df = pd.concat(dfs) group = big_df.groupby("ngraph") return (all_scores, group.agg(["mean", "std"]).T, all_metrics)
def create_video_with_all_detections( config, videos, shuffle=1, trainingsetindex=0, displayedbodyparts="all", destfolder=None, modelprefix="", ): """ Create a video labeled with all the detections stored in a '*_full.pickle' file. Parameters ---------- config : str Absolute path to the config.yaml file videos : list of str A list of strings containing the full paths to videos for analysis or a path to the directory, where all the videos with same extension are stored. shuffle : int, optional Number of shuffles of training dataset. Default is set to 1. trainingsetindex: int, optional Integer specifying which TrainingsetFraction to use. By default the first (note that TrainingFraction is a list in config.yaml). displayedbodyparts: list of strings, optional This selects the body parts that are plotted in the video. Either ``all``, then all body parts from config.yaml are used orr a list of strings that are a subset of the full list. E.g. ['hand','Joystick'] for the demo Reaching-Mackenzie-2018-08-30/config.yaml to select only these two body parts. destfolder: string, optional Specifies the destination folder that was used for storing analysis data (default is the path of the video). """ from deeplabcut.pose_estimation_tensorflow.lib.inferenceutils import Assembler import pickle, re cfg = auxiliaryfunctions.read_config(config) trainFraction = cfg["TrainingFraction"][trainingsetindex] DLCscorername, _ = auxiliaryfunctions.GetScorerName( cfg, shuffle, trainFraction, modelprefix=modelprefix) for video in videos: videofolder = os.path.splitext(video)[0] if destfolder is None: outputname = "{}_full.mp4".format(videofolder + DLCscorername) full_pickle = os.path.join(videofolder + DLCscorername + "_full.pickle") else: auxiliaryfunctions.attempttomakefolder(destfolder) outputname = os.path.join( destfolder, str(Path(video).stem) + DLCscorername + "_full.mp4") full_pickle = os.path.join( destfolder, str(Path(video).stem) + DLCscorername + "_full.pickle") if not (os.path.isfile(outputname)): print("Creating labeled video for ", str(Path(video).stem)) with open(full_pickle, "rb") as file: data = pickle.load(file) header = data.pop("metadata") all_jointnames = header["all_joints_names"] if displayedbodyparts == "all": numjoints = len(all_jointnames) bpts = range(numjoints) else: # select only "displayedbodyparts" bpts = [] for bptindex, bp in enumerate(all_jointnames): if bp in displayedbodyparts: bpts.append(bptindex) numjoints = len(bpts) frame_names = list(data) frames = [int(re.findall(r"\d+", name)[0]) for name in frame_names] colorclass = plt.cm.ScalarMappable(cmap=cfg["colormap"]) C = colorclass.to_rgba(np.linspace(0, 1, numjoints)) colors = (C[:, :3] * 255).astype(np.uint8) pcutoff = cfg["pcutoff"] dotsize = cfg["dotsize"] clip = vp(fname=video, sname=outputname, codec="mp4v") ny, nx = clip.height(), clip.width() for n in trange(clip.nframes): frame = clip.load_frame() try: ind = frames.index(n) dets = Assembler._flatten_detections( data[frame_names[ind]]) for det in dets: if det.label not in bpts or det.confidence < pcutoff: continue x, y = det.pos rr, cc = disk((y, x), dotsize, shape=(ny, nx)) frame[rr, cc] = colors[bpts.index(det.label)] except ValueError: # No data stored for that particular frame print(n, "no data") pass try: clip.save_frame(frame) except: print(n, "frame writing error.") pass clip.close() else: print("Detections already plotted, ", outputname)