Spaces:
Running
Running
| import cv2 | |
| import numpy as np | |
| from skimage.util import view_as_windows | |
| from skimage.exposure import match_histograms | |
| class ImgProcess: | |
| def __init__(self, img1 = None, img2 = None): | |
| self.img1 = None | |
| self.img2 = None | |
| self.img1LAB = None | |
| self.img2LAB = None | |
| self.load(img1, img2) | |
| def load(self, img1 = None, img2 = None): | |
| if img1 != None: | |
| self.img1 = cv2.cvtColor(np.array(img1), cv2.COLOR_RGB2BGR) | |
| if img2 != None: | |
| self.img2 = cv2.cvtColor(np.array(img2), cv2.COLOR_RGB2BGR) | |
| def isImageSet(self): | |
| if type(self.img1) == np.ndarray and type(self.img2) == np.ndarray: | |
| # if self.img1 != None and self.img2 != None: | |
| return True | |
| else: | |
| print(f"Error! Images not set!\nImage1Set = {self.img1 != None}, Image1Type = {type(self.img1)}\nImage2Set = {self.img2 != None}, Image2Type = {type(self.img2)}") | |
| return False | |
| def __changeGrayaxisToRGB(self): | |
| if len(self.img1.shape) == 3: | |
| if self.img1.shape[2] == 1: | |
| self.img1 = cv2.cvtColor(self.img1,cv2.COLOR_GRAY2BGR) | |
| else: | |
| self.img1 = cv2.cvtColor(self.img1,cv2.COLOR_GRAY2BGR) | |
| if len(self.img2.shape) == 3: | |
| if self.img2.shape[2] == 1: | |
| self.img2 = cv2.cvtColor(self.img2,cv2.COLOR_GRAY2BGR) | |
| else: | |
| self.img2 = cv2.cvtColor(self.img2,cv2.COLOR_GRAY2BGR) | |
| def __convertRGBToLAB(self): | |
| self.img1LAB = cv2.cvtColor(self.img1,cv2.COLOR_BGR2LAB) | |
| self.img2LAB = cv2.cvtColor(self.img2,cv2.COLOR_BGR2LAB) | |
| def __calculateLABColorMeanAndStd(self): | |
| img1mean_lst = [] | |
| img1std_lst = [] | |
| img2mean_lst = [] | |
| img2std_lst = [] | |
| for c in range(self.img1LAB.shape[2]): | |
| img1mean_lst.append(np.mean(self.img1LAB[:, :, c])) | |
| std1 = np.std(self.img1LAB[:, :, c]) | |
| img2mean_lst.append(np.mean(self.img2LAB[:, :, c])) | |
| std2 = np.std(self.img2LAB[:, :, c]) | |
| img1std_lst.append(std1 if std1 != 0 else (std2 if std2 != 0 else 1)) | |
| img2std_lst.append(std2 if std2 != 0 else 1) | |
| return (img1mean_lst, img1std_lst), (img2mean_lst, img2std_lst) | |
| def histogramMatching(self): | |
| matched = match_histograms(self.img1, self.img2, channel_axis=-1) | |
| matched_rgb = cv2.cvtColor(matched.astype(np.uint8), cv2.COLOR_BGR2RGB) | |
| return matched_rgb | |
| def transferColor(self, funPercent = 1.0): | |
| print("Transferring ...") | |
| output_img = None | |
| if self.isImageSet(): | |
| self.__changeGrayaxisToRGB() | |
| self.__convertRGBToLAB() | |
| (img1mean_lst, img1std_lst), (img2mean_lst, img2std_lst) = self.__calculateLABColorMeanAndStd() | |
| img1_trnsfrmd = [] | |
| for c in range(self.img1LAB.shape[2]): | |
| img1_star = self.img1LAB[:, :, c] - img1mean_lst[c] | |
| img1_dash = img1_star * (img2std_lst[c] / (funPercent * img1std_lst[c])) | |
| img1_trnsfrmd.append(img1_dash + img2mean_lst[c]) | |
| img1_trnsfrmd_arr = np.uint8(np.clip(np.rint(np.array(img1_trnsfrmd)), 0, 255)) | |
| img1_trnsfrmd_arr = np.moveaxis(img1_trnsfrmd_arr, 0, -1) | |
| output_img = cv2.cvtColor(img1_trnsfrmd_arr,cv2.COLOR_LAB2RGB) | |
| return output_img | |
| def transferTexture1(self, img1, img2): | |
| img2 = cv2.resize(img2, (img1.shape[1], img1.shape[0])) | |
| # Generate Gaussian pyramid for img1 | |
| g1 = img1.copy() | |
| gp1 = [g1] | |
| for i in range(6): | |
| g1 = cv2.pyrDown(g1) | |
| gp1.append(g1) | |
| # Generate Gaussian pyramid for img2 | |
| g2 = img2.copy() | |
| gp2 = [g2] | |
| for i in range(6): | |
| g2 = cv2.pyrDown(g2) | |
| gp2.append(g2) | |
| # Generate Laplacian Pyramid for img1 | |
| lp1 = [gp1[5]] | |
| for i in range(5, 0, -1): | |
| GE = cv2.pyrUp(gp1[i]) | |
| print(GE.shape, gp1[i - 1].shape) | |
| GE = cv2.resize(GE, (gp1[i-1].shape[1], gp1[i-1].shape[0])) | |
| print(GE.shape, gp1[i - 1].shape) | |
| L = cv2.subtract(gp1[i - 1], GE) | |
| lp1.append(L) | |
| # Generate Laplacian Pyramid for img2 | |
| lp2 = [gp2[5]] | |
| for i in range(5, 0, -1): | |
| GE = cv2.pyrUp(gp2[i]) | |
| GE = cv2.resize(GE, (gp2[i-1].shape[1], gp2[i-1].shape[0])) | |
| L = cv2.subtract(gp2[i - 1], GE) | |
| lp2.append(L) | |
| # Now add left and right halves of images in each level | |
| LS = [] | |
| for l1, l2 in zip(lp1, lp2): | |
| rows, cols, dpt = l1.shape | |
| ls = np.hstack((l1[:, 0:cols // 2], l2[:, cols // 2:])) | |
| LS.append(ls) | |
| # Reconstruct the image | |
| ls_ = LS[0] | |
| for i in range(1, 6): | |
| ls_ = cv2.pyrUp(ls_) | |
| ls_ = cv2.resize(ls_, (LS[i].shape[1], LS[i].shape[0])) | |
| ls_ = cv2.add(ls_, LS[i]) | |
| # Convert back to RGB for Gradio to display | |
| texture_transferred = cv2.cvtColor(ls_, cv2.COLOR_BGR2RGB) | |
| return texture_transferred | |
| def transferTexture2(self, img1, img2): | |
| img2 = cv2.resize(np.array(img2), (img1.shape[1], img1.shape[0])) | |
| # Convert both images to numpy arrays (in BGR format for OpenCV) | |
| # img1 = np.array(img1) | |
| # img2 = np.array(img2) | |
| patch_size = 24 # Patch size in pixels | |
| overlap_size = 8 | |
| texture_patches = view_as_windows(img2, (patch_size, patch_size, 3), step=patch_size - overlap_size) | |
| texture_patches = texture_patches.reshape(-1, patch_size, patch_size, 3) | |
| # Create output image | |
| output_image = np.zeros_like(img1) | |
| # Place patches in the output image by matching overlaps | |
| for i in range(0, img1.shape[0] - patch_size + 1, patch_size - overlap_size): | |
| for j in range(0, img1.shape[1] - patch_size + 1, patch_size - overlap_size): | |
| # Select a random patch to blend (for simplicity) | |
| best_patch = texture_patches[np.random.randint(len(texture_patches))] | |
| # Blend best patch into output with minimal seams | |
| output_image[i:i + patch_size, j:j + patch_size] = best_patch | |
| self.blend_patches(output_image, best_patch, i, j, overlap_size, patch_size) | |
| return output_image | |
| def loadAndTransfer(self, img1, img2, funPercent = 1.0): | |
| self.load(img1, img2) | |
| color_img1 = self.transferColor(funPercent) | |
| color_img2 = self.histogramMatching() | |
| return color_img1, color_img2 | |
| #return self.transferTexture1(color_img, self.img2) | |
| #return self.transferTexture2(color_img, self.img2) |