deeplabcut.pose_estimation_tensorflow.predict_multianimal

Functions:

Name	Description
AnalyzeMultiAnimalVideo	Helper function for analyzing a video with multiple individuals
GetPoseandCostsF	Batchwise prediction of pose
GetPoseandCostsF_from_assemblies	Batchwise prediction of pose
GetPoseandCostsS	Non batch wise pose estimation for video cap.

AnalyzeMultiAnimalVideo

AnalyzeMultiAnimalVideo(
    video,
    DLCscorer,
    trainFraction,
    cfg,
    dlc_cfg,
    sess,
    inputs,
    outputs,
    destfolder=None,
    robust_nframes=False,
    use_shelve=False,
)

Helper function for analyzing a video with multiple individuals

Source code in deeplabcut/pose_estimation_tensorflow/predict_multianimal.py

def AnalyzeMultiAnimalVideo(
    video,
    DLCscorer,
    trainFraction,
    cfg,
    dlc_cfg,
    sess,
    inputs,
    outputs,
    destfolder=None,
    robust_nframes=False,
    use_shelve=False,
):
    """Helper function for analyzing a video with multiple individuals"""

    print("Starting to analyze % ", video)
    video = Path(video)
    destfolder = video.parent if destfolder is None else Path(destfolder)
    destfolder.mkdir(exist_ok=True, parents=True)
    basename = f"{video.stem}{DLCscorer}"
    full_pickle = destfolder / f"{basename}_full.pickle"

    if full_pickle.is_file():
        print("Video already analyzed!", full_pickle)
        return None

    print("Loading ", video)
    vid = VideoWriter(str(video))
    if robust_nframes:
        nframes = vid.get_n_frames(robust=True)
        duration = vid.calc_duration(robust=True)
        fps = nframes / duration
    else:
        nframes = len(vid)
        duration = vid.calc_duration(robust=False)
        fps = vid.fps

    print(
        "Duration of video [s]: ",
        round(duration, 2),
        ", recorded with ",
        round(fps, 2),
        "fps!",
    )
    print(
        "Overall # of frames: ",
        nframes,
        " found with (before cropping) frame dimensions: ",
        vid.dimensions,
    )
    start = time.time()

    print(
        "Starting to extract posture from the video(s) with batchsize:",
        dlc_cfg["batch_size"],
    )

    shelf_path = str(full_pickle) if use_shelve else ""
    if int(dlc_cfg["batch_size"]) > 1:
        predicted_data, nframes = GetPoseandCostsF(
            cfg,
            dlc_cfg,
            sess,
            inputs,
            outputs,
            vid,
            nframes,
            int(dlc_cfg["batch_size"]),
            shelf_path,
        )
    else:
        predicted_data, nframes = GetPoseandCostsS(
            cfg,
            dlc_cfg,
            sess,
            inputs,
            outputs,
            vid,
            nframes,
            shelf_path,
        )

    stop = time.time()

    nx, ny = vid.dimensions
    if cfg["cropping"]:
        coords = [cfg["x1"], cfg["x2"], cfg["y1"], cfg["y2"]]
    else:
        coords = [0, nx, 0, ny]

    dictionary = {
        "start": start,
        "stop": stop,
        "run_duration": stop - start,
        "Scorer": DLCscorer,
        "DLC-model-config file": dlc_cfg,
        "fps": fps,
        "batch_size": dlc_cfg["batch_size"],
        "frame_dimensions": (ny, nx),
        "nframes": nframes,
        "iteration (active-learning)": cfg["iteration"],
        "training set fraction": trainFraction,
        "cropping": cfg["cropping"],
        "cropping_parameters": coords,
    }
    metadata = {"data": dictionary}
    print(f"Video Analyzed. Saving results in {destfolder}")

    if use_shelve:
        with open(destfolder / f"{basename}_meta.pickle", "wb") as f:
            pickle.dump(metadata, f, pickle.HIGHEST_PROTOCOL)
    else:
        auxfun_multianimal.SaveFullMultiAnimalData(predicted_data, metadata, str(destfolder / f"{basename}.h5"))

GetPoseandCostsF

GetPoseandCostsF(cfg, dlc_cfg, sess, inputs, outputs, cap, nframes, batchsize, shelf_path)

Batchwise prediction of pose

Source code in deeplabcut/pose_estimation_tensorflow/predict_multianimal.py

def GetPoseandCostsF(
    cfg,
    dlc_cfg,
    sess,
    inputs,
    outputs,
    cap,
    nframes,
    batchsize,
    shelf_path,
):
    """Batchwise prediction of pose"""
    strwidth = int(np.ceil(np.log10(nframes)))  # width for strings
    batch_ind = 0  # keeps track of which image within a batch should be written to
    batch_num = 0  # keeps track of which batch you are at
    if cfg["cropping"]:
        cap.set_bbox(cfg["x1"], cfg["x2"], cfg["y1"], cfg["y2"])
    nx, ny = cap.dimensions

    frames = np.empty((batchsize, ny, nx, 3), dtype="ubyte")  # this keeps all frames in a batch
    pbar = tqdm(total=nframes)
    counter = 0
    inds = []

    if shelf_path:
        db = shelve.open(
            shelf_path,
            protocol=pickle.DEFAULT_PROTOCOL,
        )
    else:
        db = dict()
    db["metadata"] = {
        "nms radius": dlc_cfg["nmsradius"],
        "minimal confidence": dlc_cfg["minconfidence"],
        "sigma": dlc_cfg.get("sigma", 1),
        "PAFgraph": dlc_cfg["partaffinityfield_graph"],
        "PAFinds": dlc_cfg.get("paf_best", np.arange(len(dlc_cfg["partaffinityfield_graph"]))),
        "all_joints": [[i] for i in range(len(dlc_cfg["all_joints"]))],
        "all_joints_names": [dlc_cfg["all_joints_names"][i] for i in range(len(dlc_cfg["all_joints"]))],
        "nframes": nframes,
    }
    while cap.video.isOpened():
        frame = cap.read_frame(crop=cfg["cropping"])
        key = "frame" + str(counter).zfill(strwidth)
        if frame is not None:
            # Avoid overwriting data already on the shelf
            if isinstance(db, shelve.Shelf) and key in db:
                continue
            frame = img_as_ubyte(frame)
            if frame.shape[-1] == 4:
                frame = rgba2rgb(frame)
            frames[batch_ind] = frame
            inds.append(counter)
            if batch_ind == batchsize - 1:
                D = predict.predict_batched_peaks_and_costs(
                    dlc_cfg,
                    frames,
                    sess,
                    inputs,
                    outputs,
                )
                for ind, data in zip(inds, D, strict=False):
                    db["frame" + str(ind).zfill(strwidth)] = data
                del D
                batch_ind = 0
                inds.clear()
                batch_num += 1
            else:
                batch_ind += 1
        elif counter >= nframes:
            if batch_ind > 0:
                D = predict.predict_batched_peaks_and_costs(
                    dlc_cfg,
                    frames,
                    sess,
                    inputs,
                    outputs,
                )
                for ind, data in zip(inds, D, strict=False):
                    db["frame" + str(ind).zfill(strwidth)] = data
                del D
            break
        counter += 1
        pbar.update(1)

    cap.close()
    pbar.close()
    try:
        db.close()
    except AttributeError:
        pass
    return db, nframes

GetPoseandCostsF_from_assemblies

GetPoseandCostsF_from_assemblies(
    cfg, dlc_cfg, sess, inputs, outputs, cap, nframes, batchsize, assemblies, feature_dict, extra_dict
)

Batchwise prediction of pose

Source code in deeplabcut/pose_estimation_tensorflow/predict_multianimal.py

def GetPoseandCostsF_from_assemblies(
    cfg,
    dlc_cfg,
    sess,
    inputs,
    outputs,
    cap,
    nframes,
    batchsize,
    assemblies,
    feature_dict,
    extra_dict,
):
    """Batchwise prediction of pose"""
    strwidth = int(np.ceil(np.log10(nframes)))  # width for strings
    batch_ind = 0  # keeps track of which image within a batch should be written to
    batch_num = 0  # keeps track of which batch you are at
    if cfg["cropping"]:
        cap.set_bbox(cfg["x1"], cfg["x2"], cfg["y1"], cfg["y2"])
    nx, ny = cap.dimensions

    frames = np.empty((batchsize, ny, nx, 3), dtype="ubyte")  # this keeps all frames in a batch
    pbar = tqdm(total=nframes)
    counter = 0
    inds = []

    predicted_data = {}

    while cap.video.isOpened():
        frame = cap.read_frame(crop=cfg["cropping"])
        key = "frame" + str(counter).zfill(strwidth)
        if frame is not None:
            # Avoid overwriting data already on the shelf
            if key in feature_dict:
                continue

            frame = img_as_ubyte(frame)
            if frame.shape[-1] == 4:
                frame = rgba2rgb(frame)
            frames[batch_ind] = frame
            inds.append(counter)

            if batch_ind == batchsize - 1:
                preds = predict.predict_batched_peaks_and_costs(
                    dlc_cfg, frames, sess, inputs, outputs, extra_dict=extra_dict
                )
                if not preds:
                    continue

                D, features = preds
                for i, (ind, data) in enumerate(zip(inds, D, strict=False)):
                    predicted_data["frame" + str(ind).zfill(strwidth)] = data
                    raw_coords = assemblies.get(ind)
                    if raw_coords is None:
                        continue
                    fname = "frame" + str(ind).zfill(strwidth)
                    feature_dict[fname] = _get_features_dict(
                        raw_coords,
                        features[i],
                        dlc_cfg["stride"],
                    )

                batch_ind = 0
                inds.clear()
                batch_num += 1
            else:
                batch_ind += 1
        elif counter >= nframes:
            if batch_ind > 0:
                preds = predict.predict_batched_peaks_and_costs(
                    dlc_cfg, frames, sess, inputs, outputs, extra_dict=extra_dict
                )
                if not preds:
                    continue

                D, features = preds
                for i, (ind, data) in enumerate(zip(inds, D, strict=False)):
                    predicted_data["frame" + str(ind).zfill(strwidth)] = data
                    raw_coords = assemblies.get(ind)
                    if raw_coords is None:
                        continue
                    fname = "frame" + str(ind).zfill(strwidth)
                    feature_dict[fname] = _get_features_dict(
                        raw_coords,
                        features[i],
                        dlc_cfg["stride"],
                    )

            break
        counter += 1
        pbar.update(1)

    cap.close()
    pbar.close()
    feature_dict.close()
    predicted_data["metadata"] = {
        "nms radius": dlc_cfg["nmsradius"],
        "minimal confidence": dlc_cfg["minconfidence"],
        "sigma": dlc_cfg.get("sigma", 1),
        "PAFgraph": dlc_cfg["partaffinityfield_graph"],
        "PAFinds": dlc_cfg.get("paf_best", np.arange(len(dlc_cfg["partaffinityfield_graph"]))),
        "all_joints": [[i] for i in range(len(dlc_cfg["all_joints"]))],
        "all_joints_names": [dlc_cfg["all_joints_names"][i] for i in range(len(dlc_cfg["all_joints"]))],
        "nframes": nframes,
    }
    return predicted_data, nframes

GetPoseandCostsS

GetPoseandCostsS(cfg, dlc_cfg, sess, inputs, outputs, cap, nframes, shelf_path)

Non batch wise pose estimation for video cap.

Source code in deeplabcut/pose_estimation_tensorflow/predict_multianimal.py

def GetPoseandCostsS(cfg, dlc_cfg, sess, inputs, outputs, cap, nframes, shelf_path):
    """Non batch wise pose estimation for video cap."""
    strwidth = int(np.ceil(np.log10(nframes)))  # width for strings
    if cfg["cropping"]:
        cap.set_bbox(cfg["x1"], cfg["x2"], cfg["y1"], cfg["y2"])

    if shelf_path:
        db = shelve.open(
            shelf_path,
            protocol=pickle.DEFAULT_PROTOCOL,
        )
    else:
        db = dict()
    db["metadata"] = {
        "nms radius": dlc_cfg["nmsradius"],
        "minimal confidence": dlc_cfg["minconfidence"],
        "sigma": dlc_cfg.get("sigma", 1),
        "PAFgraph": dlc_cfg["partaffinityfield_graph"],
        "PAFinds": dlc_cfg.get("paf_best", np.arange(len(dlc_cfg["partaffinityfield_graph"]))),
        "all_joints": [[i] for i in range(len(dlc_cfg["all_joints"]))],
        "all_joints_names": [dlc_cfg["all_joints_names"][i] for i in range(len(dlc_cfg["all_joints"]))],
        "nframes": nframes,
    }
    pbar = tqdm(total=nframes)
    counter = 0
    while cap.video.isOpened():
        frame = cap.read_frame(crop=cfg["cropping"])
        key = "frame" + str(counter).zfill(strwidth)
        if frame is not None:
            # Avoid overwriting data already on the shelf
            if isinstance(db, shelve.Shelf) and key in db:
                continue
            frame = img_as_ubyte(frame)
            if frame.shape[-1] == 4:
                frame = rgba2rgb(frame)
            dets = predict.predict_batched_peaks_and_costs(
                dlc_cfg,
                np.expand_dims(frame, axis=0),
                sess,
                inputs,
                outputs,
            )
            if not dets:
                continue
            db[key] = dets[0]
            del dets
        elif counter >= nframes:
            break
        counter += 1
        pbar.update(1)

    pbar.close()
    try:
        db.close()
    except AttributeError:
        pass
    return db, nframes