jo-chemla/kapture-cropper.py

## kapture-cropper.py
# py kapture-cropper.py -i dataset-kapture\ --border_px 0 --scale_factor 1 -v

import kapture
import kapture.io.csv as csv
from PIL import Image
from kapture.io.csv import kapture_to_dir
import os, logging, argparse
import kapture.utils.logging

logger = logging.getLogger("kapture-cropper")

DEFAULT_FOCAL_LENGTH_FACTOR = 1.2
EPSILON = 1e-5


def crop_command_line() -> None:
    """
    Crop a set of raster images from a kapture project, adapting their intrinsics using the parameters given on the command line.
    """
    parser = argparse.ArgumentParser(
        description="Crop a set of raster images from a kapture project, adapting their intrinsics"
    )
    parser.add_argument(
        "-i", "--input", type=str, required=True, help="input path to kapture project"
    )
    parser.add_argument(
        "-b",
        "--border_px",
        type=int,
        required=False,
        default=0,
        help="Crop amount in pixels, for each side the same",
    )
    parser.add_argument(
        "-s",
        "--scale_factor",
        type=check_positive_strictly,
        required=False,
        default=1,
        help="Rescale factor, to change the coordinate system of cams and points3d",
    )
    # Logging
    parser_verbosity = parser.add_mutually_exclusive_group()
    parser_verbosity.add_argument(
        "-v",
        "--verbose",
        nargs="?",
        default=logging.WARNING,
        const=logging.INFO,
        action=kapture.utils.logging.VerbosityParser,
        help="verbosity level (debug, info, warning, critical, ... or int value) [warning]",
    )
    parser_verbosity.add_argument(
        "-q",
        "--silent",
        "--quiet",
        action="store_const",
        dest="verbose",
        const=logging.CRITICAL,
    )

    # Parse args and do the crop
    args = parser.parse_args()
    logger.setLevel(args.verbose)
    if args.verbose <= logging.DEBUG:  # also let kapture express its logs
        kapture.utils.logging.getLogger().setLevel(args.verbose)
    kapture_crop(args.input, args.border_px, args.scale_factor)


def check_positive_strictly(value):
    try:
        value = float(value)
        if value <= 0:
            raise argparse.ArgumentTypeError(
                "{} is not a positive number".format(value)
            )
    except ValueError:
        raise Exception("{} is not an integer".format(value))
    return value


def rescale_kapture_coords(kapture_data, scale_factor):
    # Rescale sensors' trajectories
    trajectories = kapture_data.trajectories
    for traj_idx in trajectories:
        for poseTransform in trajectories[traj_idx].values():
            poseTransform.rescale(scale_factor)
    # rescale points3d observations
    points3d = kapture_data.points3d
    points3d[:, :3] *= scale_factor


def kapture_crop(kapture_dir_path, border_px=0, scale_factor=1):
    # Create dir and move other files
    if border_px != 0:
        image_dir_path = os.path.join(kapture_dir_path, "sensors/records_data")
        original_dir_path = os.path.join(
            kapture_dir_path, "sensors/records_data_original"
        )
        if os.path.exists(original_dir_path):
            logger.error(
                f"{original_dir_path} already exists, you probably already applied the cropping"
            )
            return
        os.rename(image_dir_path, original_dir_path)
        os.makedirs(image_dir_path, exist_ok=True)
    #
    # with csv.get_all_tar_handlers(kapture_dir_path) as tar_handlers:
    pairsfile_path = None
    logger.info(f"Reading Kapture project from {kapture_dir_path}...")
    kapture_data = csv.kapture_from_dir(
        kapture_dir_path,
        pairsfile_path,
        # tar_handlers=tar_handlers
    )

    records_camera = kapture_data.records_camera
    sensors = kapture_data.sensors
    logger.info(
        f"Found {len(records_camera)} records_camera and {len(sensors)} sensors"
    )
    #
    # Could do the conversion in place if no crop via PILLOW
    # converted_sensors = [
    #     (cam_id, get_cropped_sensor(cam, border_px)) for cam_id, cam in sensors.items()
    # ]
    # cropped_kapture = kapture.Kapture(sensors=converted_sensors, records_camera=records_camera)
    #
    # Loop for image crop via pillow + sensor crop via kapture
    logger.info("Cropping image rasters + computing cropped kapture sensors...")
    # handler = logging.StreamHandler()
    for handler in logger.handlers:
        handler.terminator = "\r"
    for cam_id in records_camera:  # cam = records_camera[0]
        cam = records_camera[cam_id]
        for (
            sensor_id,
            img_fp,
        ) in cam.items():  # sensor_id, img_fp = list(cam.items())[0]
            sensor = sensors[sensor_id]
            logger.info(
                f"Cropping image with id {sensor_id} and fp {img_fp}, sensor {sensor}",
                # end="\r",
            )
            # Always simplify sensor model if radial distortion coefficients are zero
            cropped_sensor = simplify_sensors_model(sensor)
            # Crop sensor if user asks for it
            if border_px != 0:
                cropped_sensor = get_cropped_sensor(cropped_sensor, border_px)
                # Crop via Pillow
                with Image.open(os.path.join(original_dir_path, img_fp)) as im:
                    width, height = im.size
                    left, top, right, bottom = (
                        border_px,
                        border_px,
                        width - border_px,
                        height - border_px,
                    )
                    cropped_im = im.crop((left, top, right, bottom))
                    cropped_im.save(os.path.join(image_dir_path, img_fp))
            #
            # write back that cropped sensor to the current kapture db
            sensors[sensor_id] = cropped_sensor
    # Rescale coordinate system
    if scale_factor != 1:
        rescale_kapture_coords(kapture_data, scale_factor)
    # Export to disk
    for handler in logger.handlers:
        handler.terminator = "\n"
    logger.info("\nImage raster files cropped, Writing cropped data to Kapture file...")
    cropped_kapture = kapture.Kapture(
        sensors=sensors,
        records_camera=records_camera,
        trajectories=kapture_data.trajectories,
        points3d=kapture_data.points3d,
    )
    kapture_to_dir(kapture_dir_path, cropped_kapture)


# Adapted from get_colmap_camera in https://github.com/naver/kapture/blob/e58e244f35fe8db47dbb2b149178456f513ef6f8/kapture/converter/colmap/cameras.py#L42
def get_cropped_sensor(camera: kapture.Camera, border: int):
    """
    Compute the cropped camera definition - uniform border width given in pixels
    #
    :param camera: a kapture camera definition
    :param border: a pixel crop count
    :return: cropped camera parameters.
    """
    assert isinstance(camera, kapture.Camera)
    assert len(camera.camera_params) >= 2
    #
    old_width = camera.camera_params[0]
    old_height = camera.camera_params[1]
    #
    # Update width with cropped value (in pixels units)
    width = old_width - 2 * border
    height = old_height - 2 * border
    # will apply to cx,cy and focal-length
    if camera.camera_type in [
        kapture.CameraType.SIMPLE_PINHOLE,
        kapture.CameraType.SIMPLE_RADIAL,
        kapture.CameraType.RADIAL,
    ]:
        # [SIMPLE_]RADIAL params: w, h, f, cx, cy, k1 [, k2]
        params = camera.camera_params[2:]
        params[0] *= width / old_width  # focal f
        params[1] -= border  # cx
        params[2] -= border  # cy
    elif camera.camera_type in [
        kapture.CameraType.PINHOLE,
        kapture.CameraType.OPENCV,
        kapture.CameraType.FULL_OPENCV,
    ]:
        # PINHOLE/OPENCV params: w, h, fx, fy, cx, cy, k1 [, k2]
        params = camera.camera_params[2:]
        params[0] *= width / old_width  # focal f
        params[1] *= width / old_width  # focal f
        params[2] -= border  # cx
        params[3] -= border  # cy
    else:
        raise ValueError(
            f"This sensor model: {camera.camera_type} is not supported by the intrinsics img cropper yet"
        )
    return kapture.Camera(camera.camera_type, [width, height, *params])


# The following table stores distortion_idx and camera_type for a given input camera_type for the simplified version
# { input_camera_type: (distortion_idx, simplified_camera_type) }
# [SIMPLE_]RADIAL params: w, h, f, cx, cy | k1 [, k2]
# [FULL_]OPENCV params: w, h, fx, fy, cx, cy | k1, k2, p1, p2 [, k3, k4, k5, k6]
#
# See https://github.com/colmap/colmap/blob/main/src/colmap/sensor/models.h
# UNKNOWN_CAMERA w, h
# SIMPLE_PINHOLE w, h, f, cx, cy
# PINHOLE w, h, fx, fy, cx, cy
# SIMPLE_RADIAL  w, h, f, cx, cy, k
# RADIAL w, h, f, cx, cy, k1, k2
# OPENCV w, h, fx, fy, cx, cy, k1, k2, p1, p2
# FULL_OPENCV w, h, fx, fy, cx, cy, k1, k2, p1, p2, k3, k4, k5, k6
SENSOR_SIMPLIFICATION_TABLE = {
    kapture.CameraType.SIMPLE_RADIAL: (5, kapture.CameraType.SIMPLE_PINHOLE),
    kapture.CameraType.RADIAL: (5, kapture.CameraType.SIMPLE_PINHOLE),
    kapture.CameraType.OPENCV: (6, kapture.CameraType.PINHOLE),
    kapture.CameraType.FULL_OPENCV: (6, kapture.CameraType.PINHOLE),
}


def simplify_sensors_model(camera: kapture.Camera):
    """
    Convert the sensor to pinhole if it is a radial which has zero k1 k2 distortion coefficients
    #
    :param camera: a kapture camera definition
    :return: cropped camera parameters.
    """
    assert isinstance(camera, kapture.Camera)
    assert len(camera.camera_params) >= 2
    #
    if camera.camera_type in SENSOR_SIMPLIFICATION_TABLE:
        distortion_idx, camera_type = SENSOR_SIMPLIFICATION_TABLE[camera.camera_type]
    else:
        raise ValueError(
            f"This sensor model: {camera.camera_type} is not supported by the sensor simplifier yet"
        )
    # Distangle pinhole params from distortion params, check if distortion is zero, otherwise return original camera
    pinhole_params = camera.camera_params[:distortion_idx]  # w, h, f[x, fy], cx, cy
    distortion_params = camera.camera_params[distortion_idx:]  # k...
    if sum([x**2 for x in distortion_params]) <= EPSILON:
        camera_params = pinhole_params
        return kapture.Camera(camera_type, camera_params)
    else:
        return camera


if __name__ == "__main__":
    crop_command_line()


"""
:: Export from RC an imagelist.lst + bundle.out
:: Convert RC imagelist.lst to imagelist-local.lst
:: Install Kapture
pip install kapture
:: Bundler -> Kapture
py C:\Python310\Scripts\kapture_import_bundler.py -v debug -i dataset-bundler\bundle.out
  -l dataset-bundler\imagelist-local.lst -im dataset-bundler\images
  --image_transfer link_absolute -o dataset-kapture --add-reconstruction
:: Kapture crop remove black strips
py kapture-cropper.py -v info -i dataset-kapture\ --border_px 10
:: Kapture -> Colmap
py C:\Python310\Scripts\kapture_export_colmap.py -v debug -f -i dataset-kapture
  -db dataset-colmap\colmap.db --reconstruction dataset-colmap\reconstruction-txt
:: Optional Colmap txt to bin conversion
mkdir dataset-colmap\sparse\0
COLMAP.bat model_converter --input_path dataset-colmap\reconstruction-txt
  --output_path dataset-colmap\sparse\0 --output_type BIN
:: Convert cameras from RADIAL to PINHOLE NOT NEEDED ANYMORE
:: xcopy dataset-bundler\images dataset-colmap\images /i
:: curl https://raw.githubusercontent.com/graphdeco-inria/gaussian-splatting/main/convert.py -o gaussian-splatting-pinhole-cameras-convert.py
:: py gaussian-splatting-pinhole-cameras-convert.py -s dataset-colmap --colmap_executable ..\COLMAP-3.8-windows-cuda
:: [--resize] #If not resizing, ImageMagick is not needed
"""

# Convert RC imagelist.lst to imagelist-local.lst
"""import os
def keep_filename(imagelist):
  with open(imagelist) as f:
    filenames = [os.path.split(fp)[1] for fp in f.readlines()]
    with open(imagelist[:-4] + '-local.lst', 'w') as f_out:
      f_out.write(''.join(filenames))
keep_filename('dataset-bundler\\imagelist.lst')
"""
	# py kapture-cropper.py -i dataset-kapture\ --border_px 0 --scale_factor 1 -v

	import kapture
	import kapture.io.csv as csv
	from PIL import Image
	from kapture.io.csv import kapture_to_dir
	import os, logging, argparse
	import kapture.utils.logging

	logger = logging.getLogger("kapture-cropper")

	DEFAULT_FOCAL_LENGTH_FACTOR = 1.2
	EPSILON = 1e-5


	def crop_command_line() -> None:
	"""
	Crop a set of raster images from a kapture project, adapting their intrinsics using the parameters given on the command line.
	"""
	parser = argparse.ArgumentParser(
	description="Crop a set of raster images from a kapture project, adapting their intrinsics"
	)
	parser.add_argument(
	"-i", "--input", type=str, required=True, help="input path to kapture project"
	)
	parser.add_argument(
	"-b",
	"--border_px",
	type=int,
	required=False,
	default=0,
	help="Crop amount in pixels, for each side the same",
	)
	parser.add_argument(
	"-s",
	"--scale_factor",
	type=check_positive_strictly,
	required=False,
	default=1,
	help="Rescale factor, to change the coordinate system of cams and points3d",
	)
	# Logging
	parser_verbosity = parser.add_mutually_exclusive_group()
	parser_verbosity.add_argument(
	"-v",
	"--verbose",
	nargs="?",
	default=logging.WARNING,
	const=logging.INFO,
	action=kapture.utils.logging.VerbosityParser,
	help="verbosity level (debug, info, warning, critical, ... or int value) [warning]",
	)
	parser_verbosity.add_argument(
	"-q",
	"--silent",
	"--quiet",
	action="store_const",
	dest="verbose",
	const=logging.CRITICAL,
	)

	# Parse args and do the crop
	args = parser.parse_args()
	logger.setLevel(args.verbose)
	if args.verbose <= logging.DEBUG: # also let kapture express its logs
	kapture.utils.logging.getLogger().setLevel(args.verbose)
	kapture_crop(args.input, args.border_px, args.scale_factor)


	def check_positive_strictly(value):
	try:
	value = float(value)
	if value <= 0:
	raise argparse.ArgumentTypeError(
	"{} is not a positive number".format(value)
	)
	except ValueError:
	raise Exception("{} is not an integer".format(value))
	return value


	def rescale_kapture_coords(kapture_data, scale_factor):
	# Rescale sensors' trajectories
	trajectories = kapture_data.trajectories
	for traj_idx in trajectories:
	for poseTransform in trajectories[traj_idx].values():
	poseTransform.rescale(scale_factor)
	# rescale points3d observations
	points3d = kapture_data.points3d
	points3d[:, :3] *= scale_factor


	def kapture_crop(kapture_dir_path, border_px=0, scale_factor=1):
	# Create dir and move other files
	if border_px != 0:
	image_dir_path = os.path.join(kapture_dir_path, "sensors/records_data")
	original_dir_path = os.path.join(
	kapture_dir_path, "sensors/records_data_original"
	)
	if os.path.exists(original_dir_path):
	logger.error(
	f"{original_dir_path} already exists, you probably already applied the cropping"
	)
	return
	os.rename(image_dir_path, original_dir_path)
	os.makedirs(image_dir_path, exist_ok=True)
	#
	# with csv.get_all_tar_handlers(kapture_dir_path) as tar_handlers:
	pairsfile_path = None
	logger.info(f"Reading Kapture project from {kapture_dir_path}...")
	kapture_data = csv.kapture_from_dir(
	kapture_dir_path,
	pairsfile_path,
	# tar_handlers=tar_handlers
	)

	records_camera = kapture_data.records_camera
	sensors = kapture_data.sensors
	logger.info(
	f"Found {len(records_camera)} records_camera and {len(sensors)} sensors"
	)
	#
	# Could do the conversion in place if no crop via PILLOW
	# converted_sensors = [
	# (cam_id, get_cropped_sensor(cam, border_px)) for cam_id, cam in sensors.items()
	# ]
	# cropped_kapture = kapture.Kapture(sensors=converted_sensors, records_camera=records_camera)
	#
	# Loop for image crop via pillow + sensor crop via kapture
	logger.info("Cropping image rasters + computing cropped kapture sensors...")
	# handler = logging.StreamHandler()
	for handler in logger.handlers:
	handler.terminator = "\r"
	for cam_id in records_camera: # cam = records_camera[0]
	cam = records_camera[cam_id]
	for (
	sensor_id,
	img_fp,
	) in cam.items(): # sensor_id, img_fp = list(cam.items())[0]
	sensor = sensors[sensor_id]
	logger.info(
	f"Cropping image with id {sensor_id} and fp {img_fp}, sensor {sensor}",
	# end="\r",
	)
	# Always simplify sensor model if radial distortion coefficients are zero
	cropped_sensor = simplify_sensors_model(sensor)
	# Crop sensor if user asks for it
	if border_px != 0:
	cropped_sensor = get_cropped_sensor(cropped_sensor, border_px)
	# Crop via Pillow
	with Image.open(os.path.join(original_dir_path, img_fp)) as im:
	width, height = im.size
	left, top, right, bottom = (
	border_px,
	border_px,
	width - border_px,
	height - border_px,
	)
	cropped_im = im.crop((left, top, right, bottom))
	cropped_im.save(os.path.join(image_dir_path, img_fp))
	#
	# write back that cropped sensor to the current kapture db
	sensors[sensor_id] = cropped_sensor
	# Rescale coordinate system
	if scale_factor != 1:
	rescale_kapture_coords(kapture_data, scale_factor)
	# Export to disk
	for handler in logger.handlers:
	handler.terminator = "\n"
	logger.info("\nImage raster files cropped, Writing cropped data to Kapture file...")
	cropped_kapture = kapture.Kapture(
	sensors=sensors,
	records_camera=records_camera,
	trajectories=kapture_data.trajectories,
	points3d=kapture_data.points3d,
	)
	kapture_to_dir(kapture_dir_path, cropped_kapture)


	# Adapted from get_colmap_camera in https://github.com/naver/kapture/blob/e58e244f35fe8db47dbb2b149178456f513ef6f8/kapture/converter/colmap/cameras.py#L42
	def get_cropped_sensor(camera: kapture.Camera, border: int):
	"""
	Compute the cropped camera definition - uniform border width given in pixels
	#
	:param camera: a kapture camera definition
	:param border: a pixel crop count
	:return: cropped camera parameters.
	"""
	assert isinstance(camera, kapture.Camera)
	assert len(camera.camera_params) >= 2
	#
	old_width = camera.camera_params[0]
	old_height = camera.camera_params[1]
	#
	# Update width with cropped value (in pixels units)
	width = old_width - 2 * border
	height = old_height - 2 * border
	# will apply to cx,cy and focal-length
	if camera.camera_type in [
	kapture.CameraType.SIMPLE_PINHOLE,
	kapture.CameraType.SIMPLE_RADIAL,
	kapture.CameraType.RADIAL,
	]:
	# [SIMPLE_]RADIAL params: w, h, f, cx, cy, k1 [, k2]
	params = camera.camera_params[2:]
	params[0] *= width / old_width # focal f
	params[1] -= border # cx
	params[2] -= border # cy
	elif camera.camera_type in [
	kapture.CameraType.PINHOLE,
	kapture.CameraType.OPENCV,
	kapture.CameraType.FULL_OPENCV,
	]:
	# PINHOLE/OPENCV params: w, h, fx, fy, cx, cy, k1 [, k2]
	params = camera.camera_params[2:]
	params[0] *= width / old_width # focal f
	params[1] *= width / old_width # focal f
	params[2] -= border # cx
	params[3] -= border # cy
	else:
	raise ValueError(
	f"This sensor model: {camera.camera_type} is not supported by the intrinsics img cropper yet"
	)
	return kapture.Camera(camera.camera_type, [width, height, *params])


	# The following table stores distortion_idx and camera_type for a given input camera_type for the simplified version
	# { input_camera_type: (distortion_idx, simplified_camera_type) }
	# [SIMPLE_]RADIAL params: w, h, f, cx, cy \| k1 [, k2]
	# [FULL_]OPENCV params: w, h, fx, fy, cx, cy \| k1, k2, p1, p2 [, k3, k4, k5, k6]
	#
	# See https://github.com/colmap/colmap/blob/main/src/colmap/sensor/models.h
	# UNKNOWN_CAMERA w, h
	# SIMPLE_PINHOLE w, h, f, cx, cy
	# PINHOLE w, h, fx, fy, cx, cy
	# SIMPLE_RADIAL w, h, f, cx, cy, k
	# RADIAL w, h, f, cx, cy, k1, k2
	# OPENCV w, h, fx, fy, cx, cy, k1, k2, p1, p2
	# FULL_OPENCV w, h, fx, fy, cx, cy, k1, k2, p1, p2, k3, k4, k5, k6
	SENSOR_SIMPLIFICATION_TABLE = {
	kapture.CameraType.SIMPLE_RADIAL: (5, kapture.CameraType.SIMPLE_PINHOLE),
	kapture.CameraType.RADIAL: (5, kapture.CameraType.SIMPLE_PINHOLE),
	kapture.CameraType.OPENCV: (6, kapture.CameraType.PINHOLE),
	kapture.CameraType.FULL_OPENCV: (6, kapture.CameraType.PINHOLE),
	}


	def simplify_sensors_model(camera: kapture.Camera):
	"""
	Convert the sensor to pinhole if it is a radial which has zero k1 k2 distortion coefficients
	#
	:param camera: a kapture camera definition
	:return: cropped camera parameters.
	"""
	assert isinstance(camera, kapture.Camera)
	assert len(camera.camera_params) >= 2
	#
	if camera.camera_type in SENSOR_SIMPLIFICATION_TABLE:
	distortion_idx, camera_type = SENSOR_SIMPLIFICATION_TABLE[camera.camera_type]
	else:
	raise ValueError(
	f"This sensor model: {camera.camera_type} is not supported by the sensor simplifier yet"
	)
	# Distangle pinhole params from distortion params, check if distortion is zero, otherwise return original camera
	pinhole_params = camera.camera_params[:distortion_idx] # w, h, f[x, fy], cx, cy
	distortion_params = camera.camera_params[distortion_idx:] # k...
	if sum([x**2 for x in distortion_params]) <= EPSILON:
	camera_params = pinhole_params
	return kapture.Camera(camera_type, camera_params)
	else:
	return camera


	if __name__ == "__main__":
	crop_command_line()


	"""
	:: Export from RC an imagelist.lst + bundle.out
	:: Convert RC imagelist.lst to imagelist-local.lst
	:: Install Kapture
	pip install kapture
	:: Bundler -> Kapture
	py C:\Python310\Scripts\kapture_import_bundler.py -v debug -i dataset-bundler\bundle.out
	-l dataset-bundler\imagelist-local.lst -im dataset-bundler\images
	--image_transfer link_absolute -o dataset-kapture --add-reconstruction
	:: Kapture crop remove black strips
	py kapture-cropper.py -v info -i dataset-kapture\ --border_px 10
	:: Kapture -> Colmap
	py C:\Python310\Scripts\kapture_export_colmap.py -v debug -f -i dataset-kapture
	-db dataset-colmap\colmap.db --reconstruction dataset-colmap\reconstruction-txt
	:: Optional Colmap txt to bin conversion
	mkdir dataset-colmap\sparse\0
	COLMAP.bat model_converter --input_path dataset-colmap\reconstruction-txt
	--output_path dataset-colmap\sparse\0 --output_type BIN
	:: Convert cameras from RADIAL to PINHOLE NOT NEEDED ANYMORE
	:: xcopy dataset-bundler\images dataset-colmap\images /i
	:: curl https://raw.githubusercontent.com/graphdeco-inria/gaussian-splatting/main/convert.py -o gaussian-splatting-pinhole-cameras-convert.py
	:: py gaussian-splatting-pinhole-cameras-convert.py -s dataset-colmap --colmap_executable ..\COLMAP-3.8-windows-cuda
	:: [--resize] #If not resizing, ImageMagick is not needed
	"""

	# Convert RC imagelist.lst to imagelist-local.lst
	"""import os
	def keep_filename(imagelist):
	with open(imagelist) as f:
	filenames = [os.path.split(fp)[1] for fp in f.readlines()]
	with open(imagelist[:-4] + '-local.lst', 'w') as f_out:
	f_out.write(''.join(filenames))
	keep_filename('dataset-bundler\\imagelist.lst')
	"""