Intelligent(?) image sharpness filtering for optimized NeRF training datasets
import os
import argparse
import shutil
import subprocess
from ImageSelector import ImageSelector

def extract_frames(input_vid, output_path):
    if not args.yes:
        answer = input(f"About to extract all frames from '{input_vid}' into '{output_path}' (even with --pretend!). This folder will persist unless manually removed. Continue? [y/N]: ").lower()
        if answer not in ["y", "yes"]:
            print("Aborting.")
            exit()
    if not os.path.exists(output_path):
        os.makedirs(output_path)
    cmd = [
        'ffmpeg',
        '-i', input_vid,
        '-vsync', 'vfr',
        '-v', 'quiet',
        '-stats',
        '-q:v', '1',  # Use highest quality for JPEG
        os.path.join(output_path, 'frame%05d.jpg')
    ]
    subprocess.run(cmd)

def main(input_path, output_path, img_exts, target_count, groups=None, scalar=None):
    # Check if input_path is a folder or video
    if os.path.isdir(input_path):
        images = [os.path.join(input_path, img) for img in os.listdir(input_path) if img.lower().endswith(img_exts)]
        images.sort()
    else:
        extract_frames(input_path, output_path)
        images = [os.path.join(output_path, img) for img in os.listdir(output_path) if img.endswith('.jpg')]
        images.sort()

    total_images = len(images)
    print(f"Found a total of {total_images} images to work on.")

    if total_images == 0:
        print("Error: No images/frames found to work on. If you specified a folder, make sure the images in it are in --exts (default: jpg, jpeg, png).")
        print("Aborting.")
        return

    if target_count is None:
        target_count = int(total_images * (args.target_percentage / 100))

    selector = ImageSelector(images)
    selected_images = selector.filter_sharpest_images(target_count, groups, scalar)

    if args.pretend:
        print(f"Would have retained {len(selected_images)} sharpest images. (--pretend)")
        if not os.path.isdir(input_path):
            print(f"Warning: Folder '{output_path}' persists, with all extracted frames in it.")
        return

    if not args.yes:
        if not os.path.isdir(input_path):
            print(f"Folder '{output_path}' will persist, with all extracted frames in it. Responding 'y' will prune it down to the desired output, but responding 'n' will keep it as-is, with all extracted frames.")
        answer = input(f"About to delete all but {len(selected_images)} sharpest images. Continue? [y/N]: ").lower()
        if answer not in ["y", "yes"]:
            print("Aborting.")
            return

    if not os.path.exists(output_path):
        os.makedirs(output_path)

    if os.path.isdir(input_path) and os.path.normcase(os.path.abspath(os.path.normpath(input_path))) != os.path.normcase(os.path.abspath(os.path.normpath(output_path))):
        print(os.path.normcase(os.path.abspath(os.path.normpath(input_path))))
        print(os.path.normcase(os.path.abspath(os.path.normpath(output_path))))
        for img in selected_images:
            new_location = os.path.join(output_path, os.path.basename(img))
            shutil.copy2(img, new_location)
    else:
        for img in images:
            if img not in selected_images:
                os.remove(img)

    print(f"Retained {len(selected_images)} sharpest images.")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Select and retain only the sharpest frames from a video or folder of images.")
    parser.add_argument('--input_path', required=True, help="Path to the input_path video or folder of images.")
    parser.add_argument('--output_path', help="Directory to save the preserved images. Mandatory for video, optional for images. Will delete images in-place if not specified.")
    parser.add_argument('--exts', default='jpg,jpeg,png', help="If --input_path is a dir, image extensions to check for. Default is 'jpg,jpeg,png'. Anything supported by OpenCV should work.")
    group_target = parser.add_mutually_exclusive_group(required=True)
    group_target.add_argument('--target_count', type=int, help="Target number of images to retain.")
    group_target.add_argument('--target_percentage', type=float, help="Target percentage of top quality images to retain. I.e. --target_percentage 95 removes the 5%% worst quality images.")
    group_division = parser.add_mutually_exclusive_group()
    group_division.add_argument('--groups', type=int, help="Specify the number of groups to divide the images into.")
    group_division.add_argument('--scalar', type=int, help="Specify the scalar value to determine group division if groups is not provided.")
    parser.add_argument('--pretend', action='store_true', help="Pretend mode. Do not delete anything, just show what would have been deleted. Warning: Will still create and populate the output dir if input_path is a video!")
    parser.add_argument('--yes', '-y', action='store_true', help="Automatically answer 'yes' to all prompts and execute actions.")

    args = parser.parse_args()

    if not os.path.exists(args.input_path):
        parser.error(f"--input_path does not exist: \"{args.input_path}\"")
    elif not os.path.isdir(args.input_path) and not args.output_path:
        parser.error("The --output_path argument is mandatory when the input_path is a video.")
    elif not args.output_path or (os.path.normcase(os.path.abspath(os.path.normpath(args.input_path))) is os.path.normcase(os.path.abspath(os.path.normpath(args.output_path)))):
        args.output_path = args.input_path

    img_exts = supported_extensions = tuple('.' + ext.lower() for ext in args.exts.split(','))

    main(args.input_path, args.output_path, img_exts, args.target_count, args.groups, args.scalar)

01_filter_raw_data.py
SharkWipf/nerf_dataset_preprocessing_helper
import os
import json
import argparse
from ImageSelector import ImageSelector

def main(transforms_path, target_count, output_file, groups=None, scalar=None, pretend=False):
    # Determine whether we're working with a directory or a specific JSON file
    if os.path.isdir(transforms_path):
        main_directory = transforms_path
        json_path = os.path.join(main_directory, "transforms.json")
    else:
        json_path = transforms_path
        main_directory = os.path.dirname(json_path)

    # Load transform data from the provided JSON path
    with open(json_path, "r") as file:
        transforms_data = json.load(file)

    frames = sorted(transforms_data["frames"], key=lambda x: x["colmap_im_id"])
    images = [os.path.join(main_directory, frame["file_path"]) for frame in frames]

    if target_count is None:
        total_images = len(images)
        target_count = int(total_images * (args.target_percentage / 100))

    selector = ImageSelector(images)
    selected_images = selector.filter_sharpest_images(target_count, groups, scalar)

    new_frames = [frame for frame in transforms_data["frames"] if os.path.join(main_directory, frame["file_path"]) in selected_images]
    transforms_data["frames"] = new_frames

    # If pretend mode, just print what would be done
    if pretend:
        print(f"Would have retained {len(selected_images)} sharpest images. Would have saved to {output_file} (--pretend)")
    else:
        if not args.yes:
            answer = input(f"About to remove all but {len(selected_images)} images from '{output_file}', continue? [y/N]: ").lower()
            if answer not in ["y", "yes"]:
                print("Aborting.")
                return

        output_path = os.path.join(os.getcwd(), output_file)
        with open(output_path, "w") as file:
            json.dump(transforms_data, file, indent=4)
        print(f"Retained {len(selected_images)} sharpest images. Saved to '{output_path}'.")

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description="Select and retain only the sharpest images from a COLMAP output.")
    parser.add_argument('--transforms_path', required=True, help="Path to the main COLMAP output directory or the transforms.json file.")
    group_target = parser.add_mutually_exclusive_group(required=True)
    group_target.add_argument('--target_count', type=int, help="Target number of images to retain.")
    group_target.add_argument('--target_percentage', type=float, help="Target percentage of top quality images to retain. I.e. --target_percentage 95 removes the 5%% worst quality images.")
    parser.add_argument('--output_file', default=None, help="Path to save the output JSON. If not specified, the default is transforms_filtered.json in the same directory as the transforms file.")
    group_division = parser.add_mutually_exclusive_group()
    group_division.add_argument('--groups', type=int, help="Specify the number of groups to divide the images into.")
    group_division.add_argument('--scalar', type=int, help="Specify the scalar value to determine group division if groups is not provided.")

    parser.add_argument('--pretend', action='store_true', help="Pretend mode. Do not write or delete anything, just show what would have been done.")
    parser.add_argument('--yes', '-y', action='store_true', help="Automatically answer 'yes' to all prompts and execute actions.")

    args = parser.parse_args()

    if args.output_file is None:
        args.output_file = os.path.join(os.path.dirname(args.transforms_path), "transforms_filtered.json")

    main(args.transforms_path, args.target_count, args.output_file, args.groups, args.scalar, args.pretend)


02_filter_colmap_data.py
SharkWipf/nerf_dataset_preprocessing_helper
import cv2
from tqdm import tqdm
from graphlib import draw_graph


class ImageSelector:
    def __init__(self, images):
        self.images = images
        self.image_fm = self._compute_sharpness_values()

    def _compute_sharpness_values(self):
        print("Calculating image sharpness...")
        return [(self.variance_of_laplacian(cv2.cvtColor(cv2.imread(img), cv2.COLOR_BGR2GRAY)), img) for img in tqdm(self.images)]

    @staticmethod
    def variance_of_laplacian(image):
        return cv2.Laplacian(image, cv2.CV_64F).var()

    @staticmethod
    def distribute_evenly(total, num_of_groups):
        ideal_per_group = total / num_of_groups
        accumulated_error = 0.0
        distribution = [0] * num_of_groups

        for i in range(num_of_groups):
            distribution[i] = int(ideal_per_group)
            accumulated_error += ideal_per_group - distribution[i]

            while accumulated_error >= 1.0:
                distribution[i] += 1
                accumulated_error -= 1.0

        return distribution, ideal_per_group

    def generate_deleted_images_graph(self, selected_images):
        bins = 100
        step = len(self.images) // bins
        percentages = []

        for i in range(bins):
            start_idx = i * step
            end_idx = (i + 1) * step if i != bins - 1 else len(self.images)
            current_bin = self.images[start_idx:end_idx]
            deleted_count = sum(1 for img in current_bin if img not in selected_images)
            avg = deleted_count / len(current_bin)
            percentages.append(avg * 100)

        draw_graph(percentages, "Distribution of to-be-deleted images")

    def generate_quality_graph(self):
        draw_graph([quality for quality, _ in self.image_fm], "Distribution of image quality")

    def filter_sharpest_images(self, target_count, group_count=None, scalar=1):
        if scalar is None:
            scalar = 1
        if group_count is None:
            group_count = target_count // (2 ** (scalar - 1))
            group_count = max(1, group_count)

        split = len(self.images) / target_count
        ratio = target_count / len(self.images)
        formatted_ratio = "{:.1%}".format(ratio)
        print(f"Requested {target_count} out of {len(self.images)} images ({formatted_ratio}, 1 in {split:.1f}).")

        group_sizes, ideal_total_images_per_group = self.distribute_evenly(len(self.images), group_count)
        images_per_group_list, ideal_selected_images_per_group = self.distribute_evenly(target_count, group_count)

        print(f"Selecting {target_count} image{'s' if target_count != 1 else ''} across {group_count} group{'s' if group_count != 1 else ''}, with total ~{ideal_total_images_per_group:.1f} image{'s' if ideal_total_images_per_group != 1 else ''} per group and selecting ~{ideal_selected_images_per_group:.1f} image{'s' if ideal_selected_images_per_group != 1 else ''} per group (scalar {scalar}).")
        draw_graph([(i % 2) for i in range(group_count)], "Group layout", 100)

        images_per_group_list, _ = self.distribute_evenly(target_count, group_count)

        selected_images = []
        offset_index = 0
        for idx, size in enumerate(group_sizes):
            end_idx = offset_index + size
            group = sorted(self.image_fm[offset_index:end_idx], reverse=True)
            selected_images.extend([img[1] for img in group[:images_per_group_list[idx]]])
            offset_index = end_idx

        self.generate_deleted_images_graph(selected_images)
        self.generate_quality_graph()

        return selected_images

ImageSelector.py
SharkWipf/nerf_dataset_preprocessing_helper
def draw_graph(data, description, bins_count=100):
    height_chars = "▁▂▃▄▅▆▇█"

    if len(data) == bins_count:
        binned_data = data
    elif len(data) > bins_count:
        bin_size = len(data) / bins_count
        binned_data = []
        # Compute the average for each bin
        for i in range(bins_count):
            start_idx = int(i * bin_size)
            end_idx = int((i + 1) * bin_size) if i != bins_count - 1 else len(data)
            avg = sum(data[start_idx:end_idx]) / (end_idx - start_idx) if (end_idx - start_idx) != 0 else 0
            binned_data.append(avg)
    else:  # Stretching the data by replicating the data points
        repetitions = bins_count // len(data)
        binned_data = []
        for val in data:
            binned_data.extend([val] * repetitions)
        # Add the remaining bins if any
        binned_data.extend([data[-1]] * (bins_count - len(binned_data)))

    min_val, max_val = min(binned_data), max(binned_data)

    if min_val == max_val:
        min_val = min_val - 1

    normalized_binned_data = [(x - min_val) / (max_val - min_val) * (len(height_chars) - 1) for x in binned_data]
    normalized_binned_data = [max(0, min(len(height_chars) - 1, round(val))) for val in normalized_binned_data]

    graph = ''.join([height_chars[int(val)] for val in normalized_binned_data])

    print(f"{description}:\n [{graph}]\n")


def test_draw_graph():
    tests = [
        ([1, 2, 3, 4, 5, 6, 7, 8], "[▁▂▃▄▅▆▇█]", 8),
        ([1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16], "[▁▂▃▄▅▆▇█]", 8),
        ([1, 2, 3], "[▁▅█]", 3),
        ([1, 2, 3], "[▁▁▅▅██]", 6),
        ([1, 2, 3], "[▁▁▁▅▅▅███]", 9),
        ([1, 1, 1], "[███████]", 7),
        ([1], "[███████]", 7),
        ([1, 2, 3, 4, 5, 6, 7, 8, 7, 6, 5, 4, 3, 2, 1], "[▁▂▃▄▅▆▇█▇▆▅▄▃▂▁]", 15),
        ([i for i in range(1,101)], "[▁▂▃▄▅▆▇█]", 8),
        ([99]+[100]*99, "[▁███████]", 8),
        ([99]+[100]*99, "[▁███████████████████████████████████████████████████████████████████████████████████████████████████]", 100),
        ([1,0,1,0], "[█████████████████████████▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁█████████████████████████▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁▁]", 100)
    ]

    for i, (data, expected, bins) in enumerate(tests, 1):
        description = f"{i}. {data}, {bins} bins. Expectation: \n {expected}\nReality"
        draw_graph(data, description, bins)


#test_draw_graph()

graphlib.py
SharkWipf/nerf_dataset_preprocessing_helper