Spaces:
Runtime error
Runtime error
| from argparse import ( | |
| ArgumentParser, | |
| Namespace, | |
| ) | |
| import os | |
| from os.path import join as pjoin | |
| from typing import Optional | |
| import sys | |
| import numpy as np | |
| import cv2 | |
| from skimage import exposure | |
| # sys.path.append('Face_Detection') | |
| # from align_warp_back_multiple_dlib import match_histograms | |
| def calculate_cdf(histogram): | |
| """ | |
| This method calculates the cumulative distribution function | |
| :param array histogram: The values of the histogram | |
| :return: normalized_cdf: The normalized cumulative distribution function | |
| :rtype: array | |
| """ | |
| # Get the cumulative sum of the elements | |
| cdf = histogram.cumsum() | |
| # Normalize the cdf | |
| normalized_cdf = cdf / float(cdf.max()) | |
| return normalized_cdf | |
| def calculate_lookup(src_cdf, ref_cdf): | |
| """ | |
| This method creates the lookup table | |
| :param array src_cdf: The cdf for the source image | |
| :param array ref_cdf: The cdf for the reference image | |
| :return: lookup_table: The lookup table | |
| :rtype: array | |
| """ | |
| lookup_table = np.zeros(256) | |
| lookup_val = 0 | |
| for src_pixel_val in range(len(src_cdf)): | |
| lookup_val | |
| for ref_pixel_val in range(len(ref_cdf)): | |
| if ref_cdf[ref_pixel_val] >= src_cdf[src_pixel_val]: | |
| lookup_val = ref_pixel_val | |
| break | |
| lookup_table[src_pixel_val] = lookup_val | |
| return lookup_table | |
| def match_histograms(src_image, ref_image, src_mask=None, ref_mask=None): | |
| """ | |
| This method matches the source image histogram to the | |
| reference signal | |
| :param image src_image: The original source image | |
| :param image ref_image: The reference image | |
| :return: image_after_matching | |
| :rtype: image (array) | |
| """ | |
| # Split the images into the different color channels | |
| # b means blue, g means green and r means red | |
| src_b, src_g, src_r = cv2.split(src_image) | |
| ref_b, ref_g, ref_r = cv2.split(ref_image) | |
| def rv(im): | |
| if ref_mask is None: | |
| return im.flatten() | |
| return im[ref_mask] | |
| def sv(im): | |
| if src_mask is None: | |
| return im.flatten() | |
| return im[src_mask] | |
| # Compute the b, g, and r histograms separately | |
| # The flatten() Numpy method returns a copy of the array c | |
| # collapsed into one dimension. | |
| src_hist_blue, bin_0 = np.histogram(sv(src_b), 256, [0, 256]) | |
| src_hist_green, bin_1 = np.histogram(sv(src_g), 256, [0, 256]) | |
| src_hist_red, bin_2 = np.histogram(sv(src_r), 256, [0, 256]) | |
| ref_hist_blue, bin_3 = np.histogram(rv(ref_b), 256, [0, 256]) | |
| ref_hist_green, bin_4 = np.histogram(rv(ref_g), 256, [0, 256]) | |
| ref_hist_red, bin_5 = np.histogram(rv(ref_r), 256, [0, 256]) | |
| # Compute the normalized cdf for the source and reference image | |
| src_cdf_blue = calculate_cdf(src_hist_blue) | |
| src_cdf_green = calculate_cdf(src_hist_green) | |
| src_cdf_red = calculate_cdf(src_hist_red) | |
| ref_cdf_blue = calculate_cdf(ref_hist_blue) | |
| ref_cdf_green = calculate_cdf(ref_hist_green) | |
| ref_cdf_red = calculate_cdf(ref_hist_red) | |
| # Make a separate lookup table for each color | |
| blue_lookup_table = calculate_lookup(src_cdf_blue, ref_cdf_blue) | |
| green_lookup_table = calculate_lookup(src_cdf_green, ref_cdf_green) | |
| red_lookup_table = calculate_lookup(src_cdf_red, ref_cdf_red) | |
| # Use the lookup function to transform the colors of the original | |
| # source image | |
| blue_after_transform = cv2.LUT(src_b, blue_lookup_table) | |
| green_after_transform = cv2.LUT(src_g, green_lookup_table) | |
| red_after_transform = cv2.LUT(src_r, red_lookup_table) | |
| # Put the image back together | |
| image_after_matching = cv2.merge([blue_after_transform, green_after_transform, red_after_transform]) | |
| image_after_matching = cv2.convertScaleAbs(image_after_matching) | |
| return image_after_matching | |
| def convert_to_BW(im, mode): | |
| if mode == "b": | |
| gray = im[..., 0] | |
| elif mode == "gb": | |
| gray = (im[..., 0].astype(float) + im[..., 1]) / 2.0 | |
| else: | |
| gray = cv2.cvtColor(im, cv2.COLOR_BGR2GRAY) | |
| gray = gray.astype(np.uint8) | |
| return np.stack([gray] * 3, axis=-1) | |
| def parse_args(args=None, namespace: Optional[Namespace] = None): | |
| parser = ArgumentParser('match histogram of src to ref') | |
| parser.add_argument('src') | |
| parser.add_argument('ref') | |
| parser.add_argument('--out', default=None, help="converted src that matches ref") | |
| parser.add_argument('--src_mask', default=None, help="mask on which to match the histogram") | |
| parser.add_argument('--ref_mask', default=None, help="mask on which to match the histogram") | |
| parser.add_argument('--spectral_sensitivity', choices=['b', 'gb', 'g'], help="match the histogram of corresponding sensitive channel(s)") | |
| parser.add_argument('--crop', type=int, default=0, help="crop the boundary to match") | |
| return parser.parse_args(args=args, namespace=namespace) | |
| def main(args): | |
| A = cv2.imread(args.ref) | |
| A = convert_to_BW(A, args.spectral_sensitivity) | |
| B = cv2.imread(args.src, 0) | |
| B = np.stack((B,) * 3, axis=-1) | |
| mask_A = cv2.resize(cv2.imread(args.ref_mask, 0), A.shape[:2][::-1], | |
| interpolation=cv2.INTER_NEAREST) > 0 if args.ref_mask else None | |
| mask_B = cv2.resize(cv2.imread(args.src_mask, 0), B.shape[:2][::-1], | |
| interpolation=cv2.INTER_NEAREST) > 0 if args.src_mask else None | |
| if args.crop > 0: | |
| c = args.crop | |
| bc = int(c / A.shape[0] * B.shape[0] + 0.5) | |
| A = A[c:-c, c:-c] | |
| B = B[bc:-bc, bc:-bc] | |
| B = match_histograms(B, A, src_mask=mask_B, ref_mask=mask_A) | |
| # B = exposure.match_histograms(B, A, multichannel=True) | |
| if args.out: | |
| os.makedirs(os.path.dirname(args.out), exist_ok=True) | |
| cv2.imwrite(args.out, B) | |
| return B | |
| if __name__ == "__main__": | |
| main(parse_args()) | |