diff --git a/deeplabcut/__init__.py b/deeplabcut/__init__.py
index 22aad797e2..5e03f373f9 100644
--- a/deeplabcut/__init__.py
+++ b/deeplabcut/__init__.py
@@ -60,7 +60,6 @@
 from deeplabcut.generate_training_dataset import (
     create_training_model_comparison,
     create_multianimaltraining_dataset,
-    cropimagesandlabels,
 )
 from deeplabcut.generate_training_dataset import (
     dropannotationfileentriesduetodeletedimages,
diff --git a/deeplabcut/create_project/new.py b/deeplabcut/create_project/new.py
index 86f0c91f42..a6571236cf 100644
--- a/deeplabcut/create_project/new.py
+++ b/deeplabcut/create_project/new.py
@@ -220,7 +220,6 @@ def create_new_project(
         cfg_file["skeleton"] = [["bodypart1", "bodypart2"], ["objectA", "bodypart3"]]
         cfg_file["default_augmenter"] = "default"
         cfg_file["default_net_type"] = "resnet_50"
-    cfg_file["croppedtraining"] = False
 
     # common parameters:
     cfg_file["Task"] = project
diff --git a/deeplabcut/generate_training_dataset/multiple_individuals_trainingsetmanipulation.py b/deeplabcut/generate_training_dataset/multiple_individuals_trainingsetmanipulation.py
index 2a1e4fb23c..73e9fe1b2b 100755
--- a/deeplabcut/generate_training_dataset/multiple_individuals_trainingsetmanipulation.py
+++ b/deeplabcut/generate_training_dataset/multiple_individuals_trainingsetmanipulation.py
@@ -24,6 +24,7 @@
     MakeTrain_pose_yaml,
     MakeTest_pose_yaml,
     MakeInference_yaml,
+    pad_train_test_indices,
 )
 from deeplabcut.utils import auxiliaryfunctions, auxfun_models, auxfun_multianimal
 
@@ -99,6 +100,8 @@ def create_multianimaltraining_dataset(
     windows2linux=False,
     net_type=None,
     numdigits=2,
+    crop_size=(400, 400),
+    crop_sampling="hybrid",
     paf_graph=None,
     trainIndices=None,
     testIndices=None,
@@ -133,6 +136,18 @@ def create_multianimaltraining_dataset(
 
     numdigits: int, optional
 
+    crop_size: tuple of int, optional
+        Dimensions (width, height) of the crops for data augmentation.
+        Default is 400x400.
+
+    crop_sampling: str, optional
+        Crop centers sampling method. Must be either:
+        "uniform" (randomly over the image),
+        "keypoints" (randomly over the annotated keypoints),
+        "density" (weighing preferentially dense regions of keypoints),
+        or "hybrid" (alternating randomly between "uniform" and "density").
+        Default is "hybrid".
+
     paf_graph: list of lists, optional (default=None)
         If not None, overwrite the default complete graph. This is useful for advanced users who
         already know a good graph, or simply want to use a specific one. Note that, in that case,
@@ -155,6 +170,12 @@ def create_multianimaltraining_dataset(
     >>> deeplabcut.create_multianimaltraining_dataset(r'C:\\Users\\Ulf\\looming-task\\config.yaml',Shuffles=[3,17,5])
     --------
     """
+    if len(crop_size) != 2 or not all(isinstance(v, int) for v in crop_size):
+        raise ValueError("Crop size must be a tuple of two integers (width, height).")
+
+    if crop_sampling not in ("uniform", "keypoints", "density", "hybrid"):
+            raise ValueError(f"Invalid sampling {crop_sampling}. Must be "
+                             f"either 'uniform', 'keypoints', 'density', or 'hybrid.")
 
     # Loading metadata from config file:
     cfg = auxiliaryfunctions.read_config(config)
@@ -170,15 +191,6 @@ def create_multianimaltraining_dataset(
         return
     Data = Data[scorer]
 
-    def strip_cropped_image_name(path):
-        # utility function to split different crops from same image into either train or test!
-        head, filename = os.path.split(path)
-        if cfg["croppedtraining"]:
-            filename = filename.split("c")[0]
-        return os.path.join(head, filename)
-
-    img_names = Data.index.map(strip_cropped_image_name).unique()
-
     if net_type is None:  # loading & linking pretrained models
         net_type = cfg.get("default_net_type", "dlcrnet_ms5")
     elif not any(net in net_type for net in ("resnet", "eff", "dlc", "mob")):
@@ -236,19 +248,12 @@ def strip_cropped_image_name(path):
     if trainIndices is None and testIndices is None:
         splits = []
         for shuffle in Shuffles:  # Creating shuffles starting from 1
-            for trainFraction in cfg["TrainingFraction"]:
-                train_inds_temp, test_inds_temp = SplitTrials(
-                    range(len(img_names)), trainFraction
+            for train_frac in cfg["TrainingFraction"]:
+                train_inds, test_inds = SplitTrials(
+                    range(len(Data)), train_frac
                 )
-                # Map back to the original indices.
-                temp = [re.escape(name) for i, name in enumerate(img_names)
-                        if i in test_inds_temp]
-                mask = Data.index.str.contains("|".join(temp))
-                testIndices = np.flatnonzero(mask)
-                trainIndices = np.flatnonzero(~mask)
-
                 splits.append(
-                    (trainFraction, shuffle, (trainIndices, testIndices))
+                    (train_frac, shuffle, (train_inds, test_inds))
                 )
     else:
         if len(trainIndices) != len(testIndices) != len(Shuffles):
@@ -265,6 +270,12 @@ def strip_cropped_image_name(path):
             print(
                 f"You passed a split with the following fraction: {int(100 * trainFraction)}%"
             )
+            # Now that the training fraction is guaranteed to be correct,
+            # the values added to pad the indices are removed.
+            train_inds = np.asarray(train_inds)
+            train_inds = train_inds[train_inds != -1]
+            test_inds = np.asarray(test_inds)
+            test_inds = test_inds[test_inds != -1]
             splits.append(
                 (trainFraction, Shuffles[shuffle], (train_inds, test_inds))
             )
@@ -387,6 +398,8 @@ def strip_cropped_image_name(path):
                 "num_idchannel": len(cfg["individuals"])
                 if cfg.get("identity", False)
                 else 0,
+                "crop_size": list(crop_size),
+                "crop_sampling": crop_sampling,
             }
 
             trainingdata = MakeTrain_pose_yaml(
@@ -441,3 +454,115 @@ def strip_cropped_image_name(path):
             )
         else:
             pass
+
+
+def convert_cropped_to_standard_dataset(
+    config_path,
+    recreate_datasets=True,
+    delete_crops=True,
+    back_up=True,
+):
+    import pandas as pd
+    import pickle
+    import shutil
+    from deeplabcut.generate_training_dataset import trainingsetmanipulation
+    from deeplabcut.utils import read_plainconfig, write_config
+
+    cfg = auxiliaryfunctions.read_config(config_path)
+    videos_orig = cfg.pop("video_sets_original")
+    is_cropped = cfg.pop("croppedtraining")
+    if videos_orig is None or not is_cropped:
+        print("Labeled data do not appear to be cropped. "
+              "Project will remain unchanged...")
+        return
+
+    project_path = cfg["project_path"]
+
+    if back_up:
+        print("Backing up project...")
+        shutil.copytree(project_path, project_path + "_bak", symlinks=True)
+
+    if delete_crops:
+        print("Deleting crops...")
+        data_path = os.path.join(project_path, "labeled-data")
+        for video in cfg["video_sets"]:
+            _, filename, _ = trainingsetmanipulation._robust_path_split(video)
+            if "_cropped" in video:  # One can never be too safe...
+                shutil.rmtree(os.path.join(data_path, filename), ignore_errors=True)
+
+    cfg["video_sets"] = videos_orig
+    write_config(config_path, cfg)
+
+    if not recreate_datasets:
+        return
+
+    datasets_folder = os.path.join(
+        project_path, auxiliaryfunctions.GetTrainingSetFolder(cfg),
+    )
+    df_old = pd.read_hdf(
+        os.path.join(datasets_folder, "CollectedData_" + cfg["scorer"] + ".h5"),
+    )
+
+    def strip_cropped_image_name(path):
+        head, filename = os.path.split(path)
+        head = head.replace("_cropped", "")
+        file, ext = filename.split(".")
+        file = file.split("c")[0]
+        return os.path.join(head, file + "." + ext)
+
+    img_names_old = np.asarray(
+        [strip_cropped_image_name(img) for img in df_old.index.to_list()]
+    )
+    df = merge_annotateddatasets(cfg, datasets_folder, False)
+    img_names = df.index.to_numpy()
+    train_idx = []
+    test_idx = []
+    pickle_files = []
+    for filename in os.listdir(datasets_folder):
+        if filename.endswith("pickle"):
+            pickle_file = os.path.join(datasets_folder, filename)
+            pickle_files.append(pickle_file)
+            if filename.startswith("Docu"):
+                with open(pickle_file, "rb") as f:
+                    _, train_inds, test_inds, train_frac = pickle.load(f)
+                    train_inds_temp = np.flatnonzero(
+                        np.isin(img_names, img_names_old[train_inds])
+                    )
+                    test_inds_temp = np.flatnonzero(
+                        np.isin(img_names, img_names_old[test_inds])
+                    )
+                    train_inds, test_inds = pad_train_test_indices(
+                        train_inds_temp, test_inds_temp, train_frac
+                    )
+                    train_idx.append(train_inds)
+                    test_idx.append(test_inds)
+
+    # Search a pose_config.yaml file to parse missing information
+    pose_config_path = ""
+    for dirpath, dirnames, filenames in os.walk(
+            os.path.join(project_path, "dlc-models")
+    ):
+        for file in filenames:
+            if file.endswith("pose_cfg.yaml"):
+                pose_config_path = os.path.join(dirpath, file)
+                break
+    pose_cfg = read_plainconfig(pose_config_path)
+    net_type = pose_cfg["net_type"]
+    if net_type == "resnet_50" and pose_cfg.get("multi_stage", False):
+        net_type = "dlcrnet_ms5"
+
+    # Clean the training-datasets folder prior to recreating the data pickles
+    shuffle_inds = set()
+    for file in pickle_files:
+        os.remove(file)
+        shuffle_inds.add(int(re.findall(r"shuffle(\d+)", file)[0]))
+    create_multianimaltraining_dataset(
+        config_path,
+        trainIndices=train_idx,
+        testIndices=test_idx,
+        Shuffles=sorted(shuffle_inds),
+        net_type=net_type,
+        paf_graph=pose_cfg["partaffinityfield_graph"],
+        crop_size=pose_cfg.get("crop_size", [400, 400]),
+        crop_sampling=pose_cfg.get("crop_sampling", "hybrid"),
+    )
diff --git a/deeplabcut/generate_training_dataset/trainingsetmanipulation.py b/deeplabcut/generate_training_dataset/trainingsetmanipulation.py
index 45631b2396..89ec13c157 100755
--- a/deeplabcut/generate_training_dataset/trainingsetmanipulation.py
+++ b/deeplabcut/generate_training_dataset/trainingsetmanipulation.py
@@ -19,7 +19,6 @@
 import numpy as np
 import pandas as pd
 import yaml
-from skimage import io
 
 from deeplabcut.pose_estimation_tensorflow import training
 from deeplabcut.utils import (
@@ -237,237 +236,6 @@ def dropimagesduetolackofannotation(config):
         )
 
 
-def cropimagesandlabels(
-    config,
-    numcrops=10,
-    size=(400, 400),
-    userfeedback=True,
-    cropdata=True,
-    excludealreadycropped=False,
-    updatevideoentries=True,
-):
-    """
-    Crop images into multiple random crops (defined by numcrops) of size dimensions. If cropdata=True then the
-    annotation data is loaded and labels for cropped images are inherited.
-    If false, then one can make crops for unlabeled folders.
-
-    This can be helpul for large frames with multiple animals. Then a smaller set of equally sized images is created.
-
-    Parameters
-    ----------
-    config : string
-        String containing the full path of the config file in the project.
-
-    numcrops: number of random crops (around random bodypart)
-
-    size: height x width in pixels
-
-    userfeedback: bool, optional
-        If this is set to false, then all requested train/test splits are created (no matter if they already exist). If you
-        want to assure that previous splits etc. are not overwritten, then set this to True and you will be asked for each split.
-
-    cropdata: bool, default True:
-        If true creates corresponding annotation data (from ground truth)
-
-    excludealreadycropped: bool, default False:
-        If true, ignore original videos whose frames are already cropped.
-        This is only useful after adding new videos post dataset creation,
-        as folders containing no new frames are otherwise automatically ignored.
-
-    updatevideoentries, bool, default true
-        If true updates video_list entries to refer to cropped frames instead. This makes sense for subsequent processing.
-
-    Example
-    --------
-    for labeling the frames
-    >>> deeplabcut.cropimagesandlabels('/analysis/project/reaching-task/config.yaml')
-
-    --------
-    """
-    from tqdm import trange
-
-    indexlength = int(np.ceil(np.log10(numcrops)))
-    project_path = os.path.dirname(config)
-    cfg = auxiliaryfunctions.read_config(config)
-    videos = list(cfg.get("video_sets_original", []))
-    if not videos:
-        videos = list(cfg["video_sets"])
-    elif excludealreadycropped:
-        for video in videos:
-            _, ext = os.path.splitext(video)
-            s = video.replace(ext, f"_cropped{ext}")
-            if s in cfg["video_sets"]:
-                videos.remove(video)
-    if not videos:
-        return
-
-    if (
-        "video_sets_original" not in cfg.keys() and updatevideoentries
-    ):  # this dict is kept for storing links to original full-sized videos
-        cfg["video_sets_original"] = {}
-
-    for video in videos:
-        vidpath, vidname, videotype = _robust_path_split(video)
-        folder = os.path.join(project_path, "labeled-data", vidname)
-        if userfeedback:
-            print("Do you want to crop frames for folder: ", folder, "?")
-            askuser = input("(yes/no):")
-        else:
-            askuser = "y"
-        if askuser == "y" or askuser == "yes" or askuser == "Y" or askuser == "Yes":
-            new_vidname = vidname + "_cropped"
-            new_folder = os.path.join(project_path, "labeled-data", new_vidname)
-            auxiliaryfunctions.attempttomakefolder(new_folder)
-
-            AnnotationData = []
-            pd_index = []
-
-            fn = os.path.join(folder, f"CollectedData_{cfg['scorer']}.h5")
-            df = pd.read_hdf(fn)
-            data = df.values.reshape((df.shape[0], -1, 2))
-            sep = "/" if "/" in df.index[0] else "\\"
-            if sep != os.path.sep:
-                df.index = df.index.str.replace(sep, os.path.sep)
-            video_new = sep.join((vidpath, new_vidname + videotype))
-            if video_new in cfg["video_sets"]:
-                _, w, _, h = map(int, cfg["video_sets"][video_new]["crop"].split(","))
-                temp_size = (h, w)
-            else:
-                temp_size = size
-            images = project_path + os.path.sep + df.index
-            # Avoid cropping already cropped images
-            cropped_images = auxiliaryfunctions.grab_files_in_folder(new_folder, "png")
-            cropped_names = set(map(lambda x: x.split("c")[0], cropped_images))
-            imnames = [
-                im for im in images.to_list() if Path(im).stem not in cropped_names
-            ]
-            if not imnames:
-                continue
-            ic = io.imread_collection(imnames)
-            for i in trange(len(ic)):
-                frame = ic[i]
-                h, w = np.shape(frame)[:2]                
-                                
-
-                imagename = os.path.relpath(ic.files[i], project_path)
-                ind = np.flatnonzero(df.index == imagename)[0]
-                cropindex = 0
-                attempts = -1
-                while cropindex < numcrops:
-                    
-                    dd = np.array(data[ind].copy(), dtype=float)
-
-                    if temp_size[0] >= h or temp_size[1] >= w:
-                        # initialize a all zero image container with target crop size
-
-                        padded_img = np.zeros((temp_size[0],temp_size[1],3),dtype=np.uint8)
-
-                        # new upper left, border protection
-                        # be careful crop side can be smaller than one side of the image
-                        
-                        y0, x0 = (
-                            np.random.randint(max(temp_size[0]-h,1)),
-                            np.random.randint(max(temp_size[1]-w,1)),
-                        )
-                        # new bottom right, border protection
-                        # avoid to exceed the container's border
-                        
-                        y1 = min(y0 + h, temp_size[0])
-                        x1 = min(x0 + w, temp_size[1])
-
-                        # fill original image to the container image
-                        # use safe upper left and bottom right to crop original image and fill it to the container
-                        padded_img[y0:y1,x0:x1,:] = frame[:y1-y0,:x1-x0,:3]
-
-                        # all keypoints are shifted by +x0 and +y0
-                        # possibly out of border again
-                        dd += [x0,y0]
-
-                        # some keypoints are out of the borders
-                        with np.errstate(invalid="ignore"):
-                            within = np.all((dd >= [x0, y0]) & (dd < [x1, y1]), axis=1)
-
-                        if cropdata:
-                            dd[~within] = np.nan                        
-                        attempts += 1
-                        if within.any() or attempts > 10:                        
-                            newimname = str(
-                                    Path(imagename).stem
-                                    + "c"
-                                    + str(cropindex).zfill(indexlength)
-                                    + ".png"
-                                    )
-                            cropppedimgname = os.path.join(new_folder, newimname)
-                            # save the padded img
-                            io.imsave(cropppedimgname, padded_img)
-                            cropindex += 1
-                            pd_index.append(
-                                os.path.join("labeled-data", new_vidname, newimname)
-                            )
-                            AnnotationData.append(dd.flatten())
-
-                        
-                    else:
-                        y0, x0 = (
-                            np.random.randint(h - temp_size[0]),
-                            np.random.randint(w - temp_size[1]),
-                        )
-                        y1 = y0 + temp_size[0]
-                        x1 = x0 + temp_size[1]
-                        
-                        with np.errstate(invalid="ignore"):
-                            within = np.all((dd >= [x0, y0]) & (dd < [x1, y1]), axis=1)
-                        if cropdata:
-                            dd[within] -= [x0, y0]
-                            dd[~within] = np.nan
-                        attempts += 1
-                        if within.any() or attempts > 10:
-                            newimname = str(
-                                Path(imagename).stem
-                                + "c"
-                                + str(cropindex).zfill(indexlength)
-                                + ".png"
-                                )
-                            cropppedimgname = os.path.join(new_folder, newimname)
-                            io.imsave(cropppedimgname, frame[y0:y1, x0:x1])
-                            cropindex += 1
-                            pd_index.append(
-                                os.path.join("labeled-data", new_vidname, newimname)
-                            )
-                            AnnotationData.append(dd.flatten())
-
-            if cropdata:
-                df = pd.DataFrame(AnnotationData, index=pd_index, columns=df.columns)
-                fn_new = fn.replace(folder, new_folder)
-                try:
-                    df_old = pd.read_hdf(fn_new)
-                    df = pd.concat((df_old, df))
-                except FileNotFoundError:
-                    pass
-                df.to_hdf(fn_new, key="df_with_missing", mode="w")
-                df.to_csv(fn_new.replace(".h5", ".csv"))
-
-            if updatevideoentries and cropdata:
-                # moving old entry to _original, dropping it from video_set and update crop parameters
-                video_orig = sep.join((vidpath, vidname + videotype))
-                video_new = sep.join((vidpath, new_vidname + videotype))
-                if video_orig not in cfg["video_sets_original"]:
-                    cfg["video_sets_original"][video_orig] = cfg["video_sets"][
-                        video_orig
-                    ]
-                    cfg["video_sets"].pop(video_orig)
-                    cfg["video_sets"][video_new] = {
-                        "crop": ", ".join(map(str, [0, temp_size[1], 0, temp_size[0]]))
-                    }
-                elif video_new not in cfg["video_sets"]:
-                    cfg["video_sets"][video_new] = {
-                        "crop": ", ".join(map(str, [0, temp_size[1], 0, temp_size[0]]))
-                    }
-
-    cfg["croppedtraining"] = True
-    auxiliaryfunctions.write_config(config, cfg)
-
-
 def check_labels(
     config,
     Labels=["+", ".", "x"],
@@ -647,8 +415,7 @@ def merge_annotateddatasets(cfg, trainingsetfolder_full, windows2linux):
         except FileNotFoundError:
             print(
                 file_path,
-                " not found (perhaps not annotated). If training on cropped data, "
-                "make sure to call `cropimagesandlabels` prior to creating the dataset.",
+                " not found (perhaps not annotated)."
             )
 
     if not len(AnnotationData):
@@ -732,18 +499,36 @@ def SplitTrials(
         test_indices = shuffle[int(train_size):]
         train_indices = shuffle[:int(train_size)]
         if enforce_train_fraction and not train_size.is_integer():
-            # Determine the index length required to guarantee
-            # the train–test ratio is exactly the desired one.
-            min_length_req = int(100 / math.gcd(100, int(round(100 * train_fraction))))
-            length_req = math.ceil(index_len / min_length_req) * min_length_req
-            n_train = int(round(length_req * train_fraction))
-            n_test = length_req - n_train
-            # Pad indices so lengths agree
-            train_indices = np.append(train_indices, [-1] * (n_train - len(train_indices)))
-            test_indices = np.append(test_indices, [-1] * (n_test - len(test_indices)))
+            train_indices, test_indices = pad_train_test_indices(
+                train_indices, test_indices, train_fraction,
+            )
         return train_indices, test_indices
 
 
+def pad_train_test_indices(train_inds, test_inds, train_fraction):
+    n_train_inds = len(train_inds)
+    n_test_inds = len(test_inds)
+    index_len = n_train_inds + n_test_inds
+    if n_train_inds / index_len == train_fraction:
+        return
+
+    # Determine the index length required to guarantee
+    # the train–test ratio is exactly the desired one.
+    min_length_req = int(100 / math.gcd(100, int(round(100 * train_fraction))))
+    min_n_train = int(round(min_length_req * train_fraction))
+    min_n_test = min_length_req - min_n_train
+    mult = max(
+        math.ceil(n_train_inds / min_n_train),
+        math.ceil(n_test_inds / min_n_test),
+    )
+    n_train = mult * min_n_train
+    n_test = mult * min_n_test
+    # Pad indices so lengths agree
+    train_inds = np.append(train_inds, [-1] * (n_train - n_train_inds))
+    test_inds = np.append(test_inds, [-1] * (n_test - n_test_inds))
+    return train_inds, test_inds
+
+
 def mergeandsplit(config, trainindex=0, uniform=True, windows2linux=False):
     """
     This function allows additional control over "create_training_dataset".
diff --git a/deeplabcut/gui/create_training_dataset.py b/deeplabcut/gui/create_training_dataset.py
index e584d9bb15..27ab7bf027 100644
--- a/deeplabcut/gui/create_training_dataset.py
+++ b/deeplabcut/gui/create_training_dataset.py
@@ -136,38 +136,6 @@ def __init__(self, parent, gui_size, cfg):
         )
         self.userfeedback.SetSelection(1)
 
-        if config_file.get("multianimalproject", False):
-
-            self.cropandlabel = wx.RadioBox(
-                self,
-                label="Crop and Label Data (Yes is required, set crop values)",
-                choices=["Yes", "No"],
-                majorDimension=1,
-                style=wx.RA_SPECIFY_COLS,
-            )
-            self.cropandlabel.Bind(wx.EVT_RADIOBOX, self.input_crop_size)
-            self.cropandlabel.SetSelection(0)
-            self.crop_text = wx.StaticBox(
-                self, label="Crop settings (set to smaller than your input images)"
-            )
-            self.crop_sizer = wx.StaticBoxSizer(self.crop_text, wx.VERTICAL)
-            self.crop_widgets = []
-            for name, val in [
-                ("# of crops", "10"),
-                ("height", "400"),
-                ("width", "400"),
-            ]:
-                temp_sizer = wx.BoxSizer(wx.HORIZONTAL)
-                label = wx.StaticText(self, label=name)
-                text = wx.TextCtrl(self, value=val)
-                self.crop_widgets.append([label, text])
-                temp_sizer.Add(label, wx.EXPAND | wx.TOP | wx.BOTTOM, 10)
-                temp_sizer.Add(text, wx.EXPAND | wx.TOP | wx.BOTTOM, 10)
-                self.crop_sizer.Add(temp_sizer)
-            self.crop_sizer.ShowItems(True)
-            self.hbox3.Add(self.cropandlabel, 10, wx.EXPAND | wx.TOP | wx.BOTTOM, 5)
-            self.hbox3.Add(self.crop_sizer, 10, wx.EXPAND | wx.TOP | wx.BOTTOM, 5)
-
         self.hbox2.Add(shuffle_text_boxsizer, 10, wx.EXPAND | wx.TOP | wx.BOTTOM, 5)
         self.hbox2.Add(trainingindex_boxsizer, 10, wx.EXPAND | wx.TOP | wx.BOTTOM, 5)
 
@@ -291,13 +259,6 @@ def __init__(self, parent, gui_size, cfg):
         self.sizer.Fit(self)
         self.Layout()
 
-    def input_crop_size(self, event):
-        if self.cropandlabel.GetStringSelection() == "No":
-            self.crop_sizer.ShowItems(False)
-        else:
-            self.crop_sizer.ShowItems(True)
-        self.SetSizer(self.sizer)
-
     def on_focus(self, event):
         pass
 
@@ -368,15 +329,6 @@ def create_training_dataset(self, event):
             userfeedback = False
 
         if config_file.get("multianimalproject", False):
-            if self.cropandlabel.GetStringSelection() == "Yes":
-                n_crops, height, width = [
-                    int(text.GetValue()) for _, text in self.crop_widgets
-                ]
-                deeplabcut.cropimagesandlabels(
-                    self.config, n_crops, (height, width), userfeedback
-                )
-            else:
-                random = False
             deeplabcut.create_multianimaltraining_dataset(
                 self.config,
                 num_shuffles,
diff --git a/deeplabcut/pose_cfg.yaml b/deeplabcut/pose_cfg.yaml
index 66df9dc014..3d4229f29e 100644
--- a/deeplabcut/pose_cfg.yaml
+++ b/deeplabcut/pose_cfg.yaml
@@ -26,13 +26,32 @@ global_scale: 0.8
 dataset_type: imgaug
 batch_size: 1
 
+# Probability with which the augmenters will be applied to input images
+# Note some augmentations have their own probability (e.g. claheratio/rotratio/...)
+apply_prob: 0.5
+
+# Resize images prior to augmentation
+pre_resize: []  # Specify [width, height] if pre-resizing is desired
+
+# Smart, on-the-fly image cropping, replacing deeplabcut.cropimagesandlabels
+crop_size: [400, 400]  # width, height
+max_shift: 0.4  # Maximum relative shift of the position of the crop center
+crop_sampling: hybrid  # Sample crop centers either uniformly over the image or based on keypoint neighbor density, at random
+
+# Other crop_sampling variants:
+#- uniform -- spatially uniform sampling of crops (over the image)
+#- keypoints -- keypoint based sampling of crops (over the image)
+#- density -- keypoint based sampling of crops biasing towards regions with more keypoints (over the image)
+#- hybrid -- 50% density and 50% uniform
+
+
 #Data loaders, i.e. with additional data augmentation options (as of 2.0.9+):
 #default with be with no extra dataloaders. Other options: 'tensorpack, deterministic'
 #types of datasets, see factory: deeplabcut/pose_estimation_tensorflow/dataset/factory.py
 #For deterministic, see https://github.com/AlexEMG/DeepLabCut/pull/324
 #For tensorpack, see https://github.com/AlexEMG/DeepLabCut/pull/409
 
-#what is the fraction of training samples with cropping? (used for scalecrop, )
+#what is the fraction of training samples with cropping? (used for scalecrop)
 cropratio: 0.4
 # see below for cropping variables for tensorpack and scalecrop (these need to be set in pose_cfg.yaml per model)
 # for imagaug strengh is modulated by crop_by
diff --git a/deeplabcut/pose_estimation_tensorflow/core/evaluate_multianimal.py b/deeplabcut/pose_estimation_tensorflow/core/evaluate_multianimal.py
index b7afc50eca..7ba201ec50 100644
--- a/deeplabcut/pose_estimation_tensorflow/core/evaluate_multianimal.py
+++ b/deeplabcut/pose_estimation_tensorflow/core/evaluate_multianimal.py
@@ -10,6 +10,7 @@
 
 
 import os
+import pickle
 from pathlib import Path
 
 import numpy as np
@@ -202,7 +203,8 @@ def evaluate_multianimal_full(
 
             # TODO: IMPLEMENT for different batch sizes?
             dlc_cfg["batch_size"] = 1  # due to differently sized images!!!
-
+            # Ignore best edges possibly defined during a prior evaluation
+            _ = dlc_cfg.pop("paf_best", None)
             joints = dlc_cfg["all_joints_names"]
 
             # Create folder structure to store results.
@@ -379,13 +381,16 @@ def evaluate_multianimal_full(
                                 if inds_gt.size and xy.size:
                                     # Pick the predictions closest to ground truth,
                                     # rather than the ones the model has most confident in
-                                    d = cdist(xy_gt.iloc[inds_gt], xy)
-                                    rows, cols = linear_sum_assignment(d)
-                                    min_dists = d[rows, cols]
+                                    xy_gt_values = xy_gt.iloc[inds_gt].values
+                                    neighbors = _find_closest_neighbors(xy_gt_values, xy, k=3)
+                                    found = neighbors != -1
+                                    min_dists = np.linalg.norm(
+                                        xy_gt_values[found] - xy[neighbors[found]], axis=1,
+                                    )
                                     inds = np.flatnonzero(all_bpts == bpt)
-                                    sl = imageindex, inds[inds_gt[rows]]
+                                    sl = imageindex, inds[inds_gt[found]]
                                     dist[sl] = min_dists
-                                    conf[sl] = probs_pred[n_joint][cols].squeeze()
+                                    conf[sl] = probs_pred[n_joint][neighbors[found]].squeeze()
 
                             if plotting:
                                 gt = temp_xy.values.reshape(
@@ -501,9 +506,7 @@ def evaluate_multianimal_full(
                             "nms radius": dlc_cfg["nmsradius"],
                             "minimal confidence": dlc_cfg["minconfidence"],
                             "PAFgraph": dlc_cfg["partaffinityfield_graph"],
-                            "PAFinds": dlc_cfg.get(
-                                "paf_best",
-                                np.arange(len(dlc_cfg["partaffinityfield_graph"])),
+                            "PAFinds": np.arange(len(dlc_cfg["partaffinityfield_graph"]),
                             ),
                             "all_joints": [
                                 [i] for i in range(len(dlc_cfg["all_joints"]))
@@ -535,23 +538,31 @@ def evaluate_multianimal_full(
 
                     # Skip data-driven skeleton selection unless
                     # the model was trained on the full graph.
-                    max_n_edges = dlc_cfg["num_joints"] * (dlc_cfg["num_joints"] - 1) // 2
-                    if len(dlc_cfg["partaffinityfield_graph"]) == max_n_edges:
-                        uncropped_data_path = data_path.replace(
-                            ".pickle", "_uncropped.pickle"
-                        )
-                        if not os.path.isfile(uncropped_data_path):
-                            print("Selecting best skeleton...")
-                            crossvalutils._rebuild_uncropped_in(evaluationfolder)
-                            _ = crossvalutils.cross_validate_paf_graphs(
-                                config,
-                                str(path_test_config).replace("pose_", "inference_"),
-                                uncropped_data_path,
-                                uncropped_data_path.replace("_full_", "_meta_"),
-                            )
+                    n_multibpts = len(cfg["multianimalbodyparts"])
+                    max_n_edges = n_multibpts * (n_multibpts - 1) // 2
+                    n_edges = len(dlc_cfg["partaffinityfield_graph"])
+                    if n_edges == max_n_edges:
+                        print("Selecting best skeleton...")
+                        n_graphs = 10
+                        paf_inds = None
+                    else:
+                        n_graphs = 1
+                        paf_inds = [list(range(n_edges))]
+                    results, paf_scores = crossvalutils.cross_validate_paf_graphs(
+                        config,
+                        str(path_test_config).replace("pose_", "inference_"),
+                        data_path,
+                        data_path.replace("_full.", "_meta."),
+                        n_graphs=n_graphs,
+                        paf_inds=paf_inds,
+                    )
+                    df = results[1].copy()
+                    df.loc(axis=0)[('mAP', 'mean')] = [d['mAP'] for d in results[2]]
+                    df.loc(axis=0)[('mAR', 'mean')] = [d['mAR'] for d in results[2]]
+                    with open(data_path.replace("_full.", "_map."), "wb") as file:
+                        pickle.dump((df, paf_scores), file)
 
                 if len(final_result) > 0:  # Only append if results were calculated
                     make_results_file(final_result, evaluationfolder, DLCscorer)
 
-    # returning to intial folder
     os.chdir(str(start_path))
diff --git a/deeplabcut/pose_estimation_tensorflow/core/predict_multianimal.py b/deeplabcut/pose_estimation_tensorflow/core/predict_multianimal.py
index c230ef6a8d..93449e51f6 100644
--- a/deeplabcut/pose_estimation_tensorflow/core/predict_multianimal.py
+++ b/deeplabcut/pose_estimation_tensorflow/core/predict_multianimal.py
@@ -45,9 +45,7 @@ def AssociationCosts(
     ny, nx, nlimbs = np.shape(partaffinitymaps)
     graph = cfg["partaffinityfield_graph"]
     limbs = cfg.get("paf_best", np.arange(len(graph)))
-    if len(graph) != len(limbs):
-        limbs = np.arange(len(graph))
-
+    graph = [graph[l] for l in limbs]
     for l, (bp1, bp2) in zip(limbs, graph):
         # get coordinates for bp1 and bp2
         C1 = coordinates[bp1]
diff --git a/deeplabcut/pose_estimation_tensorflow/datasets/augmentation.py b/deeplabcut/pose_estimation_tensorflow/datasets/augmentation.py
new file mode 100644
index 0000000000..9e64f260b8
--- /dev/null
+++ b/deeplabcut/pose_estimation_tensorflow/datasets/augmentation.py
@@ -0,0 +1,94 @@
+import imgaug.augmenters as iaa
+import numpy as np
+from scipy.spatial.distance import pdist, squareform
+
+
+class KeypointAwareCropToFixedSize(iaa.CropToFixedSize):
+    def __init__(
+        self,
+        width,
+        height,
+        max_shift=0.4,
+        crop_sampling="hybrid",
+    ):
+        """
+        Parameters
+        ----------
+        width : int
+            Crop images down to this maximum width.
+
+        height : int
+            Crop images down to this maximum height.
+
+        max_shift : float, optional (default=0.25)
+            Maximum allowed shift of the cropping center position
+            as a fraction of the crop size.
+
+        crop_sampling : str, optional (default="hybrid")
+            Crop centers sampling method. Must be either:
+            "uniform" (randomly over the image),
+            "keypoints" (randomly over the annotated keypoints),
+            "density" (weighing preferentially dense regions of keypoints),
+            or "hybrid" (alternating randomly between "uniform" and "density").
+        """
+        super(KeypointAwareCropToFixedSize, self).__init__(
+            width, height, name="kptscrop",
+        )
+        # Clamp to 40% of crop size to ensure that at least
+        # the center keypoint remains visible after the offset is applied.
+        self.max_shift = max(0., min(max_shift, 0.4))
+        if crop_sampling not in ("uniform", "keypoints", "density", "hybrid"):
+            raise ValueError(f"Invalid sampling {crop_sampling}. Must be "
+                             f"either 'uniform', 'keypoints', 'density', or 'hybrid.")
+        self.crop_sampling = crop_sampling
+
+    @staticmethod
+    def calc_n_neighbors(xy, radius):
+        d = pdist(xy, "sqeuclidean")
+        mat = squareform(d <= radius * radius, checks=False)
+        return np.sum(mat, axis=0)
+
+    def _draw_samples(self, batch, random_state):
+        n_samples = batch.nb_rows
+        offsets = np.empty((n_samples, 2), dtype=np.float32)
+        rngs = random_state.duplicate(2)
+        shift_x = self.max_shift * self.size[0] * rngs[0].uniform(-1, 1, n_samples)
+        shift_y = self.max_shift * self.size[1] * rngs[1].uniform(-1, 1, n_samples)
+        sampling = self.crop_sampling
+        for n in range(batch.nb_rows):
+            if self.crop_sampling == "hybrid":
+                sampling = random_state.choice(["uniform", "density"])
+            if sampling == "uniform":
+                center = random_state.uniform(size=2)
+            else:
+                h, w = batch.images[n].shape[:2]
+                kpts = batch.keypoints[n].to_xy_array()
+                n_kpts = kpts.shape[0]
+                inds = np.arange(n_kpts)
+                if sampling == "density":
+                    # Points located close to one another are sampled preferentially
+                    # in order to augment crowded regions.
+                    radius = 0.1 * min(h, w)
+                    n_neighbors = self.calc_n_neighbors(kpts, radius)
+                    # Include keypoints in the count to avoid null probabilities
+                    n_neighbors += 1
+                    p = n_neighbors / n_neighbors.sum()
+                else:
+                    p = np.ones_like(inds) / n_kpts
+                center = kpts[random_state.choice(inds, p=p)]
+                # Shift the crop center in both dimensions by random amounts
+                # and normalize to the original image dimensions.
+                center[0] += shift_x[n]
+                center[0] /= w
+                center[1] += shift_y[n]
+                center[1] /= h
+            offsets[n] = center
+        offsets = np.clip(offsets, 0, 1)
+        return [self.size] * n_samples, offsets[:, 0], offsets[:, 1]
+
+
+def update_crop_size(pipeline, width, height):
+    aug = pipeline.find_augmenters_by_name("kptscrop")
+    if not aug:
+        return
+    aug[0].size = width, height
diff --git a/deeplabcut/pose_estimation_tensorflow/datasets/pose_multianimal_imgaug.py b/deeplabcut/pose_estimation_tensorflow/datasets/pose_multianimal_imgaug.py
index eb1b59ffcf..f230ae4219 100644
--- a/deeplabcut/pose_estimation_tensorflow/datasets/pose_multianimal_imgaug.py
+++ b/deeplabcut/pose_estimation_tensorflow/datasets/pose_multianimal_imgaug.py
@@ -17,10 +17,12 @@
 import imgaug.augmenters as iaa
 import numpy as np
 from imgaug.augmentables import Keypoint, KeypointsOnImage
-from .factory import PoseDatasetFactory
-from .pose_base import BasePoseDataset
-from .utils import DataItem, Batch
+from deeplabcut.pose_estimation_tensorflow.datasets import augmentation
+from deeplabcut.pose_estimation_tensorflow.datasets.factory import PoseDatasetFactory
+from deeplabcut.pose_estimation_tensorflow.datasets.pose_base import BasePoseDataset
+from deeplabcut.pose_estimation_tensorflow.datasets.utils import DataItem, Batch
 from deeplabcut.utils.auxfun_videos import imread
+from math import sqrt
 
 
 @PoseDatasetFactory.register("multi-animal-imgaug")
@@ -31,6 +33,13 @@ def __init__(self, cfg):
         self.num_images = len(self.data)
         self.batch_size = cfg["batch_size"]
         print("Batch Size is %d" % self.batch_size)
+        self.pipeline = self.build_augmentation_pipeline(
+            apply_prob=cfg.get('apply_prob', 0.5),
+        )
+
+    @property
+    def default_crop_size(self):
+        return self.cfg.get("crop_size", (400, 400))  # width, height
 
     def load_dataset(self):
         cfg = self.cfg
@@ -71,11 +80,26 @@ def load_dataset(self):
         self.has_gt = has_gt
         return data
 
-    def build_augmentation_pipeline(self, height=None, width=None, apply_prob=0.5):
+    def build_augmentation_pipeline(self, apply_prob=0.5):
+        cfg = self.cfg
+
         sometimes = lambda aug: iaa.Sometimes(apply_prob, aug)
         pipeline = iaa.Sequential(random_order=False)
-        
-        cfg = self.cfg
+
+        pre_resize = cfg.get("pre_resize")
+        if pre_resize:
+            width, height = pre_resize
+            pipeline.add(iaa.Resize({"height": height, "width": width}))
+
+        # Add smart, keypoint-aware image cropping
+        w, h = self.default_crop_size
+        pipeline.add(iaa.PadToFixedSize(w, h))
+        pipeline.add(
+            augmentation.KeypointAwareCropToFixedSize(
+                w, h, cfg.get("max_shift", 0.4), cfg.get("crop_sampling", "hybrid")
+            )
+        )
+
         if cfg.get("fliplr", False):
             opt = cfg.get("fliplr", False)
             if type(opt) == int:
@@ -194,15 +218,6 @@ def get_aug_param(cfg_value):
             opt = get_aug_param(cfg_cnv["edge"])
             pipeline.add(iaa.Sometimes(cfg_cnv["edgeratio"], iaa.EdgeDetect(**opt)))
 
-        
-        if height is not None and width is not None:
-            pipeline.add(
-                iaa.Sometimes(
-                    cfg.get("cropratio", 0.4),
-                    iaa.CropAndPad(percent=(-0.3, 0.1), keep_size=False),
-                )
-            )
-            pipeline.add(iaa.Resize({"height": height, "width": width}))
         return pipeline
 
     def get_batch(self):
@@ -211,18 +226,6 @@ def get_batch(self):
         batch_joints = []
         joint_ids = []
         data_items = []
-
-        # Scale is sampled only once (per batch) to transform all of the images into same size.
-        scale = self.get_scale()
-        while True:
-            idx = np.random.choice(self.num_images)
-            scale = self.get_scale()
-            size = self.data[idx].im_size
-            target_size = np.ceil(size[1:3] * scale).astype(int)
-            if self.is_valid_size(target_size):
-                break
-
-        stride = self.cfg["stride"]
         for i in range(self.batch_size):
             data_item = self.data[img_idx[i]]
 
@@ -243,34 +246,18 @@ def get_batch(self):
                 batch_joints.append(np.array(joint_points))
 
             batch_images.append(image)
-
-        sm_size = np.ceil(target_size / (stride * self.cfg.get("smfactor", 2))).astype(
-            int
-        ) * self.cfg.get("smfactor", 2)
-
-        if stride == 2:
-            sm_size = np.ceil(target_size / 16).astype(int)
-            sm_size *= 8
-
-        # assert len(batch_images) == self.batch_size
-        return batch_images, joint_ids, batch_joints, data_items, sm_size, target_size
+        return batch_images, joint_ids, batch_joints, data_items
 
     def get_targetmaps_update(
-        self, joint_ids, joints, data_items, sm_size, target_size
+        self, joint_ids, joints, data_items, sm_size, scale,
     ):
-        part_score_targets, part_score_weights, locref_targets, locref_masks = (
-            [],
-            [],
-            [],
-            [],
-        )
-        partaffinityfield_targets, partaffinityfield_masks = [], []
+        part_score_targets = []
+        part_score_weights = []
+        locref_targets = []
+        locref_masks = []
+        partaffinityfield_targets = []
+        partaffinityfield_masks = []
         for i in range(len(data_items)):
-            # Approximating the scale
-            scale = min(
-                target_size[0] / data_items[i].im_size[1],
-                target_size[1] / data_items[i].im_size[2],
-            )
             if self.cfg.get("scmap_type", None) == "gaussian":
                 assert 0 == 1  # not implemented for pafs!
                 (
@@ -290,7 +277,7 @@ def get_targetmaps_update(
                     partaffinityfield_target,
                     partaffinityfield_mask,
                 ) = self.compute_target_part_scoremap_numpy(
-                    joint_ids[i], [joints[i]], data_items[i], sm_size, scale
+                    joint_ids[i], joints[i], data_items[i], sm_size, scale
                 )
 
             part_score_targets.append(part_score_target)
@@ -309,6 +296,20 @@ def get_targetmaps_update(
             Batch.pairwise_mask: partaffinityfield_masks,
         }
 
+    def calc_target_and_scoremap_sizes(self):
+        target_size = np.asarray(self.default_crop_size) * self.get_scale()
+        target_size = np.ceil(target_size).astype(int)
+        if not self.is_valid_size(target_size):
+            target_size = self.default_crop_size
+        stride = self.cfg["stride"]
+        sm_size = np.ceil(target_size / (stride * self.cfg.get("smfactor", 2))).astype(
+            int
+        ) * self.cfg.get("smfactor", 2)
+        if stride == 2:
+            sm_size = np.ceil(target_size / 16).astype(int)
+            sm_size *= 8
+        return target_size, sm_size
+
     def next_batch(self, plotting=False):
         while True:
             (
@@ -316,23 +317,43 @@ def next_batch(self, plotting=False):
                 joint_ids,
                 batch_joints,
                 data_items,
-                sm_size,
-                target_size,
             ) = self.get_batch()
 
-            pipeline = self.build_augmentation_pipeline(
-                height=target_size[0], width=target_size[1], apply_prob=0.5
-            )
-
-            batch_images, batch_joints = pipeline(
+            # Scale is sampled only once (per batch) to transform all of the images into same size.
+            target_size, sm_size = self.calc_target_and_scoremap_sizes()
+            scale = np.mean(target_size / self.default_crop_size)
+            augmentation.update_crop_size(self.pipeline, *target_size)
+            batch_images, batch_joints = self.pipeline(
                 images=batch_images, keypoints=batch_joints
             )
+            batch_images = np.asarray(batch_images)
+            image_shape = batch_images.shape[1:3]
+            # Discard keypoints whose coordinates lie outside the cropped image
+            batch_joints_valid = []
+            joint_ids_valid = []
+            for joints, ids in zip(batch_joints, joint_ids):
+                inside = np.logical_and.reduce(
+                    (
+                        joints[:, 0] < image_shape[1],
+                        joints[:, 0] > 0,
+                        joints[:, 1] < image_shape[0],
+                        joints[:, 1] > 0,
+                    )
+                )
+                batch_joints_valid.append(joints[inside])
+                temp = []
+                start = 0
+                for array in ids:
+                    end = start + array.size
+                    temp.append(array[inside[start:end]])
+                    start = end
+                joint_ids_valid.append(temp)
 
             # If you would like to check the augmented images, script for saving
             # the images with joints on:
             if plotting:
                 for i in range(self.batch_size):
-                    joints = batch_joints[i]
+                    joints = batch_joints_valid[i]
                     kps = KeypointsOnImage(
                         [Keypoint(x=joint[0], y=joint[1]) for joint in joints],
                         shape=batch_images[i].shape,
@@ -343,18 +364,18 @@ def next_batch(self, plotting=False):
                         os.path.join(self.cfg["project_path"], str(i) + ".png"), im
                     )
 
-            image_shape = np.array(batch_images).shape[1:3]
-            batch = {Batch.inputs: np.array(batch_images).astype(np.float64)}
+            batch = {Batch.inputs: batch_images.astype(np.float64)}
             if self.has_gt:
                 targetmaps = self.get_targetmaps_update(
-                    joint_ids, batch_joints, data_items, sm_size, image_shape
+                    joint_ids_valid,
+                    batch_joints_valid,
+                    data_items,
+                    (sm_size[1], sm_size[0]),
+                    scale,
                 )
                 batch.update(targetmaps)
 
-            # if returndata:
-            #        return batch_images,batch_joints,targetmaps
-
-            batch = {key: np.array(data) for (key, data) in batch.items()}
+            batch = {key: np.asarray(data) for (key, data) in batch.items()}
             batch[Batch.data_item] = data_items
             return batch
 
@@ -376,9 +397,8 @@ def get_scale(self):
         return scale
 
     def is_valid_size(self, target_size):
-        im_width = target_size[1]
-        im_height = target_size[0]
-        min_input_size = 100
+        im_width, im_height = target_size
+        min_input_size = self.cfg.get("min_input_size", 100)
         if im_height < min_input_size or im_width < min_input_size:
             return False
         if hasattr(self.cfg, "max_input_size"):
@@ -387,7 +407,7 @@ def is_valid_size(self, target_size):
                 return False
         return True
 
-    def compute_scmap_weights(self, scmap_shape, joint_id, data_item):
+    def compute_scmap_weights(self, scmap_shape, joint_id):
         cfg = self.cfg
         if cfg["weigh_only_present_joints"]:
             weights = np.zeros(scmap_shape)
@@ -403,20 +423,19 @@ def compute_target_part_scoremap_numpy(
         self, joint_id, coords, data_item, size, scale
     ):
         stride = self.cfg["stride"]
+        half_stride = stride // 2
         dist_thresh = float(self.cfg["pos_dist_thresh"] * scale)
         num_idchannel = self.cfg.get("num_idchannel", 0)
 
         num_joints = self.cfg["num_joints"]
-        half_stride = stride / 2
-
-        scmap = np.zeros(np.concatenate([size, np.array([num_joints + num_idchannel])]))
-        locref_size = np.concatenate([size, np.array([num_joints * 2])])
 
+        scmap = np.zeros((*size, num_joints + num_idchannel))
+        locref_size = *size, num_joints * 2
         locref_map = np.zeros(locref_size)
         locref_scale = 1.0 / self.cfg["locref_stdev"]
         dist_thresh_sq = dist_thresh ** 2
 
-        partaffinityfield_shape = np.concatenate([size, np.array([self.cfg["num_limbs"] * 2])])
+        partaffinityfield_shape = *size, self.cfg["num_limbs"] * 2
         partaffinityfield_map = np.zeros(partaffinityfield_shape)
         if self.cfg["weigh_only_present_joints"]:
             partaffinityfield_mask = np.zeros(partaffinityfield_shape)
@@ -425,38 +444,33 @@ def compute_target_part_scoremap_numpy(
             partaffinityfield_mask = np.ones(partaffinityfield_shape)
             locref_mask = np.ones(locref_size)
 
-        width = size[1]
-        height = size[0]
+        height, width = size
         grid = np.mgrid[:height, :width].transpose((1, 2, 0))
-
-        # the animal id plays no role for scoremap + locref!
-        # so let's just loop over all bpts.
-        for k, j_id in enumerate(np.concatenate(joint_id)):
-            joint_pt = coords[0][k, :]
-            j_x = np.asscalar(joint_pt[0])
-            j_x_sm = round((j_x - half_stride) / stride)
-            j_y = np.asscalar(joint_pt[1])
-            j_y_sm = round((j_y - half_stride) / stride)
-
-            min_x = round(max(j_x_sm - dist_thresh - 1, 0))
-            max_x = round(min(j_x_sm + dist_thresh + 1, width - 1))
-            min_y = round(max(j_y_sm - dist_thresh - 1, 0))
-            max_y = round(min(j_y_sm + dist_thresh + 1, height - 1))
-            x = grid.copy()[:, :, 1]
-            y = grid.copy()[:, :, 0]
-            dx = j_x - x * stride - half_stride
-            dy = j_y - y * stride - half_stride
-            dist = dx ** 2 + dy ** 2
-            mask1 = dist <= dist_thresh_sq
-            mask2 = (x >= min_x) & (x <= max_x)
-            mask3 = (y >= min_y) & (y <= max_y)
-            mask = mask1 & mask2 & mask3
-            scmap[mask, j_id] = 1
+        xx = np.expand_dims(grid[..., 1], axis=2)
+        yy = np.expand_dims(grid[..., 0], axis=2)
+
+        # Produce score maps and location refinement fields
+        coords_sm = np.round((coords - half_stride) / stride).astype(int)
+        mins = np.round(np.maximum(coords_sm - dist_thresh - 1, 0)).astype(int)
+        maxs = np.round(
+            np.minimum(coords_sm + dist_thresh + 1, [width - 1, height - 1])
+        ).astype(int)
+        dx = coords[:, 0] - xx * stride - half_stride
+        dx_ = dx * locref_scale
+        dy = coords[:, 1] - yy * stride - half_stride
+        dy_ = dy * locref_scale
+        dist = dx ** 2 + dy ** 2
+        mask1 = dist <= dist_thresh_sq
+        mask2 = (xx >= mins[:, 0]) & (xx <= maxs[:, 0])
+        mask3 = (yy >= mins[:, 1]) & (yy <= maxs[:, 1])
+        mask = mask1 & mask2 & mask3
+        for n, ind in enumerate(np.concatenate(joint_id).tolist()):
+            mask_ = mask[..., n]
+            scmap[mask_, ind] = 1
             if self.cfg["weigh_only_present_joints"]:
-                locref_mask[mask, j_id * 2 + 0] = 1.0
-                locref_mask[mask, j_id * 2 + 1] = 1.0
-            locref_map[mask, j_id * 2 + 0] = (dx * locref_scale)[mask]
-            locref_map[mask, j_id * 2 + 1] = (dy * locref_scale)[mask]
+                locref_mask[mask_, [ind * 2 + 0, ind * 2 + 1]] = 1.0
+            locref_map[mask_, ind * 2 + 0] = dx_[mask_, n]
+            locref_map[mask_, ind * 2 + 1] = dy_[mask_, n]
 
         if num_idchannel > 0:
             coordinateoffset = 0
@@ -464,108 +478,71 @@ def compute_target_part_scoremap_numpy(
             idx = [i for i, id_ in enumerate(data_item.joints)
                    if id_ < num_idchannel]
             for person_id in idx:
-                joint_ids = joint_id[person_id].copy()
-                if len(joint_ids) > 1:
-                    for k, j_id in enumerate(joint_ids):
-                        joint_pt = coords[0][k + coordinateoffset, :]
-                        j_x = np.asscalar(joint_pt[0])
-                        j_x_sm = round((j_x - half_stride) / stride)
-                        j_y = np.asscalar(joint_pt[1])
-                        j_y_sm = round((j_y - half_stride) / stride)
-
-                        min_x = round(max(j_x_sm - dist_thresh - 1, 0))
-                        max_x = round(min(j_x_sm + dist_thresh + 1, width - 1))
-                        min_y = round(max(j_y_sm - dist_thresh - 1, 0))
-                        max_y = round(min(j_y_sm + dist_thresh + 1, height - 1))
-                        x = grid.copy()[:, :, 1]
-                        y = grid.copy()[:, :, 0]
-                        dx = j_x - x * stride - half_stride
-                        dy = j_y - y * stride - half_stride
-                        dist = dx ** 2 + dy ** 2
-                        mask1 = dist <= dist_thresh_sq
-                        mask2 = (x >= min_x) & (x <= max_x)
-                        mask3 = (y >= min_y) & (y <= max_y)
-                        mask = mask1 & mask2 & mask3
-                        scmap[mask, person_id + num_joints] = 1
-                coordinateoffset += len(joint_ids)
-
-        x = grid.copy()[:, :, 1]
-        y = grid.copy()[:, :, 0]
-
-        # if self.cfg.partaffinityfield_predict:
-        # print("hello",joint_id)
-        # print(np.concatenate(joint_id)) #this is all joint_ids for all individuals!
+                joint_ids = joint_id[person_id]
+                n_joints = joint_ids.size
+                if n_joints:
+                    inds = np.arange(n_joints) + coordinateoffset
+                    mask_ = mask[..., inds].sum(axis=2)
+                    scmap[mask_, person_id + num_joints] = 1
+                coordinateoffset += n_joints
+
         coordinateoffset = 0  # the offset based on
+        y, x = np.rollaxis(grid * stride + half_stride, 2)
         for person_id in range(len(joint_id)):
-            # for k, joint_ids in enumerate(joint_id[person_id]):
-            joint_ids = joint_id[person_id].copy()
-            if len(joint_ids) > 1:  # otherwise there cannot be a joint!
-                # CONSIDER SMARTER SEARCHES here... (i.e. calculate the bpts beforehand?)
-                for l in range(self.cfg["num_limbs"]):
-                    bp1, bp2 = self.cfg["partaffinityfield_graph"][l]
-                    I1 = np.where(np.array(joint_ids) == bp1)[0]
-                    I2 = np.where(np.array(joint_ids) == bp2)[0]
-                    if (len(I1) > 0) * (len(I2) > 0):
-                        indbp1 = np.asscalar(I1[0])
-                        indbp2 = np.asscalar(I2[0])
-                        j_x = np.asscalar(coords[0][indbp1 + coordinateoffset, 0])
-                        j_y = np.asscalar(coords[0][indbp1 + coordinateoffset, 1])
-
-                        linkedj_x = np.asscalar(coords[0][indbp2 + coordinateoffset, 0])
-                        linkedj_y = np.asscalar(coords[0][indbp2 + coordinateoffset, 1])
-
-                        dist = np.sqrt((linkedj_x - j_x) ** 2 + (linkedj_y - j_y) ** 2)
-                        if dist > 0:
-                            Dx = (linkedj_x - j_x) * 1.0 / dist  # x-axis UNIT VECTOR
-                            Dy = (linkedj_y - j_y) * 1.0 / dist
-
-                            d1 = [
-                                np.asscalar(Dx * j_x + Dy * j_y),
-                                np.asscalar(Dx * linkedj_x + Dy * linkedj_y),
-                            ]  # in-line with direct axis
-                            d1lowerboundary = min(d1)
-                            d1upperboundary = max(d1)
-                            d2mid = np.asscalar(
-                                j_y * Dx - j_x * Dy
-                            )  # orthogonal direction
-
-                            distance_along = Dx * (x * stride + half_stride) + Dy * (
-                                y * stride + half_stride
-                            )
-                            distance_across = (
+            joint_ids = joint_id[person_id].tolist()
+            if len(joint_ids) >= 2:  # there is a possible edge
+                for l, (bp1, bp2) in enumerate(self.cfg["partaffinityfield_graph"]):
+                    try:
+                        ind1 = joint_ids.index(bp1)
+                    except ValueError:
+                        continue
+                    try:
+                        ind2 = joint_ids.index(bp2)
+                    except ValueError:
+                        continue
+                    j_x, j_y = coords[ind1 + coordinateoffset]
+                    linkedj_x, linkedj_y = coords[ind2 + coordinateoffset]
+                    dist = sqrt((linkedj_x - j_x) ** 2 + (linkedj_y - j_y) ** 2)
+                    if dist > 0:
+                        Dx = (linkedj_x - j_x) / dist  # x-axis UNIT VECTOR
+                        Dy = (linkedj_y - j_y) / dist
+                        d1 = [
+                            Dx * j_x + Dy * j_y,
+                            Dx * linkedj_x + Dy * linkedj_y,
+                        ]  # in-line with direct axis
+                        d1lowerboundary = min(d1)
+                        d1upperboundary = max(d1)
+                        d2mid = j_y * Dx - j_x * Dy  # orthogonal direction
+
+                        distance_along = Dx * x + Dy * y
+                        distance_across = (
+                            (
                                 (
-                                    (
-                                        (y * stride + half_stride) * Dx
-                                        - (x * stride + half_stride) * Dy
-                                    )
-                                    - d2mid
+                                    y * Dx
+                                    - x * Dy
                                 )
-                                * 1.0
-                                / self.cfg["pafwidth"]
-                                * scale
+                                - d2mid
                             )
+                            * 1.0
+                            / self.cfg["pafwidth"]
+                            * scale
+                        )
 
-                            mask1 = (distance_along >= d1lowerboundary) & (
-                                distance_along <= d1upperboundary
-                            )
-                            mask2 = np.abs(distance_across) <= 1
-                            # mask3 = ((x >= 0) & (x <= width-1))
-                            # mask4 = ((y >= 0) & (y <= height-1))
-                            mask = mask1 & mask2  # &mask3 &mask4
-                            if self.cfg["weigh_only_present_joints"]:
-                                partaffinityfield_mask[mask, l * 2 + 0] = 1.0
-                                partaffinityfield_mask[mask, l * 2 + 1] = 1.0
-
-                            partaffinityfield_map[mask, l * 2 + 0] = (
-                                Dx * (1 - abs(distance_across))
-                            )[mask]
-                            partaffinityfield_map[mask, l * 2 + 1] = (
-                                Dy * (1 - abs(distance_across))
-                            )[mask]
+                        mask1 = (distance_along >= d1lowerboundary) & (
+                            distance_along <= d1upperboundary
+                        )
+                        distance_across_abs = np.abs(distance_across)
+                        mask2 = distance_across_abs <= 1
+                        mask = mask1 & mask2
+                        temp = 1 - distance_across_abs[mask]
+                        if self.cfg["weigh_only_present_joints"]:
+                            partaffinityfield_mask[mask, [l * 2 + 0, l * 2 + 1]] = 1.0
+                        partaffinityfield_map[mask, l * 2 + 0] = Dx * temp
+                        partaffinityfield_map[mask, l * 2 + 1] = Dy * temp
 
             coordinateoffset += len(joint_ids)  # keeping track of the blocks
 
-        weights = self.compute_scmap_weights(scmap.shape, joint_id, data_item)
+        weights = self.compute_scmap_weights(scmap.shape, joint_id)
         return (
             scmap,
             weights,
@@ -611,9 +588,9 @@ def gaussian_scmap(self, joint_id, coords, data_item, size, scale):
         # so let's just loop over all bpts.
         for k, j_id in enumerate(np.concatenate(joint_id)):
             joint_pt = coords[0][k, :]
-            j_x = np.asscalar(joint_pt[0])
+            j_x = joint_pt[0].item()
             j_x_sm = round((j_x - half_stride) / stride)
-            j_y = np.asscalar(joint_pt[1])
+            j_y = joint_pt[1].item()
             j_y_sm = round((j_y - half_stride) / stride)
 
             map_j = grid.copy()
@@ -643,13 +620,13 @@ def gaussian_scmap(self, joint_id, coords, data_item, size, scale):
                     I1 = np.where(np.array(joint_ids) == bp1)[0]
                     I2 = np.where(np.array(joint_ids) == bp2)[0]
                     if (len(I1) > 0) * (len(I2) > 0):
-                        indbp1 = np.asscalar(I1[0])
-                        indbp2 = np.asscalar(I2[0])
-                        j_x = np.asscalar(coords[0][indbp1 + coordinateoffset, 0])
-                        j_y = np.asscalar(coords[0][indbp1 + coordinateoffset, 1])
+                        indbp1 = I1[0].item()
+                        indbp2 = I2[0].item()
+                        j_x = (coords[0][indbp1 + coordinateoffset, 0]).item()
+                        j_y = (coords[0][indbp1 + coordinateoffset, 1]).item()
 
-                        linkedj_x = np.asscalar(coords[0][indbp2 + coordinateoffset, 0])
-                        linkedj_y = np.asscalar(coords[0][indbp2 + coordinateoffset, 1])
+                        linkedj_x = (coords[0][indbp2 + coordinateoffset, 0]).item()
+                        linkedj_y = (coords[0][indbp2 + coordinateoffset, 1]).item()
 
                         dist = np.sqrt((linkedj_x - j_x) ** 2 + (linkedj_y - j_y) ** 2)
                         if dist > 0:
@@ -657,14 +634,12 @@ def gaussian_scmap(self, joint_id, coords, data_item, size, scale):
                             Dy = (linkedj_y - j_y) * 1.0 / dist
 
                             d1 = [
-                                np.asscalar(Dx * j_x + Dy * j_y),
-                                np.asscalar(Dx * linkedj_x + Dy * linkedj_y),
+                                Dx * j_x + Dy * j_y,
+                                Dx * linkedj_x + Dy * linkedj_y,
                             ]  # in-line with direct axis
                             d1lowerboundary = min(d1)
                             d1upperboundary = max(d1)
-                            d2mid = np.asscalar(
-                                j_y * Dx - j_x * Dy
-                            )  # orthogonal direction
+                            d2mid = j_y * Dx - j_x * Dy  # orthogonal direction
 
                             distance_along = Dx * (x * stride + half_stride) + Dy * (
                                 y * stride + half_stride
@@ -701,5 +676,5 @@ def gaussian_scmap(self, joint_id, coords, data_item, size, scale):
 
             coordinateoffset += len(joint_ids)  # keeping track of the blocks
 
-        weights = self.compute_scmap_weights(scmap.shape, joint_id, data_item)
+        weights = self.compute_scmap_weights(scmap.shape, joint_id)
         return scmap, weights, locref_map, locref_mask
diff --git a/deeplabcut/pose_estimation_tensorflow/lib/crossvalutils.py b/deeplabcut/pose_estimation_tensorflow/lib/crossvalutils.py
index 8cf4574d01..5985ab2e1a 100644
--- a/deeplabcut/pose_estimation_tensorflow/lib/crossvalutils.py
+++ b/deeplabcut/pose_estimation_tensorflow/lib/crossvalutils.py
@@ -11,9 +11,7 @@
 import os
 import pickle
 import shutil
-import warnings
 from collections import defaultdict
-from itertools import groupby
 from tqdm import tqdm
 
 import networkx as nx
@@ -57,219 +55,6 @@ def _unsorted_unique(array):
     return np.asarray(array)[np.sort(inds)]
 
 
-def _rebuild_uncropped_metadata(metadata, image_paths, output_name=""):
-    train_inds_orig = set(metadata["data"]["trainIndices"])
-    train_inds, test_inds = [], []
-    for k, (_, group) in tqdm(enumerate(groupby(image_paths, _form_original_path))):
-        if image_paths.index(next(group)) in train_inds_orig:
-            train_inds.append(k)
-        else:
-            test_inds.append(k)
-    meta_new = metadata.copy()
-    meta_new["data"]["trainIndices"] = train_inds
-    meta_new["data"]["testIndices"] = test_inds
-
-    if output_name:
-        with open(output_name, "wb") as file:
-            pickle.dump(meta_new, file)
-
-    return meta_new
-
-
-def _rebuild_uncropped_data(data, params, output_name=""):
-    """
-    Reconstruct predicted data as if they had been obtained on full size images.
-    This is required to evaluate part affinity fields and cross-validate
-    and animal assembly.
-
-    Parameters
-    ----------
-    data : dict
-        Dictionary of predicted data as loaded from
-        _full.pickle files under evaluation-results.
-
-    params : dict
-        Evaluation settings. Formed from the metadata using _set_up_evaluation().
-
-    output_name : str
-        If passed, dump the uncropped data into a pickle file of the same name.
-
-    Returns
-    -------
-    (uncropped data, list of image paths)
-
-    """
-    image_paths = params["imnames"]
-    bodyparts = params["joint_names"]
-    idx = (
-        data[image_paths[0]]["groundtruth"][2]
-        .unstack("coords")
-        .reindex(bodyparts, level="bodyparts")
-        .index
-    )
-    individuals = idx.get_level_values("individuals").unique()
-    has_single = "single" in individuals
-    n_individuals = len(individuals) - has_single
-
-    data_new = dict()
-    with warnings.catch_warnings():
-        warnings.simplefilter("ignore", category=RuntimeWarning)
-        for basename, group in tqdm(groupby(image_paths, _form_original_path)):
-            imnames_ = list(group)
-            n_crops = len(imnames_)
-
-            # Sort crop patches to maximize the probability they overlap with others
-            all_coords_gt = [
-                data[imname]["groundtruth"][2].to_numpy().reshape((-1, 2))
-                for imname in imnames_
-            ]
-            overlap = np.zeros((n_crops, n_crops))
-            inds = list(zip(*np.triu_indices(n_crops, k=1)))
-            masks = dict()  # Cache boolean masks
-            for i1, i2 in inds:
-                if i1 not in masks:
-                    masks[i1] = np.isfinite(all_coords_gt[i1]).any(axis=1)
-                if i2 not in masks:
-                    masks[i2] = np.isfinite(all_coords_gt[i2]).any(axis=1)
-                overlap[i1, i2] = overlap[i2, i1] = np.sum(masks[i1] & masks[i2])
-            count = np.count_nonzero(overlap, axis=1)
-            imnames = [imnames_[i] for i in np.argsort(count)[::-1]]
-
-            # Form the ground truth back
-            ref_gt = None
-            all_trans = np.zeros(
-                (len(imnames), 2)
-            )  # Store translations w.r.t. first ref crop
-            for i, imname in enumerate(imnames):
-                coords_gt = data[imname]["groundtruth"][2].to_numpy().reshape((-1, 2))
-                if ref_gt is None:
-                    ref_gt = coords_gt
-                    continue
-                trans = np.nanmean(coords_gt - ref_gt, axis=0)
-                if np.all(~np.isnan(trans)):
-                    all_trans[i] = trans
-                empty = np.isnan(ref_gt)
-                has_value = ~np.isnan(coords_gt)
-                mask = np.any(empty & has_value, axis=1)
-                if mask.any():
-                    coords_gt_trans = coords_gt - all_trans[i]
-                    ref_gt[mask] = coords_gt_trans[mask]
-
-            # Match detections across crops
-            temp = pd.DataFrame(ref_gt, index=idx, columns=["x", "y"])
-            temp.columns.names = ["coords"]
-            if has_single:
-                temp.drop("single", level="individuals", inplace=True)
-            ref_pred = np.full(
-                (n_individuals, len(temp) // n_individuals, 4 + n_individuals), np.nan
-            )  # Hold x, y, prob, dist, ids
-            costs = dict()
-            shape = n_individuals, n_individuals
-            for ind in params["paf"]:
-                costs[ind] = {"m1": np.zeros(shape), "distance": np.full(shape, np.inf)}
-            if not np.isnan(temp.to_numpy()).all():
-                ref_gt_ = dict()
-                for bpt, df_ in temp.groupby("bodyparts"):
-                    values = df_.to_numpy()
-                    inds = np.flatnonzero(np.all(~np.isnan(values), axis=1))
-                    ref_gt_[bpt] = values, inds
-                for i, imname in enumerate(imnames):
-                    coords_pred = data[imname]["prediction"]["coordinates"][0]
-                    probs_pred = data[imname]["prediction"]["confidence"]
-                    costs_pred = data[imname]["prediction"]["costs"]
-                    try:
-                        ids_pred = data[imname]["prediction"]["identity"]
-                    except KeyError:
-                        ids_pred = None
-                    map_ = dict()
-                    for n, bpt in enumerate(ref_gt_):
-                        xy_gt, inds_gt = ref_gt_[bpt]
-                        ind = bodyparts.index(bpt)
-                        xy = coords_pred[ind]
-                        prob = probs_pred[ind]
-                        ids = None if ids_pred is None else ids_pred[ind]
-                        if inds_gt.size and xy.size:
-                            xy_trans = xy - all_trans[i]
-                            d = cdist(xy_gt[inds_gt], xy_trans)
-                            rows, cols = linear_sum_assignment(d)
-                            probs_ = prob[cols]
-                            ids_ = ids[cols] if ids is not None else None
-                            dists_ = d[rows, cols]
-                            inds_rows = inds_gt[rows]
-                            map_[n] = inds_rows, cols
-                            is_free = np.isnan(ref_pred[inds_rows, n]).all(axis=1)
-                            closer = dists_ < ref_pred[inds_rows, n, 3]
-                            mask = np.logical_or(is_free, closer)
-                            if mask.any():
-                                coords_ = xy_trans[cols]
-                                sl = inds_rows[mask]
-                                ref_pred[sl, n, :2] = coords_[mask]
-                                ref_pred[sl, n, 2] = probs_[mask].squeeze()
-                                ref_pred[sl, n, 3] = dists_[mask]
-                                if ids_ is not None:
-                                    ref_pred[sl, n, 4:] = ids_[mask]
-                    # Store the costs associated with the retained candidates
-                    for n, (ind1, ind2) in enumerate(params["paf_graph"]):
-                        if ind1 in map_ and ind2 in map_:
-                            sl1 = np.ix_(map_[ind1][0], map_[ind2][0])
-                            sl2 = np.ix_(map_[ind1][1], map_[ind2][1])
-                            mask = costs_pred[n]["m1"][sl2] > costs[n]["m1"][sl1]
-                            if mask.any():
-                                inds_lin = (sl1[0] * n_individuals + sl1[1])[mask]
-                                costs[n]["m1"].ravel()[inds_lin] = costs_pred[n]["m1"][
-                                    sl2
-                                ][mask]
-                                costs[n]["distance"].ravel()[inds_lin] = costs_pred[n][
-                                    "distance"
-                                ][sl2][mask]
-
-            ref_pred_ = ref_pred.swapaxes(0, 1)
-            coordinates = ref_pred_[..., :2]
-            confidence = ref_pred_[..., 2]
-            pred_dict = {
-                "prediction": {
-                    "coordinates": (coordinates,),
-                    "confidence": confidence,
-                    "costs": costs,
-                },
-                "groundtruth": (None, None, temp.stack(dropna=False)),
-            }
-            identities = ref_pred_[..., 4:]
-            if ~np.all(np.isnan(identities)):
-                pred_dict["prediction"]["identity"] = identities
-            data_new[basename] = pred_dict
-
-    image_paths = list(data_new)
-    data_new["metadata"] = data["metadata"]
-
-    if output_name:
-        with open(output_name, "wb") as file:
-            pickle.dump(data_new, file)
-
-    return data_new, image_paths
-
-
-def _rebuild_uncropped_in(base_folder,):
-    for dirpath, dirnames, filenames in os.walk(base_folder):
-        for file in filenames:
-            if file.endswith("_full.pickle"):
-                full_data_file = os.path.join(dirpath, file)
-                metadata_file = full_data_file.replace("_full.", "_meta.")
-                with open(full_data_file, "rb") as file:
-                    data = pickle.load(file)
-                with open(metadata_file, "rb") as file:
-                    metadata = pickle.load(file)
-                params = _set_up_evaluation(data)
-                _rebuild_uncropped_data(
-                    data, params, full_data_file.replace(".pickle", "_uncropped.pickle")
-                )
-                _rebuild_uncropped_metadata(
-                    metadata,
-                    params["imnames"],
-                    metadata_file.replace(".pickle", "_uncropped.pickle"),
-                )
-
-
 def _find_closest_neighbors(query, ref, k=3):
     n_preds = ref.shape[0]
     tree = cKDTree(ref)
@@ -313,34 +98,68 @@ def _calc_separability(
     return sep, threshold
 
 
-def _calc_within_between_pafs(data, metadata, per_bodypart=True, train_set_only=True):
+def _calc_within_between_pafs(
+    data,
+    metadata,
+    per_edge=True,
+    train_set_only=True,
+):
     train_inds = set(metadata["data"]["trainIndices"])
+    graph = data["metadata"]["PAFgraph"]
     within_train = defaultdict(list)
     within_test = defaultdict(list)
     between_train = defaultdict(list)
     between_test = defaultdict(list)
-    mask_diag = None
     for i, (key, dict_) in enumerate(data.items()):
         if key == "metadata":
             continue
+
         is_train = i in train_inds
         if train_set_only and not is_train:
             continue
+
+        df = dict_["groundtruth"][2]
+        try:
+            df.drop("single", level="individuals", inplace=True)
+        except KeyError:
+            pass
+        bpts = df.index.get_level_values("bodyparts").unique().to_list()
+        coords_gt = (df.unstack(["individuals", "coords"])
+                     .reindex(bpts, level="bodyparts")
+                     .to_numpy()
+                     .reshape((len(bpts), -1, 2)))
+        coords = dict_["prediction"]["coordinates"][0]
+        # Get animal IDs and corresponding indices in the arrays of detections
+        lookup = dict()
+        for i, (coord, coord_gt) in enumerate(zip(coords, coords_gt)):
+            inds = np.flatnonzero(np.all(~np.isnan(coord), axis=1))
+            inds_gt = np.flatnonzero(np.all(~np.isnan(coord_gt), axis=1))
+            if inds.size and inds_gt.size:
+                neighbors = _find_closest_neighbors(coord_gt[inds_gt], coord[inds], k=3)
+                found = neighbors != -1
+                lookup[i] = dict(zip(inds_gt[found], inds[neighbors[found]]))
+
         costs = dict_["prediction"]["costs"]
         for k, v in costs.items():
             paf = v["m1"]
-            nonzero = paf != 0
-            if mask_diag is None:
-                mask_diag = np.eye(paf.shape[0], dtype=bool)
-            within_vals = paf[np.logical_and(mask_diag, nonzero)]
-            between_vals = paf[np.logical_and(~mask_diag, nonzero)]
+            mask_within = np.zeros(paf.shape, dtype=bool)
+            s, t = graph[k]
+            if s not in lookup or t not in lookup:
+                continue
+            lu_s = lookup[s]
+            lu_t = lookup[t]
+            common_id = set(lu_s).intersection(lu_t)
+            for id_ in common_id:
+                mask_within[lu_s[id_], lu_t[id_]] = True
+            within_vals = paf[mask_within]
+            between_vals = paf[~mask_within]
             if is_train:
                 within_train[k].extend(within_vals)
                 between_train[k].extend(between_vals)
             else:
                 within_test[k].extend(within_vals)
                 between_test[k].extend(between_vals)
-    if not per_bodypart:
+    if not per_edge:
         within_train = np.concatenate([*within_train.values()])
         within_test = np.concatenate([*within_test.values()])
         between_train = np.concatenate([*between_train.values()])
@@ -412,7 +231,6 @@ def _benchmark_paf_graphs(
         print(f"Graph {j}|{n_graphs}")
         graph = [paf_graph[i] for i in paf]
         ass.paf_inds = paf
-        ass.graph = graph
         ass.assemble()
         oks = evaluate_assembly(ass.assemblies, ass_true_dict, oks_sigma)
         all_metrics.append(oks)
@@ -469,11 +287,11 @@ def _get_n_best_paf_graphs(
     if which not in ("best", "worst"):
         raise ValueError('`which` must be either "best" or "worst"')
 
-    (within_train, within_test), (between_train, _) = _calc_within_between_pafs(
-        data, metadata, train_set_only=False
+    (within_train, _), (between_train, _) = _calc_within_between_pafs(
+        data, metadata, train_set_only=True,
     )
     # Handle unlabeled bodyparts...
-    existing_edges = set(k for k, v in within_test.items() if v)
+    existing_edges = set(k for k, v in within_train.items() if v)
     if ignore_inds is not None:
         existing_edges = existing_edges.difference(ignore_inds)
     existing_edges = list(existing_edges)
@@ -485,26 +303,24 @@ def _get_n_best_paf_graphs(
             dict(zip(existing_edges, [0] * len(existing_edges)))
         )
 
-    scores, thresholds = zip(
+    scores, _ = zip(
         *[
-            _calc_separability(b_train, w_train, metric=metric)
-            for n, (w_train, b_train) in enumerate(
-                zip(within_train.values(), between_train.values())
-            )
-            if n in existing_edges
+            _calc_separability(between_train[n], within_train[n], metric=metric)
+            for n in existing_edges
         ]
     )
 
     # Find minimal skeleton
     G = nx.Graph()
     for edge, score in zip(existing_edges, scores):
-        G.add_edge(*full_graph[edge], weight=score)
+        if np.isfinite(score):
+            G.add_edge(*full_graph[edge], weight=score)
     if which == "best":
         order = np.asarray(existing_edges)[np.argsort(scores)[::-1]]
         if root is None:
             root = []
             for edge in nx.maximum_spanning_edges(G, data=False):
-                root.append(full_graph.index(list(edge)))
+                root.append(full_graph.index(sorted(edge)))
     else:
         order = np.asarray(existing_edges)[np.argsort(scores)]
         if root is None:
@@ -532,6 +348,8 @@ def cross_validate_paf_graphs(
     add_discarded=True,
     calibrate=False,
     overwrite_config=True,
+    n_graphs=10,
+    paf_inds=None,
 ):
     cfg = auxiliaryfunctions.read_config(config)
     inf_cfg = auxiliaryfunctions.read_plainconfig(inference_config)
@@ -547,9 +365,14 @@ def cross_validate_paf_graphs(
     to_ignore = auxfun_multianimal.filter_unwanted_paf_connections(
         cfg, params["paf_graph"]
     )
-    paf_inds, paf_scores = _get_n_best_paf_graphs(
-        data, metadata, params["paf_graph"], ignore_inds=to_ignore
+    best_graphs = _get_n_best_paf_graphs(
+        data, metadata, params["paf_graph"],
+        ignore_inds=to_ignore,
+        n_graphs=n_graphs,
     )
+    paf_scores = best_graphs[1]
+    if paf_inds is None:
+        paf_inds = best_graphs[0]
 
     if calibrate:
         trainingsetfolder = auxiliaryfunctions.GetTrainingSetFolder(cfg)
@@ -583,4 +406,4 @@ def cross_validate_paf_graphs(
     if output_name:
         with open(output_name, "wb") as file:
             pickle.dump([results], file)
-    return results
+    return results, paf_scores
diff --git a/deeplabcut/pose_estimation_tensorflow/lib/inferenceutils.py b/deeplabcut/pose_estimation_tensorflow/lib/inferenceutils.py
index 965b5faf59..b26f50d9b0 100644
--- a/deeplabcut/pose_estimation_tensorflow/lib/inferenceutils.py
+++ b/deeplabcut/pose_estimation_tensorflow/lib/inferenceutils.py
@@ -347,7 +347,8 @@ def _flatten_detections(data_dict):
 
     def extract_best_links(self, joints_dict, costs, trees=None):
         links = []
-        for (s, t), ind in zip(self.graph, self.paf_inds):
+        for ind in self.paf_inds:
+            s, t = self.graph[ind]
             dets_s = joints_dict.get(s, None)
             dets_t = joints_dict.get(t, None)
             if dets_s is None or dets_t is None:
@@ -592,6 +593,8 @@ def _assemble(self, data_dict, ind_frame):
         for joint in joints:
             bag[joint.label].append(joint)
 
+        assembled = set()
+
         if self.n_uniquebodyparts:
             unique = np.full((self.n_uniquebodyparts, 3), np.nan)
             for n, ind in enumerate(range(self.n_multibodyparts, self.n_keypoints)):
@@ -602,6 +605,11 @@ def _assemble(self, data_dict, ind_frame):
                     det = max(dets, key=lambda x: x.confidence)
                 else:
                     det = dets[0]
+                # Mark the unique body parts as assembled anyway so
+                # they are not used later on to fill assemblies.
+                assembled.update(d.idx for d in dets)
+                if det.confidence <= self.pcutoff and not self.add_discarded:
+                    continue
                 unique[n] = *det.pos, det.confidence
             if np.isnan(unique).all():
                 unique = None
@@ -613,7 +621,6 @@ def _assemble(self, data_dict, ind_frame):
 
         if self.identity_only:
             assemblies = []
-            assembled = set()
             get_attr = operator.attrgetter("group")
             temp = sorted(
                 (joint for joint in joints if np.isfinite(joint.confidence)),
@@ -648,7 +655,8 @@ def _assemble(self, data_dict, ind_frame):
                 vecs = np.vstack([link.to_vector() for link in links])
                 self._trees[ind_frame] = cKDTree(vecs)
 
-            assemblies, assembled = self.build_assemblies(links)
+            assemblies, assembled_ = self.build_assemblies(links)
+            assembled.update(assembled_)
 
         # Remove invalid assemblies
         discarded = set(
diff --git a/deeplabcut/pose_estimation_tensorflow/predict_videos.py b/deeplabcut/pose_estimation_tensorflow/predict_videos.py
index 9d232516bc..c17e2f33d4 100644
--- a/deeplabcut/pose_estimation_tensorflow/predict_videos.py
+++ b/deeplabcut/pose_estimation_tensorflow/predict_videos.py
@@ -290,18 +290,6 @@ def analyze_videos(
             from deeplabcut.pose_estimation_tensorflow.predict_multianimal import (
                 AnalyzeMultiAnimalVideo,
             )
-
-            # Re-use data-driven PAF graph for video analysis. Note that this must
-            # happen after setting up the TF session to avoid graph mismatch.
-            best_edges = dlc_cfg.get("paf_best")
-            if best_edges is not None:
-                best_graph = [dlc_cfg["partaffinityfield_graph"][i] for i in best_edges]
-            else:
-                best_graph = dlc_cfg["partaffinityfield_graph"]
-
-            dlc_cfg["partaffinityfield_graph"] = best_graph
-            dlc_cfg["num_limbs"] = len(best_graph)
-
             for video in Videos:
                 AnalyzeMultiAnimalVideo(
                     video,
diff --git a/deeplabcut/utils/auxfun_videos.py b/deeplabcut/utils/auxfun_videos.py
index 82b58a43b3..2897613dfa 100644
--- a/deeplabcut/utils/auxfun_videos.py
+++ b/deeplabcut/utils/auxfun_videos.py
@@ -346,7 +346,7 @@ def check_video_integrity(video_path):
     
 # Historically DLC used: from scipy.misc import imread, imresize >> deprecated functions
 def imread(path, mode=None):
-    return cv2.cvtColor(cv2.imread(path), cv2.COLOR_BGR2RGB)
+    return cv2.imread(path)[..., ::-1]  # ~10% faster than using cv2.cvtColor
 
 
 # https://docs.opencv.org/3.4.0/da/d54/group__imgproc__transform.html#ga5bb5a1fea74ea38e1a5445ca803ff121
diff --git a/deeplabcut/utils/auxiliaryfunctions.py b/deeplabcut/utils/auxiliaryfunctions.py
index d7144b5f20..c860eccce1 100644
--- a/deeplabcut/utils/auxiliaryfunctions.py
+++ b/deeplabcut/utils/auxiliaryfunctions.py
@@ -62,7 +62,6 @@ def create_config_template(multianimal=False):
         \n
     # Cropping Parameters (for analysis and outlier frame detection)
         cropping:
-        croppedtraining:
     #if cropping is true for analysis, then set the values here:
         x1:
         x2:
@@ -110,7 +109,6 @@ def create_config_template(multianimal=False):
         \n
     # Cropping Parameters (for analysis and outlier frame detection)
         cropping:
-        croppedtraining:
     #if cropping is true for analysis, then set the values here:
         x1:
         x2:
diff --git a/deeplabcut/utils/conversioncode.py b/deeplabcut/utils/conversioncode.py
index d29b6eaac8..fd47bc7064 100644
--- a/deeplabcut/utils/conversioncode.py
+++ b/deeplabcut/utils/conversioncode.py
@@ -199,13 +199,9 @@ def merge_windowsannotationdataONlinuxsystem(cfg):
 
     AnnotationData = []
     data_path = Path(cfg["project_path"], "labeled-data")
-    use_cropped = cfg.get("croppedtraining", False)
     annotationfolders = []
     for elem in auxiliaryfunctions.grab_files_in_folder(data_path, relative=False):
-        if os.path.isdir(elem) and (
-            (use_cropped and elem.endswith("_cropped"))
-            or not (use_cropped or "_cropped" in elem)
-        ):
+        if os.path.isdir(elem):
             annotationfolders.append(elem)
     print("The following folders were found:", annotationfolders)
     for folder in annotationfolders:
diff --git a/docs/maDLC_UserGuide.md b/docs/maDLC_UserGuide.md
index dc79ca93d1..a51e3087c1 100644
--- a/docs/maDLC_UserGuide.md
+++ b/docs/maDLC_UserGuide.md
@@ -224,21 +224,9 @@ deeplabcut.check_labels(config_path, visualizeindividuals=True/False)
 
 For each video directory in labeled-data this function creates a subdirectory with **labeled** as a suffix. Those directories contain the frames plotted with the annotated body parts. The user can double check if the body parts are labeled correctly. If they are not correct, the user can reload the frames (i.e. `deeplabcut.label_frames`), move them around, and click save again.
 
-**CROP+LABEL:** When you are done checking the label quality and adjusting if needed, please then use this new function to crop frames /labels for more efficient training. PLEASE call this before you create a training dataset by running:
-```python
-deeplabcut.cropimagesandlabels(path_config_file, userfeedback=False)
-```
-### Create Training Dataset:
-
-Ideally for training DNNs, one uses large batch sizes. Thus, for mutli-animal training batch processing is preferred. This means that we'd like the images to be similarly sized. You can of course have differing size of images you label (but we suggest cropping out useless pixels!). So, we have a new function that can pre-process your data to be compatible with batch training. As noted above, please run this function before you `create_multianmialtraining_dataset`. This function assures that each crop is "small", by default 400 x 400, which allows larger batchsizes and that there are multiple crops so that different parts of larger images are covered.
-
-You **should also first run** `deeplabcut.cropimagesandlabels(config_path)` before creating a training set, as we use batch processing and many users have smaller GPUs that cannot accommodate larger images + larger batchsizes. This is also a type of data augmentation.
 
-NOTE: you can edit the crop size. If your images are very large (2k, 4k pixels), consider increasing this size, but be aware unless you have a lagre GPU (24 GB or more), you will hit memory errors. _You can lower the batchsize, but this may affect performance._
+### Create Training Dataset:
 
-```python
-deeplabcut.cropimagesandlabels(path_config_file, size=(400, 400), userfeedback=False)
-```
 At this point you also select your neural network type. Please see Lauer et al. 2021 for options. For **create_multianimaltraining_dataset** we already changed this such that by default you will use imgaug, ADAM optimization, our new DLCRNet, and batch training. We suggest these defaults at this time. Then run:
 
 ```python
@@ -265,7 +253,7 @@ are listed in Box 2.
 **OPTIONAL POINTS:**
 
 With the data-driven skeleton selection introduced in 2.2rc1, DLC networks are trained by default
-on complete skeletons (i.e., they learn all possible redundant connections), before being optimially pruned
+on complete skeletons (i.e., they learn all possible redundant connections), before being optimally pruned
 at model evaluation. Although this procedure is by far superior to manually defining a graph,
 we leave manually-defining a skeleton as an option for the advanced user:
 
@@ -288,6 +276,14 @@ Our recent [A Primer on Motion Capture with Deep Learning: Principles, Pitfalls,
 
 Alternatively, you can set the loader (as well as other training parameters) in the **pose_cfg.yaml** file of the model that you want to train. Note, to get details on the options, look at the default file: [**pose_cfg.yaml**](https://github.com/AlexEMG/DeepLabCut/blob/master/deeplabcut/pose_cfg.yaml).
 
+Importantly, image cropping as previously done with `deeplabcut.cropimagesandlabels` in multi-animal projects
+is now part of the augmentation pipeline. In other words, image crops are no longer stored in labeled-data/..._cropped
+folders. Crop size still defaults to (400, 400); if your images are very large (e.g. 2k, 4k pixels), consider increasing the crop size, but be aware unless you have a strong GPU (24 GB memory or more), you will hit memory errors. You can lower the batch size, but this may affect performance.
+
+In addition, one can specify a crop sampling strategy: crop centers can either be taken at random over the image (`uniform`) or the annotated keypoints (`keypoints`); with a focus on regions of the scene with high body part density (`density`); last, combining `uniform` and `density` for a `hybrid` balanced strategy (this is the default strategy). Note that both parameters can be easily edited prior to training in the **pose_cfg.yaml** configuration file.
+As a reminder, cropping images into smaller patches is a form of data augmentation that simultaneously
+allows the use of batch processing even on small GPUs that could not otherwise accommodate larger images + larger batchsizes (this usually increases performance and decreasing training time). 
+
 
 ### Train The Network:
 
diff --git a/docs/tutorial.md b/docs/tutorial.md
index 5b8ea7064f..45fe1917ab 100644
--- a/docs/tutorial.md
+++ b/docs/tutorial.md
@@ -51,18 +51,6 @@ deeplabcut.check_labels(
 )
 ```
 
-
-**Crop frames to augment the dataset**
-```python
-deeplabcut.cropimagesandlabels(
-    config_path,
-    numcrops=10,
-    size=(400, 400),
-    userfeedback=False,
-)
-```
-
-
 **Create the training dataset**
 ```python
 deeplabcut.create_multianimaltraining_dataset(
diff --git a/examples/COLAB_maDLC_TrainNetwork_VideoAnalysis.ipynb b/examples/COLAB_maDLC_TrainNetwork_VideoAnalysis.ipynb
index d38d47df6c..9f1b68c2eb 100644
--- a/examples/COLAB_maDLC_TrainNetwork_VideoAnalysis.ipynb
+++ b/examples/COLAB_maDLC_TrainNetwork_VideoAnalysis.ipynb
@@ -200,14 +200,6 @@
         "id": "eMeUwgxPoEJP"
       },
       "source": [
-        "# ATTENTION:\n",
-        "\n",
-        "#which shuffle do you want to create and train?\n",
-        "shuffle = 1 #edit if needed; 1 is the default.\n",
-        "\n",
-        "#Note, you must run this. If your images are smaller than 400 by 400, please make these numbers smaller.\n",
-        "deeplabcut.cropimagesandlabels(path_config_file, size=(400, 400), userfeedback=False)\n",
-        "\n",
         "#if you labeled on Windows, please set the windows2linux=True:\n",
         "deeplabcut.create_multianimaltraining_dataset(path_config_file, Shuffles=[shuffle], net_type=\"dlcrnet_ms5\",windows2linux=False)"
       ],
diff --git a/examples/testscript_multianimal.py b/examples/testscript_multianimal.py
index d623e70fd2..285c3685f4 100644
--- a/examples/testscript_multianimal.py
+++ b/examples/testscript_multianimal.py
@@ -70,7 +70,7 @@
         for image in auxiliaryfunctions.grab_files_in_folder(image_folder, "png")
     ]
     fake_data = np.tile(
-        np.repeat(50 * np.arange(len(animals_id)) + 100, 2), (len(index), 1)
+        np.repeat(50 * np.arange(len(animals_id)) + 50, 2), (len(index), 1)
     )
     df = pd.DataFrame(fake_data, index=index, columns=columns)
     output_path = os.path.join(image_folder, f"CollectedData_{SCORER}.csv")
@@ -80,9 +80,6 @@
     )
     print("Artificial data created.")
 
-    print("Cropping and exchanging")
-    deeplabcut.cropimagesandlabels(config_path, userfeedback=False)
-
     print("Checking labels...")
     deeplabcut.check_labels(config_path, draw_skeleton=False)
     print("Labels checked.")
@@ -95,12 +92,13 @@
     model_folder = auxiliaryfunctions.GetModelFolder(
         TRAIN_SIZE, 1, cfg, cfg["project_path"]
     )
-    pose_config_path = os.path.join(model_folder, "train/pose_cfg.yaml")
+    pose_config_path = os.path.join(model_folder, "train", "pose_cfg.yaml")
     edits = {
         "global_scale": 0.5,
         "batch_size": 1,
         "save_iters": N_ITER,
         "display_iters": N_ITER // 2,
+        "crop_size": [200, 200],
         # "multi_step": [[0.001, N_ITER]],
     }
     deeplabcut.auxiliaryfunctions.edit_config(pose_config_path, edits)
diff --git a/tests/conftest.py b/tests/conftest.py
index 8354e75360..44a67b90cb 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -3,11 +3,24 @@
 import pickle
 import pytest
 from deeplabcut.pose_estimation_tensorflow.lib import inferenceutils, crossvalutils
+from PIL import Image
 
 
 TEST_DATA_DIR = os.path.join(os.path.dirname(os.path.realpath(__file__)), "data")
 
 
+@pytest.fixture(scope="session")
+def sample_image():
+    return np.asarray(Image.open(os.path.join(TEST_DATA_DIR, "image.png")))
+
+
+@pytest.fixture(scope="session")
+def sample_keypoints():
+    with open(os.path.join(TEST_DATA_DIR, "trimouse_assemblies.pickle"), "rb") as file:
+        temp = pickle.load(file)
+    return np.concatenate(temp[0])[:, :2]
+
+
 @pytest.fixture(scope="session")
 def real_assemblies():
     with open(os.path.join(TEST_DATA_DIR, "trimouse_assemblies.pickle"), "rb") as file:
@@ -23,14 +36,11 @@ def real_tracklets():
 
 
 @pytest.fixture(scope="session")
-def uncropped_data_and_metadata():
+def evaluation_data_and_metadata():
     full_data_file = os.path.join(TEST_DATA_DIR, "trimouse_eval.pickle")
     metadata_file = full_data_file.replace("eval", "meta")
     with open(full_data_file, "rb") as file:
         data = pickle.load(file)
     with open(metadata_file, "rb") as file:
         metadata = pickle.load(file)
-    params = crossvalutils._set_up_evaluation(data)
-    data_unc, _ = crossvalutils._rebuild_uncropped_data(data, params)
-    meta_unc = crossvalutils._rebuild_uncropped_metadata(metadata, params["imnames"])
-    return data_unc, meta_unc
+    return data, metadata
diff --git a/tests/data/image.png b/tests/data/image.png
index f08bf8719c..3e45ff7840 100644
Binary files a/tests/data/image.png and b/tests/data/image.png differ
diff --git a/tests/data/trimouse_eval.pickle b/tests/data/trimouse_eval.pickle
index 4d5a346e90..3f90bf4c95 100644
Binary files a/tests/data/trimouse_eval.pickle and b/tests/data/trimouse_eval.pickle differ
diff --git a/tests/data/trimouse_meta.pickle b/tests/data/trimouse_meta.pickle
index c62dc6b89f..5673195aa4 100644
Binary files a/tests/data/trimouse_meta.pickle and b/tests/data/trimouse_meta.pickle differ
diff --git a/tests/test_crossvalutils.py b/tests/test_crossvalutils.py
index 7cc67a4f34..506f287757 100644
--- a/tests/test_crossvalutils.py
+++ b/tests/test_crossvalutils.py
@@ -2,11 +2,11 @@
 from deeplabcut.pose_estimation_tensorflow.lib import crossvalutils
 
 
-BEST_GRAPH = [2, 56, 7, 31, 38, 63, 65, 60, 54, 1, 13]
+BEST_GRAPH = [14, 15, 16, 11, 22, 31, 61, 7, 59, 62, 64]
 
 
-def test_get_n_best_paf_graphs(uncropped_data_and_metadata):
-    data, metadata = uncropped_data_and_metadata
+def test_get_n_best_paf_graphs(evaluation_data_and_metadata):
+    data, metadata = evaluation_data_and_metadata
     params = crossvalutils._set_up_evaluation(data)
     n_graphs = 5
     paf_inds, dict_ = crossvalutils._get_n_best_paf_graphs(
@@ -19,8 +19,8 @@ def test_get_n_best_paf_graphs(uncropped_data_and_metadata):
     assert len(paf_inds[-1]) == len(params["paf_graph"])
 
 
-def test_benchmark_paf_graphs(uncropped_data_and_metadata):
-    data, _ = uncropped_data_and_metadata
+def test_benchmark_paf_graphs(evaluation_data_and_metadata):
+    data, _ = evaluation_data_and_metadata
     cfg = {
         "individuals": ["mickey", "minnie", "bianca"],
         "uniquebodyparts": [],
@@ -47,5 +47,5 @@ def test_benchmark_paf_graphs(uncropped_data_and_metadata):
     assert len(all_scores) == 1
     assert all_scores[0][1] == BEST_GRAPH
     miss, purity = results[1].xs("mean", level=1).to_numpy().squeeze()
-    assert np.isclose(miss, 0.0)
-    assert np.isclose(purity, 1.0)
+    assert np.isclose(miss, 0.02, atol=1e-2)
+    assert np.isclose(purity, 0.98, atol=1e-2)
diff --git a/tests/test_dataset_augmentation.py b/tests/test_dataset_augmentation.py
new file mode 100644
index 0000000000..ff32fda697
--- /dev/null
+++ b/tests/test_dataset_augmentation.py
@@ -0,0 +1,77 @@
+import imgaug.augmenters as iaa
+import pytest
+from deeplabcut.pose_estimation_tensorflow.datasets import augmentation
+
+
+@pytest.mark.parametrize(
+    "width, height",
+    [
+        (200, 200),
+        (300, 300),
+        (400, 400),
+    ]
+)
+def test_keypoint_aware_cropping(
+    sample_image,
+    sample_keypoints,
+    width,
+    height,
+):
+    aug = augmentation.KeypointAwareCropToFixedSize(
+        width=width,
+        height=height,
+    )
+    images_aug, keypoints_aug = aug(
+        images=[sample_image],
+        keypoints=[sample_keypoints],
+    )
+    assert len(images_aug) == len(keypoints_aug) == 1
+    assert all(im.shape[:2] == (height, width) for im in images_aug)
+    # Ensure at least a keypoint is visible in each crop
+    assert all(len(kpts) for kpts in keypoints_aug)
+
+    # Test passing in a batch of frames
+    n_samples = 8
+    images_aug, keypoints_aug = aug(
+        images=[sample_image] * n_samples,
+        keypoints=[sample_keypoints] * n_samples,
+    )
+    assert len(images_aug) == len(keypoints_aug) == n_samples
+
+
+@pytest.mark.parametrize(
+    "width, height",
+    [
+        (200, 200),
+        (300, 300),
+        (400, 400),
+    ]
+)
+def test_sequential(
+    sample_image,
+    sample_keypoints,
+    width,
+    height,
+):
+    # Guarantee that images smaller than crop size are handled fine
+    very_small_image = sample_image[:50, :50]
+    aug = iaa.Sequential([
+        iaa.PadToFixedSize(width, height),
+        augmentation.KeypointAwareCropToFixedSize(width, height),
+    ])
+    images_aug, keypoints_aug = aug(
+        images=[very_small_image],
+        keypoints=[sample_keypoints],
+    )
+    assert len(images_aug) == len(keypoints_aug) == 1
+    assert all(im.shape[:2] == (height, width) for im in images_aug)
+    # Ensure at least a keypoint is visible in each crop
+    assert all(len(kpts) for kpts in keypoints_aug)
+
+    # Test passing in a batch of frames
+    n_samples = 8
+    images_aug, keypoints_aug = aug(
+        images=[very_small_image] * n_samples,
+        keypoints=[sample_keypoints] * n_samples,
+    )
+    assert len(images_aug) == len(keypoints_aug) == n_samples
diff --git a/tests/test_inferenceutils.py b/tests/test_inferenceutils.py
index 9ef969d55e..f43d3a3b95 100644
--- a/tests/test_inferenceutils.py
+++ b/tests/test_inferenceutils.py
@@ -135,9 +135,7 @@ def test_assembler(tmpdir_factory, real_assemblies):
         [3, 4],
         [0, 2],
     ]
-    paf_inds = [ass.graph.index(edge) for edge in naive_graph]
-    ass.graph = naive_graph
-    ass.paf_inds = paf_inds
+    ass.paf_inds = [ass.graph.index(edge) for edge in naive_graph]
     ass.assemble()
     assert not ass.unique
     assert len(ass.assemblies) == len(real_assemblies)
@@ -184,9 +182,7 @@ def test_assembler_with_identity(tmpdir_factory, real_assemblies):
         [3, 4],
         [0, 2],
     ]
-    paf_inds = [ass.graph.index(edge) for edge in naive_graph]
-    ass.graph = naive_graph
-    ass.paf_inds = paf_inds
+    ass.paf_inds = [ass.graph.index(edge) for edge in naive_graph]
     ass.assemble()
     assert not ass.unique
     assert len(ass.assemblies) == len(real_assemblies)
diff --git a/tests/test_pose_multianimal_imgaug.py b/tests/test_pose_multianimal_imgaug.py
index de29f8d6b2..64d5a3e196 100644
--- a/tests/test_pose_multianimal_imgaug.py
+++ b/tests/test_pose_multianimal_imgaug.py
@@ -27,61 +27,59 @@ def ma_dataset():
 
 
 @pytest.mark.parametrize(
-    "batch_size, scale, stride",
+    "scale, stride",
     [
-        (1, 0.6, 2),
-        (1, 0.6, 4),
-        (1, 0.6, 8),
-        (8, 0.8, 4),
-        (8, 1.0, 8),
-        (8, 1.2, 8),
-        (16, 0.6, 4),
-        (16, 0.8, 8),
+        (0.6, 2),
+        (0.6, 4),
+        (0.6, 8),
+        (0.8, 4),
+        (1.0, 8),
+        (1.2, 8),
+        (0.6, 4),
+        (0.8, 8),
     ]
 )
-def test_get_batch(
+def test_calc_target_and_scoremap_sizes(
     ma_dataset,
-    batch_size,
     scale,
     stride,
 ):
-    ma_dataset.batch_size = batch_size
     ma_dataset.cfg["global_scale"] = scale
     ma_dataset.cfg["stride"] = stride
-    (
-        batch_images,
-        joint_ids,
-        batch_joints,
-        data_items,
-        sm_size,
-        target_size,
-    ) = ma_dataset.get_batch()
-    assert len(batch_images) == len(joint_ids) == len(batch_joints) \
-           == len(data_items) == batch_size
-    for data_item, joint_id in zip(data_items, joint_ids):
-        assert len(data_item.joints) == len(joint_id)
-        for joints, id_ in zip(data_item.joints.values(), joint_id):
-            np.testing.assert_equal(joints[:, 0], id_)
+    target_size, sm_size = ma_dataset.calc_target_and_scoremap_sizes()
     np.testing.assert_equal(np.asarray([400, 400]) * scale, target_size)
     np.testing.assert_equal(target_size / stride, sm_size)
 
 
-@pytest.mark.parametrize(
-    "height, width, prob",
-    [
-        (None, None, 0.5),
-        (200, 250, 0.3),
-    ]
-)
-def test_build_augmentation_pipeline(ma_dataset, height, width, prob):
-    _ = ma_dataset.build_augmentation_pipeline(height, width, prob)
+def test_get_batch(ma_dataset):
+    for batch_size in 1, 4, 8, 16:
+        ma_dataset.batch_size = batch_size
+        (
+            batch_images,
+            joint_ids,
+            batch_joints,
+            data_items,
+        ) = ma_dataset.get_batch()
+        assert len(batch_images) == len(joint_ids) == len(batch_joints) \
+               == len(data_items) == batch_size
+        for data_item, joint_id in zip(data_items, joint_ids):
+            assert len(data_item.joints) == len(joint_id)
+            for joints, id_ in zip(data_item.joints.values(), joint_id):
+                np.testing.assert_equal(joints[:, 0], id_)
+
+
+def test_build_augmentation_pipeline(ma_dataset):
+    for prob in (0.3, 0.5):
+        _ = ma_dataset.build_augmentation_pipeline(prob)
 
 
 @pytest.mark.parametrize("num_idchannel", range(4))
 def test_get_targetmaps(ma_dataset, num_idchannel):
     ma_dataset.cfg["num_idchannel"] = num_idchannel
     batch = ma_dataset.get_batch()[1:]
-    maps = ma_dataset.get_targetmaps_update(*batch)
+    target_size, sm_size = ma_dataset.calc_target_and_scoremap_sizes()
+    scale = np.mean(target_size / ma_dataset.default_crop_size)
+    maps = ma_dataset.get_targetmaps_update(*batch, sm_size, scale)
     assert all(len(map_) == ma_dataset.batch_size for map_ in maps.values())
     assert maps[Batch.part_score_targets][0].shape \
            == maps[Batch.part_score_weights][0].shape