samples/python/snippets/chromatic_aberration_correction.py#

  • An example correcting chromatic aberration with Python

 1#!/usr/bin/env python3
 2# This file is part of OpenCV project.
 3# It is subject to the license terms in the LICENSE file found in the top-level directory
 4# of this distribution and at http://opencv.org/license.html
 5
 6import argparse
 7import sys
 8import cv2 as cv
 9
10USAGE = """\
11Chromatic Aberration Correction Sample
12Usage:
13  chromatic_aberration_correction.py <input_image> <calibration_file> [--bayer <code>] [--output <path>]
14
15Arguments:
16  input_image       Path to the input image. Can be:
17                      • a 3-channel BGR image, or
18                      • a 1-channel raw Bayer image (see bayer_pattern)
19  calibration_file  OpenCV YAML/XML file with chromatic aberration calibration:
20                      image_width, image_height, red_channel/coeffs_x, coeffs_y,
21                      blue_channel/coeffs_x, coeffs_y.
22  output            (optional) Path to save the corrected image. Default: corrected.png
23  bayer             (optional) integer code for demosaicing a 1-channel raw image
24                    If omitted or <0, input is assumed 3-channel BGR.
25
26Example:
27  python chromatic_aberration_correction.py input.png calib.yaml --bayer 46 --output corrected.png
28"""
29
30def main(argv=None):
31    parser = argparse.ArgumentParser(
32        description="Chromatic Aberration Correction Sample",
33        formatter_class=argparse.RawDescriptionHelpFormatter,
34        epilog=USAGE
35    )
36    parser.add_argument("input", help="Input image (BGR or Bayer)")
37    parser.add_argument("calibration", help="Calibration file (YAML/XML)")
38    parser.add_argument("--output", default="corrected.png", help="Output image file")
39    parser.add_argument("--bayer", type=int, default=-1, help="Bayer pattern code for demosaic")
40    parser.add_argument("--no-gui", action="store_true", help="Do not open image windows")
41
42    args = parser.parse_args(argv)
43
44    img = cv.imread(args.input, cv.IMREAD_UNCHANGED)
45    if img is None:
46        print(f"ERROR: Could not load input image: {args.input}", file=sys.stderr)
47        return 1
48
49    fs = cv.FileStorage(args.calibration, cv.FileStorage_READ)
50    if not fs.isOpened():
51        print(f"Could not calibration coefficients from {args.calibration}")
52        return 1
53
54    try:
55        coeffMat, size, degree = cv.loadChromaticAberrationParams(fs.root())
56        corrected = cv.correctChromaticAberration(img, coeffMat, size, degree, args.bayer)
57
58        if corrected is None:
59            print("ERROR: cv.correctChromaticAberration returned None", file=sys.stderr)
60            return 1
61
62        if not args.no_gui:
63            cv.namedWindow("Original", cv.WINDOW_AUTOSIZE)
64            cv.namedWindow("Corrected", cv.WINDOW_AUTOSIZE)
65            cv.imshow("Original", img)
66            cv.imshow("Corrected", corrected)
67            print("Press any key to continue...")
68            cv.waitKey(0)
69            cv.destroyAllWindows()
70
71        if not cv.imwrite(args.output, corrected):
72            print(f"WARNING: Could not write output image: {args.output}", file=sys.stderr)
73        else:
74            print(f"Saved corrected image to: {args.output}")
75
76    except cv.error as e:
77        print(f"OpenCV error: {e}", file=sys.stderr)
78        return 1
79
80    return 0
81
82if __name__ == "__main__":
83    sys.exit(main())