Example #1
0
def ExtractFramesbasedonPreselection(Index,
                                     extractionalgorithm,
                                     Dataframe,
                                     dataname,
                                     scorer,
                                     video,
                                     cfg,
                                     config,
                                     opencv=True,
                                     cluster_resizewidth=30,
                                     cluster_color=False,
                                     savelabeled=True):
    from deeplabcut.create_project import add
    start = cfg['start']
    stop = cfg['stop']
    numframes2extract = cfg['numframes2pick']
    bodyparts = cfg['bodyparts']

    videofolder = str(Path(video).parents[0])
    vname = str(Path(video).stem)
    tmpfolder = os.path.join(cfg['project_path'], 'labeled-data', vname)
    if os.path.isdir(tmpfolder):
        print("Frames from video", vname,
              " already extracted (more will be added)!")
    else:
        auxiliaryfunctions.attempttomakefolder(tmpfolder)

    nframes = np.size(Dataframe.index)
    print("Loading video...")
    if opencv:
        import cv2
        cap = cv2.VideoCapture(video)
        fps = cap.get(5)
        duration = nframes * 1. / fps
        size = (int(cap.get(4)), int(cap.get(3)))
    else:
        from moviepy.editor import VideoFileClip
        clip = VideoFileClip(video)
        fps = clip.fps
        duration = clip.duration
        size = clip.size

    if cfg['cropping']:  # one might want to adjust
        coords = (cfg['x1'], cfg['x2'], cfg['y1'], cfg['y2'])
    else:
        coords = None

    print("Duration of video [s]: ", duration, ", recorded @ ", fps, "fps!")
    print(
        "Overall # of frames: ",
        nframes,
        "with (cropped) frame dimensions: ",
    )
    if extractionalgorithm == 'uniform':
        if opencv:
            frames2pick = frameselectiontools.UniformFramescv2(
                cap, numframes2extract, start, stop, Index)
        else:
            frames2pick = frameselectiontools.UniformFrames(
                clip, numframes2extract, start, stop, Index)
    elif extractionalgorithm == 'kmeans':
        if opencv:
            frames2pick = frameselectiontools.KmeansbasedFrameselectioncv2(
                cap,
                numframes2extract,
                start,
                stop,
                cfg['cropping'],
                coords,
                Index,
                resizewidth=cluster_resizewidth,
                color=cluster_color)
        else:
            if cfg['cropping']:
                clip = clip.crop(y1=cfg['y1'],
                                 y2=cfg['x2'],
                                 x1=cfg['x1'],
                                 x2=cfg['x2'])
            frames2pick = frameselectiontools.KmeansbasedFrameselection(
                clip,
                numframes2extract,
                start,
                stop,
                Index,
                resizewidth=cluster_resizewidth,
                color=cluster_color)

    else:
        print(
            "Please implement this method yourself! Currently the options are 'kmeans', 'jump', 'uniform'."
        )
        frames2pick = []

    # Extract frames + frames with plotted labels and store them in folder (with name derived from video name) nder labeled-data
    print("Let's select frames indices:", frames2pick)
    colors = visualization.get_cmap(len(bodyparts), cfg['colormap'])
    strwidth = int(np.ceil(np.log10(nframes)))  #width for strings
    for index in frames2pick:  ##tqdm(range(0,nframes,10)):
        if opencv:
            PlottingSingleFramecv2(cap, cv2, cfg['cropping'], coords,
                                   Dataframe, bodyparts, tmpfolder, index,
                                   scorer, cfg['dotsize'], cfg['pcutoff'],
                                   cfg['alphavalue'], colors, strwidth,
                                   savelabeled)
        else:
            PlottingSingleFrame(clip, Dataframe, bodyparts, tmpfolder, index,
                                scorer, cfg['dotsize'], cfg['pcutoff'],
                                cfg['alphavalue'], colors, strwidth,
                                savelabeled)
        plt.close("all")

    #close videos
    if opencv:
        cap.release()
    else:
        clip.close()
        del clip

    # Extract annotations based on DeepLabCut and store in the folder (with name derived from video name) under labeled-data
    if len(frames2pick) > 0:
        #Dataframe = pd.read_hdf(os.path.join(videofolder,dataname+'.h5'))
        DF = Dataframe.ix[frames2pick]
        DF.index = [
            os.path.join('labeled-data', vname,
                         "img" + str(index).zfill(strwidth) + ".png")
            for index in DF.index
        ]  #exchange index number by file names.

        machinefile = os.path.join(
            tmpfolder, 'machinelabels-iter' + str(cfg['iteration']) + '.h5')
        if Path(machinefile).is_file():
            Data = pd.read_hdf(machinefile, 'df_with_missing')
            DataCombined = pd.concat([Data, DF])
            #drop duplicate labels:
            DataCombined = DataCombined[~DataCombined.index.duplicated(
                keep='first')]

            DataCombined.to_hdf(machinefile, key='df_with_missing', mode='w')
            DataCombined.to_csv(
                os.path.join(tmpfolder, "machinelabels.csv")
            )  #this is always the most current one (as reading is from h5)
        else:
            DF.to_hdf(machinefile, key='df_with_missing', mode='w')
            DF.to_csv(os.path.join(tmpfolder, "machinelabels.csv"))
        try:
            if cfg['cropping']:
                add.add_new_videos(
                    config, [video],
                    coords=[coords])  # make sure you pass coords as a list
            else:
                add.add_new_videos(config, [video], coords=None)
        except:  #can we make a catch here? - in fact we should drop indices from DataCombined if they are in CollectedData.. [ideal behavior; currently this is pretty unlikely]
            print(
                "AUTOMATIC ADDING OF VIDEO TO CONFIG FILE FAILED! You need to do this manually for including it in the config.yaml file!"
            )
            print("Videopath:", video, "Coordinates for cropping:", coords)
            pass

        print(
            "The outlier frames are extracted. They are stored in the subdirectory labeled-data\%s."
            % vname)
        print(
            "Once you extracted frames for all videos, use 'refine_labels' to manually correct the labels."
        )
    else:
        print("No frames were extracted.")
def ExtractFramesbasedonPreselection(
    Index,
    extractionalgorithm,
    data,
    video,
    cfg,
    config,
    opencv=True,
    cluster_resizewidth=30,
    cluster_color=False,
    savelabeled=True,
    with_annotations=True,
):
    from deeplabcut.create_project import add

    start = cfg["start"]
    stop = cfg["stop"]
    numframes2extract = cfg["numframes2pick"]
    bodyparts = auxiliaryfunctions.IntersectionofBodyPartsandOnesGivenbyUser(
        cfg, "all")

    videofolder = str(Path(video).parents[0])
    vname = str(Path(video).stem)
    tmpfolder = os.path.join(cfg["project_path"], "labeled-data", vname)
    if os.path.isdir(tmpfolder):
        print("Frames from video", vname,
              " already extracted (more will be added)!")
    else:
        auxiliaryfunctions.attempttomakefolder(tmpfolder, recursive=True)

    nframes = len(data)
    print("Loading video...")
    if opencv:
        vid = VideoWriter(video)
        fps = vid.fps
        duration = vid.calc_duration()
    else:
        from moviepy.editor import VideoFileClip

        clip = VideoFileClip(video)
        fps = clip.fps
        duration = clip.duration

    if cfg["cropping"]:  # one might want to adjust
        coords = (cfg["x1"], cfg["x2"], cfg["y1"], cfg["y2"])
    else:
        coords = None

    print("Duration of video [s]: ", duration, ", recorded @ ", fps, "fps!")
    print("Overall # of frames: ", nframes,
          "with (cropped) frame dimensions: ")
    if extractionalgorithm == "uniform":
        if opencv:
            frames2pick = frameselectiontools.UniformFramescv2(
                vid, numframes2extract, start, stop, Index)
        else:
            frames2pick = frameselectiontools.UniformFrames(
                clip, numframes2extract, start, stop, Index)
    elif extractionalgorithm == "kmeans":
        if opencv:
            frames2pick = frameselectiontools.KmeansbasedFrameselectioncv2(
                vid,
                numframes2extract,
                start,
                stop,
                cfg["cropping"],
                coords,
                Index,
                resizewidth=cluster_resizewidth,
                color=cluster_color,
            )
        else:
            if cfg["cropping"]:
                clip = clip.crop(y1=cfg["y1"],
                                 y2=cfg["x2"],
                                 x1=cfg["x1"],
                                 x2=cfg["x2"])
            frames2pick = frameselectiontools.KmeansbasedFrameselection(
                clip,
                numframes2extract,
                start,
                stop,
                Index,
                resizewidth=cluster_resizewidth,
                color=cluster_color,
            )

    else:
        print(
            "Please implement this method yourself! Currently the options are 'kmeans', 'jump', 'uniform'."
        )
        frames2pick = []

    # Extract frames + frames with plotted labels and store them in folder (with name derived from video name) nder labeled-data
    print("Let's select frames indices:", frames2pick)
    colors = visualization.get_cmap(len(bodyparts), cfg["colormap"])
    strwidth = int(np.ceil(np.log10(nframes)))  # width for strings
    for index in frames2pick:  ##tqdm(range(0,nframes,10)):
        if opencv:
            PlottingSingleFramecv2(
                vid,
                cfg["cropping"],
                coords,
                data,
                bodyparts,
                tmpfolder,
                index,
                cfg["dotsize"],
                cfg["pcutoff"],
                cfg["alphavalue"],
                colors,
                strwidth,
                savelabeled,
            )
        else:
            PlottingSingleFrame(
                clip,
                data,
                bodyparts,
                tmpfolder,
                index,
                cfg["dotsize"],
                cfg["pcutoff"],
                cfg["alphavalue"],
                colors,
                strwidth,
                savelabeled,
            )
        plt.close("all")

    # close videos
    if opencv:
        vid.close()
    else:
        clip.close()
        del clip

    # Extract annotations based on DeepLabCut and store in the folder (with name derived from video name) under labeled-data
    if len(frames2pick) > 0:
        try:
            if cfg["cropping"]:
                add.add_new_videos(
                    config, [video],
                    coords=[coords])  # make sure you pass coords as a list
            else:
                add.add_new_videos(config, [video], coords=None)
        except:  # can we make a catch here? - in fact we should drop indices from DataCombined if they are in CollectedData.. [ideal behavior; currently this is pretty unlikely]
            print(
                "AUTOMATIC ADDING OF VIDEO TO CONFIG FILE FAILED! You need to do this manually for including it in the config.yaml file!"
            )
            print("Videopath:", video, "Coordinates for cropping:", coords)
            pass

        if with_annotations:
            machinefile = os.path.join(
                tmpfolder,
                "machinelabels-iter" + str(cfg["iteration"]) + ".h5")
            if isinstance(data, pd.DataFrame):
                df = data.loc[frames2pick]
                df.index = [
                    os.path.join(
                        "labeled-data",
                        vname,
                        "img" + str(index).zfill(strwidth) + ".png",
                    ) for index in df.index
                ]  # exchange index number by file names.
            elif isinstance(data, dict):
                idx = [
                    os.path.join(
                        "labeled-data",
                        vname,
                        "img" + str(index).zfill(strwidth) + ".png",
                    ) for index in frames2pick
                ]
                filename = os.path.join(str(tmpfolder),
                                        f"CollectedData_{cfg['scorer']}.h5")
                try:
                    df_temp = pd.read_hdf(filename, "df_with_missing")
                    columns = df_temp.columns
                except FileNotFoundError:
                    columns = pd.MultiIndex.from_product(
                        [
                            [cfg["scorer"]],
                            cfg["individuals"],
                            cfg["multianimalbodyparts"],
                            ["x", "y"],
                        ],
                        names=["scorer", "individuals", "bodyparts", "coords"],
                    )
                    if cfg["uniquebodyparts"]:
                        columns2 = pd.MultiIndex.from_product(
                            [
                                [cfg["scorer"]],
                                ["single"],
                                cfg["uniquebodyparts"],
                                ["x", "y"],
                            ],
                            names=[
                                "scorer", "individuals", "bodyparts", "coords"
                            ],
                        )
                        df_temp = pd.concat((
                            pd.DataFrame(columns=columns),
                            pd.DataFrame(columns=columns2),
                        ))
                        columns = df_temp.columns
                array = np.full((len(frames2pick), len(columns)), np.nan)
                for i, index in enumerate(frames2pick):
                    data_temp = data.get(index)
                    if data_temp is not None:
                        vals = np.concatenate(data_temp)[:, :2].flatten()
                        array[i, :len(vals)] = vals
                df = pd.DataFrame(array, index=idx, columns=columns)
            else:
                return
            if Path(machinefile).is_file():
                Data = pd.read_hdf(machinefile, "df_with_missing")
                DataCombined = pd.concat([Data, df])
                # drop duplicate labels:
                DataCombined = DataCombined[~DataCombined.index.duplicated(
                    keep="first")]

                DataCombined.to_hdf(machinefile,
                                    key="df_with_missing",
                                    mode="w")
                DataCombined.to_csv(
                    os.path.join(tmpfolder, "machinelabels.csv")
                )  # this is always the most current one (as reading is from h5)
            else:
                df.to_hdf(machinefile, key="df_with_missing", mode="w")
                df.to_csv(os.path.join(tmpfolder, "machinelabels.csv"))

        print(
            "The outlier frames are extracted. They are stored in the subdirectory labeled-data\%s."
            % vname)
        print(
            "Once you extracted frames for all videos, use 'refine_labels' to manually correct the labels."
        )
    else:
        print("No frames were extracted.")
def extract_frames(
    config,
    mode="automatic",
    algo="kmeans",
    crop=False,
    userfeedback=True,
    cluster_step=1,
    cluster_resizewidth=30,
    cluster_color=False,
    opencv=True,
    slider_width=25,
    config3d=None,
    extracted_cam=0,
    videos_list=None,
):
    """
    Extracts frames from the videos in the config.yaml file. Only the videos in the config.yaml will be used to select the frames.\n
    Use the function ``add_new_video`` at any stage of the project to add new videos to the config file and extract their frames.

    The provided function either selects frames from the videos in a randomly and temporally uniformly distributed way (uniform), \n
    by clustering based on visual appearance (k-means), or by manual selection.

    Three important parameters for automatic extraction: numframes2pick, start and stop are set in the config file.

    After frames have been extracted from all videos from one camera, matched frames from other cameras can be extracted using mode = ``match``.
    This is necessary if you plan to use epipolar lines to improve labeling across multiple camera angles. It will overwrite previously extracted
    images from the second camera angle if necessary.

    Please refer to the user guide for more details on methods and parameters https://www.nature.com/articles/s41596-019-0176-0
    or the preprint: https://www.biorxiv.org/content/biorxiv/early/2018/11/24/476531.full.pdf

    Parameters
    ----------
    config : string
        Full path of the config.yaml file as a string.

    mode : string
        String containing the mode of extraction. It must be either ``automatic`` or ``manual`` to extract the initial set of frames. It can also be ``match`` to match frames between
        the cameras in preparation for the use of epipolar lines during labeling; namely, extract from camera_1 first, then run this to extract the matched frames in camera_2.
        WARNING: if you use match, and you previously extracted and labeled frames from the second camera, this will overwrite your data. This will require you deleting the
        collectdata.h5/.csv files before labeling.... Use with caution!

    algo : string
        String specifying the algorithm to use for selecting the frames. Currently, deeplabcut supports either ``kmeans`` or ``uniform`` based selection. This flag is
        only required for ``automatic`` mode and the default is ``uniform``. For uniform, frames are picked in temporally uniform way, kmeans performs clustering on downsampled frames (see user guide for details).
        Note: color information is discarded for kmeans, thus e.g. for camouflaged octopus clustering one might want to change this.

    crop : bool, optional
        If True, video frames are cropped according to the corresponding coordinates stored in the config.yaml.
        Alternatively, if cropping coordinates are not known yet, crop='GUI' triggers a user interface
        where the cropping area can be manually drawn and saved.

    userfeedback: bool, optional
        If this is set to false during automatic mode then frames for all videos are extracted. The user can set this to true, which will result in a dialog,
        where the user is asked for each video if (additional/any) frames from this video should be extracted. Use this, e.g. if you have already labeled
        some folders and want to extract data for new videos.

    cluster_resizewidth: number, default: 30
        For k-means one can change the width to which the images are downsampled (aspect ratio is fixed).

    cluster_step: number, default: 1
        By default each frame is used for clustering, but for long videos one could only use every nth frame (set by: cluster_step). This saves memory before clustering can start, however,
        reading the individual frames takes longer due to the skipping.

    cluster_color: bool, default: False
        If false then each downsampled image is treated as a grayscale vector (discarding color information). If true, then the color channels are considered. This increases
        the computational complexity.

    opencv: bool, default: True
        Uses openCV for loading & extractiong (otherwise moviepy (legacy))

    slider_width: number, default: 25
        Width of the video frames slider, in percent of window

    config3d: string, optional
        Path to the config.yaml file in the 3D project. This will be used to match frames extracted from all cameras present in the field 'camera_names' to the
        frames extracted from the camera given by the parameter 'extracted_cam'

    extracted_cam: number, default: 0
        The index of the camera that already has extracted frames. This will match frame numbers to extract for all other cameras.
        This parameter is necessary if you wish to use epipolar lines in the labeling toolbox. Only use if mode = 'match' and config3d is provided.

    videos_list: list, default: None
            A list of the string containing full paths to videos to extract frames for. If this is left as None all videos specified in the config file will have frames extracted. 
            Otherwise one can select a subset by passing those paths. 

    Examples
    --------
    for selecting frames automatically with 'kmeans' and want to crop the frames.
    >>> deeplabcut.extract_frames('/analysis/project/reaching-task/config.yaml','automatic','kmeans',True)
    --------
    for selecting frames automatically with 'kmeans' and defining the cropping area at runtime.
    >>> deeplabcut.extract_frames('/analysis/project/reaching-task/config.yaml','automatic','kmeans','GUI')
    --------
    for selecting frames automatically with 'kmeans' and considering the color information.
    >>> deeplabcut.extract_frames('/analysis/project/reaching-task/config.yaml','automatic','kmeans',cluster_color=True)
    --------
    for selecting frames automatically with 'uniform' and want to crop the frames.
    >>> deeplabcut.extract_frames('/analysis/project/reaching-task/config.yaml','automatic',crop=True)
    --------
    for selecting frames manually,
    >>> deeplabcut.extract_frames('/analysis/project/reaching-task/config.yaml','manual')
    --------
    for selecting frames manually, with a 60% wide frames slider
    >>> deeplabcut.extract_frames('/analysis/project/reaching-task/config.yaml','manual', slider_width=60)
    --------
    for extracting frames from a second camera that match the frames extracted from the first
    >>> deeplabcut.extract_frames('/analysis/project/reaching-task/config.yaml', mode='match', extracted_cam=0)

    While selecting the frames manually, you do not need to specify the ``crop`` parameter in the command. Rather, you will get a prompt in the graphic user interface to choose
    if you need to crop or not.
    --------

    """
    import os
    import sys
    import re
    import glob
    import numpy as np
    from pathlib import Path
    from skimage import io
    from skimage.util import img_as_ubyte
    from deeplabcut.utils import frameselectiontools
    from deeplabcut.utils import auxiliaryfunctions

    if mode == "manual":
        wd = Path(config).resolve().parents[0]
        os.chdir(str(wd))
        from deeplabcut.gui import frame_extraction_toolbox

        frame_extraction_toolbox.show(config, slider_width)

    elif mode == "automatic":
        config_file = Path(config).resolve()
        cfg = auxiliaryfunctions.read_config(config_file)
        print("Config file read successfully.")

        numframes2pick = cfg["numframes2pick"]
        start = cfg["start"]
        stop = cfg["stop"]

        # Check for variable correctness
        if start > 1 or stop > 1 or start < 0 or stop < 0 or start >= stop:
            raise Exception(
                "Erroneous start or stop values. Please correct it in the config file."
            )
        if numframes2pick < 1 and not int(numframes2pick):
            raise Exception(
                "Perhaps consider extracting more, or a natural number of frames."
            )
        if videos_list is None:
            videos = cfg.get("video_sets_original") or cfg["video_sets"]
        else:  #filter video_list by the ones in the config file
            videos = [v for v in cfg["video_sets"] if v in videos_list]

        if opencv:
            from deeplabcut.utils.auxfun_videos import VideoReader
        else:
            from moviepy.editor import VideoFileClip

        has_failed = []
        for video in videos:
            if userfeedback:
                print(
                    "Do you want to extract (perhaps additional) frames for video:",
                    video,
                    "?",
                )
                askuser = input("yes/no")
            else:
                askuser = "******"

            if (askuser == "y" or askuser == "yes" or askuser == "Ja"
                    or askuser == "ha" or askuser == "oui"
                    or askuser == "ouais"):  # multilanguage support :)

                if opencv:
                    cap = VideoReader(video)
                    nframes = len(cap)
                else:
                    # Moviepy:
                    clip = VideoFileClip(video)
                    fps = clip.fps
                    nframes = int(np.ceil(clip.duration * 1.0 / fps))
                if not nframes:
                    print("Video could not be opened. Skipping...")
                    continue

                indexlength = int(np.ceil(np.log10(nframes)))

                fname = Path(video)
                output_path = Path(
                    config).parents[0] / "labeled-data" / fname.stem

                if output_path.exists():
                    if len(os.listdir(output_path)):
                        if userfeedback:
                            askuser = input(
                                "The directory already contains some frames. Do you want to add to it?(yes/no): "
                            )
                        if not (askuser == "y" or askuser == "yes"
                                or askuser == "Y" or askuser == "Yes"):
                            sys.exit("Delete the frames and try again later!")

                if crop == "GUI":
                    cfg = select_cropping_area(config, [video])
                try:
                    coords = cfg["video_sets"][video]["crop"].split(",")
                except KeyError:
                    coords = cfg["video_sets_original"][video]["crop"].split(
                        ",")

                if crop and not opencv:
                    clip = clip.crop(
                        y1=int(coords[2]),
                        y2=int(coords[3]),
                        x1=int(coords[0]),
                        x2=int(coords[1]),
                    )
                elif not crop:
                    coords = None

                print("Extracting frames based on %s ..." % algo)
                if algo == "uniform":
                    if opencv:
                        frames2pick = frameselectiontools.UniformFramescv2(
                            cap, numframes2pick, start, stop)
                    else:
                        frames2pick = frameselectiontools.UniformFrames(
                            clip, numframes2pick, start, stop)
                elif algo == "kmeans":
                    if opencv:
                        frames2pick = frameselectiontools.KmeansbasedFrameselectioncv2(
                            cap,
                            numframes2pick,
                            start,
                            stop,
                            crop,
                            coords,
                            step=cluster_step,
                            resizewidth=cluster_resizewidth,
                            color=cluster_color,
                        )
                    else:
                        frames2pick = frameselectiontools.KmeansbasedFrameselection(
                            clip,
                            numframes2pick,
                            start,
                            stop,
                            step=cluster_step,
                            resizewidth=cluster_resizewidth,
                            color=cluster_color,
                        )
                else:
                    print(
                        "Please implement this method yourself and send us a pull request! Otherwise, choose 'uniform' or 'kmeans'."
                    )
                    frames2pick = []

                if not len(frames2pick):
                    print("Frame selection failed...")
                    return

                output_path = (Path(config).parents[0] / "labeled-data" /
                               Path(video).stem)
                is_valid = []
                if opencv:
                    for index in frames2pick:
                        cap.set_to_frame(index)  # extract a particular frame
                        frame = cap.read_frame()
                        if frame is not None:
                            image = img_as_ubyte(frame)
                            img_name = (str(output_path) + "/img" +
                                        str(index).zfill(indexlength) + ".png")
                            if crop:
                                io.imsave(
                                    img_name,
                                    image[int(coords[2]):int(coords[3]),
                                          int(coords[0]):int(coords[1]), :, ],
                                )  # y1 = int(coords[2]),y2 = int(coords[3]),x1 = int(coords[0]), x2 = int(coords[1]
                            else:
                                io.imsave(img_name, image)
                            is_valid.append(True)
                        else:
                            print("Frame", index, " not found!")
                            is_valid.append(False)
                    cap.close()
                else:
                    for index in frames2pick:
                        try:
                            image = img_as_ubyte(
                                clip.get_frame(index * 1.0 / clip.fps))
                            img_name = (str(output_path) + "/img" +
                                        str(index).zfill(indexlength) + ".png")
                            io.imsave(img_name, image)
                            if np.var(image) == 0:  # constant image
                                print(
                                    "Seems like black/constant images are extracted from your video. Perhaps consider using opencv under the hood, by setting: opencv=True"
                                )
                            is_valid.append(True)
                        except FileNotFoundError:
                            print("Frame # ", index, " does not exist.")
                            is_valid.append(False)
                    clip.close()
                    del clip

                if not any(is_valid):
                    has_failed.append(True)
                else:
                    has_failed.append(False)

            else:  # NO!
                has_failed.append(False)

        if all(has_failed):
            print("Frame extraction failed. Video files must be corrupted.")
            return
        elif any(has_failed):
            print("Although most frames were extracted, some were invalid.")
        else:
            print(
                "Frames were successfully extracted, for the videos listed in the config.yaml file."
            )
        print(
            "\nYou can now label the frames using the function 'label_frames' "
            "(Note, you should label frames extracted from diverse videos (and many videos; we do not recommend training on single videos!))."
        )

    elif mode == "match":
        import cv2

        config_file = Path(config).resolve()
        cfg = auxiliaryfunctions.read_config(config_file)
        print("Config file read successfully.")
        videos = sorted(cfg["video_sets"].keys())
        project_path = Path(config).parents[0]
        labels_path = os.path.join(project_path, "labeled-data/")
        video_dir = os.path.join(project_path, "videos/")
        try:
            cfg_3d = auxiliaryfunctions.read_config(config3d)
        except:
            raise Exception(
                "You must create a 3D project and edit the 3D config file before extracting matched frames. \n"
            )
        cams = cfg_3d["camera_names"]
        extCam_name = cams[extracted_cam]
        del cams[extracted_cam]
        label_dirs = sorted(
            glob.glob(os.path.join(labels_path, "*" + extCam_name + "*")))

        # select crop method
        crop_list = []
        for video in videos:
            if extCam_name not in video:
                if crop == "GUI":
                    cfg = select_cropping_area(config, [video])
                    print("in gui code")
                coords = cfg["video_sets"][video]["crop"].split(",")

                if crop and not opencv:
                    clip = clip.crop(
                        y1=int(coords[2]),
                        y2=int(coords[3]),
                        x1=int(coords[0]),
                        x2=int(coords[1]),
                    )
                elif not crop:
                    coords = None
                crop_list.append(coords)
        print(crop_list)

        for coords, dirPath in zip(crop_list, label_dirs):
            extracted_images = glob.glob(os.path.join(dirPath, "*png"))

            imgPattern = re.compile("[0-9]{1,10}")
            for cam in cams:
                output_path = re.sub(extCam_name, cam, dirPath)

                for fname in os.listdir(output_path):
                    if fname.endswith(".png"):
                        os.remove(os.path.join(output_path, fname))

                vid = os.path.join(video_dir,
                                   os.path.basename(output_path)) + ".avi"
                cap = cv2.VideoCapture(vid)
                print("\n extracting matched frames from " +
                      os.path.basename(output_path) + ".avi")
                for img in extracted_images:
                    imgNum = re.findall(imgPattern, os.path.basename(img))[0]
                    cap.set(1, int(imgNum))
                    ret, frame = cap.read()
                    if ret:
                        image = img_as_ubyte(
                            cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
                        img_name = str(output_path) + "/img" + imgNum + ".png"
                        if crop:
                            io.imsave(
                                img_name,
                                image[int(coords[2]):int(coords[3]),
                                      int(coords[0]):int(coords[1]), :, ],
                            )
                        else:
                            io.imsave(img_name, image)
        print(
            "\n Done extracting matched frames. You can now begin labeling frames using the function label_frames\n"
        )

    else:
        print(
            "Invalid MODE. Choose either 'manual', 'automatic' or 'match'. Check ``help(deeplabcut.extract_frames)`` on python and ``deeplabcut.extract_frames?`` \
              for ipython/jupyter notebook for more details.")
def extract_frames(config,
                   mode='automatic',
                   algo='kmeans',
                   crop=False,
                   userfeedback=True,
                   cluster_step=1,
                   cluster_resizewidth=30,
                   cluster_color=False,
                   opencv=True,
                   flymovie=False,
                   slider_width=25):
    """
    Extracts frames from the videos in the config.yaml file. Only the videos in the config.yaml will be used to select the frames.\n
    Use the function ``add_new_video`` at any stage of the project to add new videos to the config file and extract their frames.
    
    The provided function either selects frames from the videos in a randomly and temporally uniformly distributed way (uniform), \n 
    by clustering based on visual appearance (k-means), or by manual selection. 
    
    Three important parameters for automatic extraction: numframes2pick, start and stop are set in the config file. 
    
    Please refer to the user guide for more details on methods and parameters https://www.biorxiv.org/content/biorxiv/early/2018/11/24/476531.full.pdf
    
    Parameters
    ----------
    config : string
        Full path of the config.yaml file as a string.
        
    mode : string
        String containing the mode of extraction. It must be either ``automatic`` or ``manual``.
        
    algo : string 
        String specifying the algorithm to use for selecting the frames. Currently, deeplabcut supports either ``kmeans`` or ``uniform`` based selection. This flag is
        only required for ``automatic`` mode and the default is ``uniform``. For uniform, frames are picked in temporally uniform way, kmeans performs clustering on downsampled frames (see user guide for details).
        Note: color information is discarded for kmeans, thus e.g. for camouflaged octopus clustering one might want to change this. 
        
    crop : bool, optional
        If this is set to True, a user interface pops up with a frame to select the cropping parameters. Use the left click to draw a cropping area and hit the button set cropping parameters to save the cropping parameters for a video.
        The default is ``False``; if provided it must be either ``True`` or ``False``.
            
    userfeedback: bool, optional
        If this is set to false during automatic mode then frames for all videos are extracted. The user can set this to true, which will result in a dialog,
        where the user is asked for each video if (additional/any) frames from this video should be extracted. Use this, e.g. if you have already labeled
        some folders and want to extract data for new videos. 
    
    cluster_resizewidth: number, default: 30
        For k-means one can change the width to which the images are downsampled (aspect ratio is fixed).
    
    cluster_step: number, default: 1
        By default each frame is used for clustering, but for long videos one could only use every nth frame (set by: cluster_step). This saves memory before clustering can start, however, 
        reading the individual frames takes longer due to the skipping.
    
    cluster_color: bool, default: False
        If false then each downsampled image is treated as a grayscale vector (discarding color information). If true, then the color channels are considered. This increases 
        the computational complexity. 
    
    opencv: bool, default: True
        Uses openCV for loading & extractiong (otherwise moviepy (legacy))
        
    slider_width: number, default: 25
        Width of the video frames slider, in percent of window
        
    Examples
    --------
    for selecting frames automatically with 'kmeans' and want to crop the frames.
    >>> deeplabcut.extract_frames('/analysis/project/reaching-task/config.yaml','automatic','kmeans',True)
    --------
    for selecting frames automatically with 'kmeans' and considering the color information.
    >>> deeplabcut.extract_frames('/analysis/project/reaching-task/config.yaml','automatic','kmeans',cluster_color=True)
    --------
    for selecting frames automatically with 'uniform' and want to crop the frames.
    >>> deeplabcut.extract_frames('/analysis/project/reaching-task/config.yaml','automatic',crop=True)
    --------
    for selecting frames manually,
    >>> deeplabcut.extract_frames('/analysis/project/reaching-task/config.yaml','manual')
    --------
    for selecting frames manually, with a 60% wide frames slider
    >>> deeplabcut.extract_frames('/analysis/project/reaching-task/config.yaml','manual', slider_width=60)
    
    While selecting the frames manually, you do not need to specify the ``crop`` parameter in the command. Rather, you will get a prompt in the graphic user interface to choose 
    if you need to crop or not.
    --------
    
    """
    import os
    import sys
    import numpy as np
    from pathlib import Path
    from skimage import io
    import skimage
    from skimage.util import img_as_ubyte
    import matplotlib.pyplot as plt
    import matplotlib.patches as patches
    from deeplabcut.utils import frameselectiontools
    from deeplabcut.utils import auxiliaryfunctions
    from matplotlib.widgets import RectangleSelector

    if mode == "manual":
        wd = Path(config).resolve().parents[0]
        os.chdir(str(wd))
        from deeplabcut.generate_training_dataset import frame_extraction_toolbox
        from deeplabcut.utils import select_crop_parameters
        frame_extraction_toolbox.show(config, slider_width)

    elif mode == "automatic":
        config_file = Path(config).resolve()
        cfg = auxiliaryfunctions.read_config(config_file)
        print("Config file read successfully.")

        numframes2pick = cfg['numframes2pick']
        start = cfg['start']
        stop = cfg['stop']

        # Check for variable correctness
        if start > 1 or stop > 1 or start < 0 or stop < 0 or start >= stop:
            raise Exception(
                "Erroneous start or stop values. Please correct it in the config file."
            )
        if numframes2pick < 1 and not int(numframes2pick):
            raise Exception(
                "Perhaps consider extracting more, or a natural number of frames."
            )

        videos = cfg['video_sets'].keys()
        if opencv:
            import cv2
        elif flymovie:
            from motmot.FlyMovieFormat import FlyMovieFormat as FMF
            import cv2
        else:
            from moviepy.editor import VideoFileClip
        for vindex, video in enumerate(videos):
            #plt.close("all")
            global coords
            coords = cfg['video_sets'][video]['crop'].split(',')

            if userfeedback:
                print(
                    "Do you want to extract (perhaps additional) frames for video:",
                    video, "?")
                askuser = input("yes/no")
            else:
                askuser = "******"

            if askuser == 'y' or askuser == 'yes' or askuser == 'Ja' or askuser == 'ha':  # multilanguage support :)
                if opencv:
                    cap = cv2.VideoCapture(video)
                    fps = cap.get(
                        5
                    )  #https://docs.opencv.org/2.4/modules/highgui/doc/reading_and_writing_images_and_video.html#videocapture-get
                    nframes = int(cap.get(7))
                    duration = nframes * 1. / fps
                elif flymovie:
                    cap = FMF.FlyMovie(video)
                    nframes = cap.n_frames
                    while True:
                        try:
                            cap.get_frame(nframes)
                        except FMF.NoMoreFramesException:
                            nframes -= 1
                            continue
                        break
                    fps = 1. / (cap.get_frame(min(100, nframes))[1] -
                                cap.get_frame(min(100, nframes) - 1)[1])
                    duration = cap.get_frame(nframes)[1]
                else:
                    #Moviepy:
                    clip = VideoFileClip(video)
                    fps = clip.fps
                    duration = clip.duration
                    nframes = int(np.ceil(clip.duration * 1. / fps))

                indexlength = int(np.ceil(np.log10(nframes)))

                if crop == True:
                    from deeplabcut.utils import select_crop_parameters
                    if opencv:
                        cap.set(2, start * duration)
                        ret, frame = cap.read()
                        if ret:
                            image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                    elif flymovie:
                        frame = cap.get_frame(int(nframes * start))[0]
                        if frame.ndim != 3:
                            frame = skimage.color.gray2rgb(frame)
                        image = frame
                    else:
                        image = clip.get_frame(
                            start * clip.duration
                        )  #frame is accessed by index *1./clip.fps (fps cancels)

                    fname = Path(video)
                    output_path = Path(
                        config).parents[0] / 'labeled-data' / fname.stem

                    if output_path.exists():
                        fig, ax = plt.subplots(1)
                        # Call the GUI to select the cropping parameters
                        coords = select_crop_parameters.show(config, image)
                        # Update the config.yaml file with current cropping parameters
                        cfg['video_sets'][video] = {
                            'crop':
                            ', '.join(
                                map(str, [
                                    int(coords[0]),
                                    int(coords[1]),
                                    int(coords[2]),
                                    int(coords[3])
                                ]))
                        }
                        auxiliaryfunctions.write_config(config_file, cfg)

                        if len(os.listdir(output_path)) == 0:  #check if empty
                            #store full frame from random location (good for augmentation)
                            index = int(start * duration +
                                        np.random.rand() * duration *
                                        (stop - start))
                            if opencv:
                                cap.set(1, index)
                                ret, frame = cap.read()
                                if ret:
                                    image = img_as_ubyte(
                                        cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
                            elif flymovie:
                                frame = cap.get_frame(int(nframes * start))[0]
                                if frame.ndim != 3:
                                    frame = skimage.color.gray2rgb(frame)
                                image = img_as_ubyte(frame)
                            else:
                                image = img_as_ubyte(
                                    clip.get_frame(index * 1. / clip.fps))
                                clip = clip.crop(
                                    y1=int(coords[2]),
                                    y2=int(coords[3]),
                                    x1=int(coords[0]),
                                    x2=int(coords[1]))  #now crop clip

                            saveimg = str(output_path) + '/img' + str(
                                index).zfill(indexlength) + ".png"
                            io.imsave(saveimg, image)

                        else:
                            askuser = input(
                                "The directory already contains some frames. Do you want to add to it?(yes/no): "
                            )
                            if askuser == 'y' or askuser == 'yes' or askuser == 'Y' or askuser == 'Yes':
                                #clip=clip.crop(y1 = int(coords[2]),y2 = int(coords[3]),x1 = int(coords[0]), x2 = int(coords[1]))
                                index = int(start * duration +
                                            np.random.rand() * duration *
                                            (stop - start))
                                if opencv:
                                    cap.set(1, index)
                                    ret, frame = cap.read()
                                    if ret:
                                        image = img_as_ubyte(
                                            cv2.cvtColor(
                                                frame, cv2.COLOR_BGR2RGB))
                                elif flymovie:
                                    frame = cap.get_frame(int(nframes *
                                                              start))[0]
                                    if frame.ndim != 3:
                                        frame = skimage.color.gray2rgb(frame)
                                    image = img_as_ubyte(frame)
                                else:
                                    image = img_as_ubyte(
                                        clip.get_frame(index * 1. / clip.fps))
                                    clip = clip.crop(y1=int(coords[2]),
                                                     y2=int(coords[3]),
                                                     x1=int(coords[0]),
                                                     x2=int(coords[1]))

                                saveimg = str(output_path) + '/img' + str(
                                    index).zfill(indexlength) + ".png"
                                io.imsave(saveimg, image)
                                pass
                            else:
                                sys.exit(
                                    "Delete the frames and try again later!")

                else:
                    numframes2pick = cfg[
                        'numframes2pick'] + 1  # without cropping a full size frame will not be extracted >> thus one more frame should be selected in next stage.

                print("Extracting frames based on %s ..." % algo)

                if algo == 'uniform':  #extract n-1 frames (0 was already stored)
                    if opencv:
                        frames2pick = frameselectiontools.UniformFramescv2(
                            cap, numframes2pick - 1, start, stop)
                    elif flymovie:
                        frames2pick = frameselectiontools.UniformFramesfmf(
                            cap, numframes2pick - 1, start, stop)
                    else:
                        frames2pick = frameselectiontools.UniformFrames(
                            clip, numframes2pick - 1, start, stop)
                elif algo == 'kmeans':
                    if opencv:
                        frames2pick = frameselectiontools.KmeansbasedFrameselectioncv2(
                            cap,
                            numframes2pick - 1,
                            start,
                            stop,
                            crop,
                            coords,
                            step=cluster_step,
                            resizewidth=cluster_resizewidth,
                            color=cluster_color)
                    elif flymovie:
                        print("FMF not supported by kmeans as of now!")
                        frames2pick = []
                    else:
                        frames2pick = frameselectiontools.KmeansbasedFrameselection(
                            clip,
                            numframes2pick - 1,
                            start,
                            stop,
                            step=cluster_step,
                            resizewidth=cluster_resizewidth,
                            color=cluster_color)
                else:
                    print(
                        "Please implement this method yourself and send us a pull request! Otherwise, choose 'uniform' or 'kmeans'."
                    )
                    frames2pick = []

                output_path = Path(config).parents[0] / 'labeled-data' / Path(
                    video).stem
                if opencv:
                    for index in frames2pick:
                        cap.set(1, index)  #extract a particular frame
                        ret, frame = cap.read()
                        if ret:
                            image = img_as_ubyte(
                                cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
                            img_name = str(output_path) + '/img' + str(
                                index).zfill(indexlength) + ".png"
                            if crop:
                                io.imsave(
                                    img_name,
                                    image[int(coords[2]):int(coords[3]),
                                          int(coords[0]):int(coords[1]), :]
                                )  #y1 = int(coords[2]),y2 = int(coords[3]),x1 = int(coords[0]), x2 = int(coords[1]
                            else:
                                io.imsave(img_name, image)
                        else:
                            print("Frame", index, " not found!")
                    cap.release()
                elif flymovie:
                    for index in frames2pick:
                        print(index)
                        frame = cap.get_frame(int(index))[0]
                        if frame.ndim != 3:
                            frame = skimage.color.gray2rgb(frame)
                        image = img_as_ubyte(frame)
                        img_name = str(output_path) + '/img' + str(
                            index).zfill(indexlength) + ".png"
                        if crop:
                            io.imsave(
                                img_name,
                                image[int(coords[2]):int(coords[3]),
                                      int(coords[0]):int(coords[1]), :]
                            )  #y1 = int(coords[2]),y2 = int(coords[3]),x1 = int(coords[0]), x2 = int(coords[1]
                        else:
                            io.imsave(img_name, image)
                    cap.close()
                else:
                    for index in frames2pick:
                        try:
                            image = img_as_ubyte(
                                clip.get_frame(index * 1. / clip.fps))
                            img_name = str(output_path) + '/img' + str(
                                index).zfill(indexlength) + ".png"
                            io.imsave(img_name, image)
                            if np.var(image) == 0:  #constant image
                                print(
                                    "Seems like black/constant images are extracted from your video. Perhaps consider using opencv under the hood, by setting: opencv=True"
                                )

                        except FileNotFoundError:
                            print("Frame # ", index, " does not exist.")

                    #close video.
                    clip.close()
                    del clip
    else:
        print(
            "Invalid MODE. Choose either 'manual' or 'automatic'. Check ``help(deeplabcut.extract_frames)`` on python and ``deeplabcut.extract_frames?`` \
              for ipython/jupyter notebook for more details.")

    print(
        "\nFrames were selected.\nYou can now label the frames using the function 'label_frames' (if you extracted enough frames for all videos)."
    )
def extract_frames(config,
                   mode='automatic',
                   algo='kmeans',
                   crop=False,
                   userfeedback=True,
                   cluster_step=1,
                   cluster_resizewidth=30,
                   cluster_color=False,
                   opencv=True,
                   slider_width=25):
    """
    Extracts frames from the videos in the config.yaml file. Only the videos in the config.yaml will be used to select the frames.\n
    Use the function ``add_new_video`` at any stage of the project to add new videos to the config file and extract their frames.

    The provided function either selects frames from the videos in a randomly and temporally uniformly distributed way (uniform), \n
    by clustering based on visual appearance (k-means), or by manual selection.

    Three important parameters for automatic extraction: numframes2pick, start and stop are set in the config file.

    Please refer to the user guide for more details on methods and parameters https://www.nature.com/articles/s41596-019-0176-0
    or the preprint: https://www.biorxiv.org/content/biorxiv/early/2018/11/24/476531.full.pdf

    Parameters
    ----------
    config : string
        Full path of the config.yaml file as a string.

    mode : string
        String containing the mode of extraction. It must be either ``automatic`` or ``manual``.

    algo : string
        String specifying the algorithm to use for selecting the frames. Currently, deeplabcut supports either ``kmeans`` or ``uniform`` based selection. This flag is
        only required for ``automatic`` mode and the default is ``uniform``. For uniform, frames are picked in temporally uniform way, kmeans performs clustering on downsampled frames (see user guide for details).
        Note: color information is discarded for kmeans, thus e.g. for camouflaged octopus clustering one might want to change this.

    crop : bool, optional
        If True, video frames are cropped according to the corresponding coordinates stored in the config.yaml.
        Alternatively, if cropping coordinates are not known yet, crop='GUI' triggers a user interface
        where the cropping area can be manually drawn and saved.

    userfeedback: bool, optional
        If this is set to false during automatic mode then frames for all videos are extracted. The user can set this to true, which will result in a dialog,
        where the user is asked for each video if (additional/any) frames from this video should be extracted. Use this, e.g. if you have already labeled
        some folders and want to extract data for new videos.

    cluster_resizewidth: number, default: 30
        For k-means one can change the width to which the images are downsampled (aspect ratio is fixed).

    cluster_step: number, default: 1
        By default each frame is used for clustering, but for long videos one could only use every nth frame (set by: cluster_step). This saves memory before clustering can start, however,
        reading the individual frames takes longer due to the skipping.

    cluster_color: bool, default: False
        If false then each downsampled image is treated as a grayscale vector (discarding color information). If true, then the color channels are considered. This increases
        the computational complexity.

    opencv: bool, default: True
        Uses openCV for loading & extractiong (otherwise moviepy (legacy))

    slider_width: number, default: 25
        Width of the video frames slider, in percent of window

    Examples
    --------
    for selecting frames automatically with 'kmeans' and want to crop the frames.
    >>> deeplabcut.extract_frames('/analysis/project/reaching-task/config.yaml','automatic','kmeans',True)
    --------
    for selecting frames automatically with 'kmeans' and defining the cropping area at runtime.
    >>> deeplabcut.extract_frames('/analysis/project/reaching-task/config.yaml','automatic','kmeans','GUI')
    --------
    for selecting frames automatically with 'kmeans' and considering the color information.
    >>> deeplabcut.extract_frames('/analysis/project/reaching-task/config.yaml','automatic','kmeans',cluster_color=True)
    --------
    for selecting frames automatically with 'uniform' and want to crop the frames.
    >>> deeplabcut.extract_frames('/analysis/project/reaching-task/config.yaml','automatic',crop=True)
    --------
    for selecting frames manually,
    >>> deeplabcut.extract_frames('/analysis/project/reaching-task/config.yaml','manual')
    --------
    for selecting frames manually, with a 60% wide frames slider
    >>> deeplabcut.extract_frames('/analysis/project/reaching-task/config.yaml','manual', slider_width=60)

    While selecting the frames manually, you do not need to specify the ``crop`` parameter in the command. Rather, you will get a prompt in the graphic user interface to choose
    if you need to crop or not.
    --------

    """
    import os
    import sys
    import numpy as np
    from pathlib import Path
    from skimage import io
    from skimage.util import img_as_ubyte
    from deeplabcut.utils import frameselectiontools
    from deeplabcut.utils import auxiliaryfunctions

    if mode == "manual":
        wd = Path(config).resolve().parents[0]
        os.chdir(str(wd))
        from deeplabcut.generate_training_dataset import frame_extraction_toolbox
        from deeplabcut.utils import select_crop_parameters
        frame_extraction_toolbox.show(config, slider_width)

    elif mode == "automatic":
        config_file = Path(config).resolve()
        cfg = auxiliaryfunctions.read_config(config_file)
        print("Config file read successfully.")

        numframes2pick = cfg['numframes2pick']
        start = cfg['start']
        stop = cfg['stop']

        # Check for variable correctness
        if start > 1 or stop > 1 or start < 0 or stop < 0 or start >= stop:
            raise Exception(
                "Erroneous start or stop values. Please correct it in the config file."
            )
        if numframes2pick < 1 and not int(numframes2pick):
            raise Exception(
                "Perhaps consider extracting more, or a natural number of frames."
            )

        videos = cfg['video_sets'].keys()
        if opencv:
            import cv2
        else:
            from moviepy.editor import VideoFileClip

        has_failed = []
        for vindex, video in enumerate(videos):
            if userfeedback:
                print(
                    "Do you want to extract (perhaps additional) frames for video:",
                    video, "?")
                askuser = input("yes/no")
            else:
                askuser = "******"

            if askuser == 'y' or askuser == 'yes' or askuser == 'Ja' or askuser == 'ha'\
                    or askuser == 'oui' or askuser == 'ouais':  # multilanguage support :)
                if opencv:
                    cap = cv2.VideoCapture(video)
                    fps = cap.get(
                        5
                    )  # https://docs.opencv.org/2.4/modules/highgui/doc/reading_and_writing_images_and_video.html#videocapture-get
                    nframes = int(cap.get(7))
                else:
                    # Moviepy:
                    clip = VideoFileClip(video)
                    fps = clip.fps
                    nframes = int(np.ceil(clip.duration * 1. / fps))
                indexlength = int(np.ceil(np.log10(nframes)))

                fname = Path(video)
                output_path = Path(
                    config).parents[0] / 'labeled-data' / fname.stem

                if output_path.exists():
                    if len(os.listdir(output_path)):
                        askuser = input(
                            "The directory already contains some frames. Do you want to add to it?(yes/no): "
                        )
                        if not (askuser == 'y' or askuser == 'yes'
                                or askuser == 'Y' or askuser == 'Yes'):
                            sys.exit("Delete the frames and try again later!")

                if crop == 'GUI':
                    cfg = select_cropping_area(config, [video])
                coords = cfg['video_sets'][video]['crop'].split(',')
                if crop and not opencv:
                    clip = clip.crop(y1=int(coords[2]),
                                     y2=int(coords[3]),
                                     x1=int(coords[0]),
                                     x2=int(coords[1]))
                elif not crop:
                    coords = None

                print("Extracting frames based on %s ..." % algo)
                if algo == 'uniform':
                    if opencv:
                        frames2pick = frameselectiontools.UniformFramescv2(
                            cap, numframes2pick, start, stop)
                    else:
                        frames2pick = frameselectiontools.UniformFrames(
                            clip, numframes2pick, start, stop)
                elif algo == 'kmeans':
                    if opencv:
                        frames2pick = frameselectiontools.KmeansbasedFrameselectioncv2(
                            cap,
                            numframes2pick,
                            start,
                            stop,
                            crop,
                            coords,
                            step=cluster_step,
                            resizewidth=cluster_resizewidth,
                            color=cluster_color)
                    else:
                        frames2pick = frameselectiontools.KmeansbasedFrameselection(
                            clip,
                            numframes2pick,
                            start,
                            stop,
                            step=cluster_step,
                            resizewidth=cluster_resizewidth,
                            color=cluster_color)
                else:
                    print(
                        "Please implement this method yourself and send us a pull request! Otherwise, choose 'uniform' or 'kmeans'."
                    )
                    frames2pick = []

                if not len(frames2pick):
                    print('Frame selection failed...')
                    return

                output_path = Path(config).parents[0] / 'labeled-data' / Path(
                    video).stem
                is_valid = []
                if opencv:
                    for index in frames2pick:
                        cap.set(1, index)  # extract a particular frame
                        ret, frame = cap.read()
                        if ret:
                            image = img_as_ubyte(
                                cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
                            img_name = str(output_path) + '/img' + str(
                                index).zfill(indexlength) + ".png"
                            if crop:
                                io.imsave(
                                    img_name,
                                    image[int(coords[2]):int(coords[3]),
                                          int(coords[0]):int(coords[1]), :]
                                )  # y1 = int(coords[2]),y2 = int(coords[3]),x1 = int(coords[0]), x2 = int(coords[1]
                            else:
                                io.imsave(img_name, image)
                            is_valid.append(True)
                        else:
                            print("Frame", index, " not found!")
                            is_valid.append(False)
                    cap.release()
                else:
                    for index in frames2pick:
                        try:
                            image = img_as_ubyte(
                                clip.get_frame(index * 1. / clip.fps))
                            img_name = str(output_path) + '/img' + str(
                                index).zfill(indexlength) + ".png"
                            io.imsave(img_name, image)
                            if np.var(image) == 0:  # constant image
                                print(
                                    "Seems like black/constant images are extracted from your video. Perhaps consider using opencv under the hood, by setting: opencv=True"
                                )
                            is_valid.append(True)
                        except FileNotFoundError:
                            print("Frame # ", index, " does not exist.")
                            is_valid.append(False)
                    clip.close()
                    del clip

                if not any(is_valid):
                    has_failed.append(True)
                else:
                    has_failed.append(False)

        if all(has_failed):
            print('Frame extraction failed. Video files must be corrupted.')
            return
        elif any(has_failed):
            print('Although most frames were extracted, some were invalid.')
        else:
            print("Frames were successfully extracted.")
        print(
            "\nYou can now label the frames using the function 'label_frames' "
            "(if you extracted enough frames for all videos).")
    else:
        print(
            "Invalid MODE. Choose either 'manual' or 'automatic'. Check ``help(deeplabcut.extract_frames)`` on python and ``deeplabcut.extract_frames?`` \
              for ipython/jupyter notebook for more details.")
Example #6
0
def extract_frames(config,
                   algo='uniform',
                   crop=False,
                   checkcropping=False,
                   userfeedback=True,
                   Screens=1,
                   scale_w=.8,
                   scale_h=.8,
                   opencv=True):
    """
    Extracts frames from the videos in the config.yaml file. Only the videos in the config.yaml will be used to select the frames.\n
    Use the function ``add_new_video`` at any stage of the project to add new videos to the config file and extract their frames.
    
    The provided function either selects frames from the videos in a randomly and temporally uniformly distributed way (uniform), \n 
    by clustering based on visual appearance (k-means), or by manual selection. 
    
    Three important parameters for automatic extraction: numframes2pick, start and stop are set in the config file. 
    
    Please refer to the user guide for more details on methods and parameters https://www.biorxiv.org/content/biorxiv/early/2018/11/24/476531.full.pdf
    
    Parameters
    ----------
    config : string
        Full path of the config.yaml file as a string.
        
    mode : string
        String containing the mode of extraction. It must be either ``automatic`` or ``manual``.
        
    algo : string 
        String specifying the algorithm to use for selecting the frames. Currently, deeplabcut supports either ``kmeans`` or ``uniform`` based selection. This flag is
        only required for ``automatic`` mode and the default is ``uniform``. For uniform, frames are picked in temporally uniform way, kmeans performs clustering on downsampled frames (see user guide for details).
        Note: color information is discarded for kmeans, thus e.g. for camouflaged octopus clustering one might want to change this. 
        
    crop : bool, optional
        If this is set to True, the selected frames are cropped based on the ``crop`` parameters in the config.yaml file. 
        The default is ``False``; if provided it must be either ``True`` or ``False``.
        
    checkcropping: bool, optional
        If this is set to True, the cropping parameters are overlayed in a plot of the first frame to check and the user can decide if the program should proceed 
        with those parameters, or perhaps edit them. The default is ``False``; if provided it must be either ``True`` or ``False``.
    
    userfeedback: bool, optional
        If this is set to false during automatic mode then frames for all videos are extracted. The user can set this to true, which will result in a dialog,
        where the user is asked for each video if (additional/any) frames from this video should be extracted. Use this, e.g. if you have already labeled
        some folders and want to extract data for new videos. 
    
    cluster_resizewidth: number, default: 30
        For k-means one can change the width to which the images are downsampled (aspect ratio is fixed).
    
    cluster_step: number, default: 1
        By default each frame is used for clustering, but for long videos one could only use every nth frame (set by: cluster_step). This saves memory before clustering can start, however, 
        reading the individual frames takes longer due to the skipping.
    
    cluster_color: bool, default: False
        If false then each downsampled image is treated as a grayscale vector (discarding color information). If true, then the color channels are considered. This increases 
        the computational complexity. 
    
    opencv: bool, default: True
        Uses openCV for loading & extractiong (otherwise moviepy (legacy))
        

    The three parameters Screens=1,scale_w=.8,scale_h=.8 define the relative height (scale_h), relative widht (scale_w) and number of screens (horizontally) and thereby 
    affect the dimensions of the manual frame extraction GUI.
        
    Examples
    --------
    for selecting frames automatically with 'kmeans' and want to crop the frames based on the ``crop`` parameters in config.yaml
    >>> deeplabcut.extract_frames('/analysis/project/reaching-task/config.yaml','automatic','kmeans',True)
    --------
    for selecting frames automatically with 'kmeans' and considering the color information.
    >>> deeplabcut.extract_frames('/analysis/project/reaching-task/config.yaml','automatic','kmeans',cluster_color=True)
    --------
    for selecting frames automatically with 'uniform' and want to crop the frames based on the ``crop`` parameters in config.yaml
    >>> deeplabcut.extract_frames('/analysis/project/reaching-task/config.yaml','automatic',crop=True)
    --------
    for selecting frames automatically with 'uniform', want to crop the frames based on the ``crop`` parameters in config.yaml and check for cropping
    >>> deeplabcut.extract_frames('/analysis/project/reaching-task/config.yaml','automatic',crop=True,checkcropping=True)
    --------
    for selecting frames manually,
    >>> deeplabcut.extract_frames('/analysis/project/reaching-task/config.yaml','manual')
    
    While selecting the frames manually, you do not need to specify the ``crop`` parameter in the command. Rather, you will get a prompt in the graphic user interface to choose 
    if you need to crop or not.
    --------
    
    """
    import os
    import sys
    import yaml
    import numpy as np
    import math
    from pathlib import Path
    from skimage import io
    from skimage.util import img_as_ubyte
    import matplotlib.pyplot as plt
    import matplotlib.patches as patches
    from deeplabcut.utils import frameselectiontools

    config_file = Path(config).resolve()
    with open(str(config_file), 'r') as ymlfile:
        cfg = yaml.load(ymlfile)
    print("Config file read successfully.")

    numframes2pick = cfg['numframes2pick']

    if numframes2pick < 1 and not int(numframes2pick):
        raise Exception(
            "Perhaps consider extracting more, or a natural number of frames.")

    videos = cfg['video_sets'].keys()
    if opencv:
        import cv2
    else:
        from moviepy.editor import VideoFileClip
    for vindex, video in enumerate(videos):
        # plt.close("all")
        coords = cfg['video_sets'][video]['crop'].split(',')

        if userfeedback:
            print(
                "Do you want to extract (perhaps additional) frames for video:",
                video, "?")
            askuser = input("yes/no")
        else:
            askuser = "******"
        start, stop = 0, 1
        if askuser == 'y' or askuser == 'yes' or askuser == 'Ja' or askuser == 'ha':  # multilanguage support :)
            # indexlength = int(np.ceil(np.log10(clip.duration * clip.fps)))
            if opencv:
                cap = cv2.VideoCapture(video)
                fps = cap.get(
                    5
                )  # https://docs.opencv.org/2.4/modules/highgui/doc/reading_and_writing_images_and_video.html#videocapture-get
                nframes = int(cap.get(7))
                duration = nframes * 1. / fps
            else:
                # Moviepy:
                clip = VideoFileClip(video)
                fps = clip.fps
                duration = clip.duration
                nframes = int(np.ceil(clip.duration * 1. / fps))
            indexlength = int(np.ceil(np.log10(nframes)))
            if crop == True:
                print(
                    "Make sure you change the crop parameters in the config.yaml file. The default parameters are set to the video dimensions."
                )
                if opencv:
                    cap.set(2, start * duration)
                    ret, frame = cap.read()
                    if ret:
                        image = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
                else:
                    image = clip.get_frame(
                        start * clip.duration
                    )  # frame is accessed by index *1./clip.fps (fps cancels)

                fname = Path(video)
                output_path = Path(config).parents[
                    0] / 'pretrain' / 'training-datasets' / fname.stem

                if output_path.exists() and checkcropping == True:
                    fig, ax = plt.subplots(1)
                    # Display the image
                    ax.imshow(image)
                    # Create a Rectangle patch
                    rect = patches.Rectangle((int(coords[0]), int(coords[2])),
                                             int(coords[1]) - int(coords[0]),
                                             int(coords[3]) - int(coords[2]),
                                             linewidth=3,
                                             edgecolor='r',
                                             facecolor='none')
                    # Add the patch to the Axes
                    ax.add_patch(rect)
                    plt.show()

                    print(
                        "The red boundary indicates how the cropped image will look."
                    )
                    # saveimg = str(Path(config).parents[0] / Path('labeled-data','IsCroppingOK_'+fname.stem +".png"))
                    # io.imsave(saveimg, image)

                    msg = input("Is the cropping ok? (yes/no): ")
                    if msg == "yes" or msg == "y" or msg == "Yes" or msg == "Y":
                        if len(os.listdir(output_path)) == 0:  # check if empty
                            # store full frame from random location (good for augmentation)
                            index = int(start * duration +
                                        np.random.rand() * duration *
                                        (stop - start))
                            if opencv:
                                cap.set(1, index)
                                ret, frame = cap.read()
                                if ret:
                                    image = img_as_ubyte(
                                        cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
                            else:
                                image = img_as_ubyte(
                                    clip.get_frame(index * 1. / clip.fps))
                                clip = clip.crop(
                                    y1=int(coords[2]),
                                    y2=int(coords[3]),
                                    x1=int(coords[0]),
                                    x2=int(coords[1]))  # now crop clip

                            saveimg = str(output_path) + '/img' + str(
                                index).zfill(indexlength) + ".png"
                            io.imsave(saveimg, image)

                        else:
                            askuser = input(
                                "The directory already contains some frames. Do you want to add to it?(yes/no): "
                            )
                            if askuser == 'y' or askuser == 'yes' or askuser == 'Y' or askuser == 'Yes':
                                # clip=clip.crop(y1 = int(coords[2]),y2 = int(coords[3]),x1 = int(coords[0]), x2 = int(coords[1]))
                                pass
                            else:
                                sys.exit(
                                    "Delete the frames and try again later!")
                    else:
                        sys.exit(
                            "Correct the crop parameters in the config.yaml file and try again!"
                        )

                elif output_path.exists():  # cropping without checking:
                    index = int(start * duration +
                                np.random.rand() * duration * (stop - start))
                    if opencv:
                        cap.set(1, index)
                        ret, frame = cap.read()
                        if ret:
                            image = img_as_ubyte(
                                cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
                    else:
                        image = img_as_ubyte(
                            clip.get_frame(index * 1. / clip.fps))
                        clip = clip.crop(y1=int(coords[2]),
                                         y2=int(coords[3]),
                                         x1=int(coords[0]),
                                         x2=int(coords[1]))

                    saveimg = str(output_path) + '/img' + str(index).zfill(
                        indexlength) + ".png"
                    io.imsave(saveimg, image)

            else:
                numframes2pick = cfg[
                    'numframes2pick'] + 1  # without cropping a full size frame will not be extracted >> thus one more frame should be selected in next stage.

            print("Extracting frames")
            if algo == 'uniform':  # extract n-1 frames (0 was already stored)
                if opencv:
                    frames2pick = frameselectiontools.UniformFramescv2(
                        cap, numframes2pick - 1, start, stop)
                else:
                    frames2pick = frameselectiontools.UniformFrames(
                        clip, numframes2pick - 1, start, stop)
            elif algo == 'all':
                if opencv:
                    # nframes = int(cap.get(7))
                    frames2pick = range(
                        0, numframes2pick)  # math.ceil(nframes * stop)
                else:
                    frames2pick = range(
                        0, numframes2pick
                    )  # math.ceil(clip.duration * clip.fps * stop)

            output_path = Path(config).parents[
                0] / 'pretrain' / 'training-datasets' / Path(video).stem
            if opencv:
                for index in frames2pick:
                    cap.set(1, index)  # extract a particular frame
                    ret, frame = cap.read()
                    if ret:
                        image = img_as_ubyte(
                            cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
                        img_name = str(output_path) + '/img' + str(
                            index).zfill(indexlength) + ".png"
                        if crop:
                            io.imsave(
                                img_name,
                                image[int(coords[2]):int(coords[3]),
                                      int(coords[0]):int(coords[1]), :]
                            )  # y1 = int(coords[2]),y2 = int(coords[3]),x1 = int(coords[0]), x2 = int(coords[1]
                        else:
                            io.imsave(img_name, image)
                    else:
                        print("Frame", index, " not found!")
                cap.release()
            else:
                for index in frames2pick:
                    try:
                        image = img_as_ubyte(
                            clip.get_frame(index * 1. / clip.fps))
                        img_name = str(output_path) + '/img' + str(
                            index).zfill(indexlength) + ".png"
                        io.imsave(img_name, image)
                        if np.var(image) == 0:  # constant image
                            print(
                                "Seems like black/constant images are extracted from your video. Perhaps consider using opencv under the hood, by setting: opencv=True"
                            )

                    except FileNotFoundError:
                        print("Frame # ", index, " does not exist.")

                # close video.
                clip.close()
                del clip

    print(
        "\nFrames were selected.\nYou can now label the frames using the function 'label_frames' (if you extracted enough frames for all videos)."
    )