ColorCorrectionML is a Python package for color correction of sRGB images using machine learning. It uses ML regression methods (linear, least sqaure, and partial least squares regression) to learn the color correction function from a training image with a color checker. The learned function is then applied to correct the color of a test image.
pip install colorcorrectionML
from ColorCorrectionML import ColorCorrectionML
import cv2
img = cv2.imread('Images/img2.png')
cc = ColorCorrectionML(img, chart='Classic', illuminant='D50')
method = 'pls' # 'linear', 'lstsq', 'pls'
# for linear regression, least square regression, and partial least square regression respectively
show = True
kwargs = {
'method': method,
'degree': 3, # degree of polynomial
'interactions_only': False, # only interactions terms,
'ncomp': 10, # number of components for PLS only
'max_iter': 5000, # max iterations for PLS only
'white_balance_mtd': 0 # 0: no white balance, 1: learningBasedWB, 2: simpleWB, 3: grayWorldWB,
}
M, patch_size = cc.compute_correction(
show=show,
**kwargs
)
# resize img by 2
# img = cv2.resize(img, (0,0), fx=0.3, fy=0.3, interpolation=cv2.INTER_AREA)
img_corr = cc.correct_img(img, show=True)
# img_corr = cc.Parallel_correct_img(img, chunks_=50000, show=True)
Scatter plot of the original image and the corrected image color values
- Add other reference color values (D55, D65, D70, D75)
- Add other color charts (ColorChecker24, ColorCheckerSG, ColorCheckerDC)
- Add other color spaces (CIELab, XYZ, etc.)
- Add other regression methods (Ridge, Lasso, ElasticNet, etc.)
- Refine the white balance methods