import concurrent.futures import numpy as np import cv2 import glob def shift_img(img, x, y): M = np.float32([[1, 0, x], [0, 1, y]]) shifted = cv2.warpAffine(img, M, (img.shape[1], img.shape[0])) return shifted def binary_diff_mask(clean, dirty, thresold=0.3): # this parts corrects gamma, and always remember, sRGB values are not in linear scale with lights intensity, clean = np.power(clean, 1/2.2) dirty = np.power(dirty, 1/2.2) # averaged_per_pixel = np.abs(dirty / shift_img(clean, 5, 0) - 1) # print(averaged_per_pixel) diff = np.abs(clean - dirty) bin_diff = (diff > thresold).astype(np.uint8) return bin_diff clean = glob.glob("data/source/Oxford_raindrop_dataset/clean/*.png") clean = sorted(clean) dirty = glob.glob("data/source/Oxford_raindrop_dataset/dirty/*.png") dirty = sorted(dirty) clean_img = cv2.imread(clean[34]) dirty_img = cv2.imread(dirty[34]) binary_diff_mask_img = binary_diff_mask(dirty_img/255, clean_img/255, thresold=0.05) k = 20 def process(i, j): print(i) clean_img_copy = shift_img(clean_img, (i-k)/4, (j-k)/4) binary_diff_mask_img = binary_diff_mask(dirty_img / 255, clean_img_copy / 255, threshold=0.3) success = cv2.imwrite(f"test/test_img_x{(i-k)/4}-y{(j-k)/4}.png", binary_diff_mask_img*255) if not success: print(f"Failed to save image at test/test_img{(i - k) / 4}-{(j - k) / 4}.png") with concurrent.futures.ProcessPoolExecutor() as executor: for i in range(k*2): for j in range(k*2): executor.submit(process, i, j)