Spaces:
Sleeping
Sleeping
| #imports | |
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| import os | |
| from math import floor, ceil | |
| from random import randint | |
| from sklearn.neighbors import KDTree | |
| from skimage.util.shape import view_as_windows | |
| from skimage import io | |
| from PIL import Image, ImageDraw | |
| from IPython.display import clear_output | |
| class patchBasedTextureSynthesis: | |
| def __init__(self, exampleMapPath, in_outputPath, in_outputSize, in_patchSize, in_overlapSize, in_windowStep = 5, in_mirror_hor = True, in_mirror_vert = True, in_shapshots = True): | |
| self.exampleMap = self.loadExampleMap(exampleMapPath) | |
| self.snapshots = in_shapshots | |
| self.outputPath = in_outputPath | |
| self.outputSize = in_outputSize | |
| self.patchSize = in_patchSize | |
| self.overlapSize = in_overlapSize | |
| self.mirror_hor = in_mirror_hor | |
| self.mirror_vert = in_mirror_vert | |
| self.total_patches_count = 0 #excluding mirrored versions | |
| self.windowStep = 5 | |
| self.iter = 0 | |
| self.checkIfDirectoryExists() #check if output directory exists | |
| self.examplePatches = self.prepareExamplePatches() | |
| self.canvas, self.filledMap, self.idMap = self.initCanvas() | |
| self.initFirstPatch() #place random block to start with | |
| self.kdtree_topOverlap, self.kdtree_leftOverlap, self.kdtree_combined = self.initKDtrees() | |
| self.PARM_truncation = 0.8 | |
| self.PARM_attenuation = 2 | |
| def checkIfDirectoryExists(self): | |
| if not os.path.exists(self.outputPath): | |
| os.makedirs(self.outputPath) | |
| def resolveAll(self): | |
| self.saveParams() | |
| #resolve all unresolved patches | |
| for i in range(np.sum(1-self.filledMap).astype(int)): | |
| self.resolveNext() | |
| if not self.snapshots: | |
| img = Image.fromarray(np.uint8(self.canvas*255)) | |
| img = img.resize((self.outputSize[0], self.outputSize[1]), resample=0, box=None) | |
| img.save(self.outputPath + 'out.jpg') | |
| # else: | |
| # self.visualize([0,0], [], [], showCandidates=False) | |
| return img | |
| def saveParams(self): | |
| #write | |
| text_file = open(self.outputPath + 'params.txt', "w") | |
| text_file.write("PatchSize: %d \nOverlapSize: %d \nMirror Vert: %d \nMirror Hor: %d" % (self.patchSize, self.overlapSize, self.mirror_vert, self.mirror_hor)) | |
| text_file.close() | |
| def resolveNext(self): | |
| #coordinate of the next one to resolve | |
| coord = self.idCoordTo2DCoord(np.sum(self.filledMap), np.shape(self.filledMap)) #get 2D coordinate of next to resolve patch | |
| #get overlap areas of the patch we want to resolve | |
| overlapArea_Top = self.getOverlapAreaTop(coord) | |
| overlapArea_Left = self.getOverlapAreaLeft(coord) | |
| #find most similar patch from the examples | |
| dist, ind = self.findMostSimilarPatches(overlapArea_Top, overlapArea_Left, coord) | |
| if self.mirror_hor or self.mirror_vert: | |
| #check that top and left neighbours are not mirrors | |
| dist, ind = self.checkForMirrors(dist, ind, coord) | |
| #choose random valid patch | |
| probabilities = self.distances2probability(dist, self.PARM_truncation, self.PARM_attenuation) | |
| chosenPatchId = np.random.choice(ind, 1, p=probabilities) | |
| #update canvas | |
| blend_top = (overlapArea_Top is not None) | |
| blend_left = (overlapArea_Left is not None) | |
| self.updateCanvas(chosenPatchId, coord[0], coord[1], blend_top, blend_left) | |
| #update filledMap and id map ;) | |
| self.filledMap[coord[0], coord[1]] = 1 | |
| self.idMap[coord[0], coord[1]] = chosenPatchId | |
| #visualize | |
| # self.visualize(coord, chosenPatchId, ind) | |
| self.iter += 1 | |
| def visualize(self, coord, chosenPatchId, nonchosenPatchId, showCandidates = True): | |
| #full visualization includes both example and generated img | |
| canvasSize = np.shape(self.canvas) | |
| #insert generated image | |
| vis = np.zeros((canvasSize[0], canvasSize[1] * 2, 3)) + 0.2 | |
| vis[:, 0:canvasSize[1]] = self.canvas | |
| #insert example | |
| exampleHighlited = np.copy(self.exampleMap) | |
| if showCandidates: | |
| exampleHighlited = self.hightlightPatchCandidates(chosenPatchId, nonchosenPatchId) | |
| h = floor(canvasSize[0] / 2) | |
| w = floor(canvasSize[1] / 2) | |
| exampleResized = self.resize(exampleHighlited, [h, w]) | |
| offset_h = floor((canvasSize[0] - h) / 2) | |
| offset_w = floor((canvasSize[1] - w) / 2) | |
| vis[offset_h:offset_h+h, canvasSize[1]+offset_w:canvasSize[1]+offset_w+w] = exampleResized | |
| #show live update | |
| plt.imshow(vis) | |
| clear_output(wait=True) | |
| display(plt.show()) | |
| if self.snapshots: | |
| img = Image.fromarray(np.uint8(vis*255)) | |
| img = img.resize((self.outputSize[0]*2, self.outputSize[1]), resample=0, box=None) | |
| img.save(self.outputPath + 'out' + str(self.iter) + '.jpg') | |
| def hightlightPatchCandidates(self, chosenPatchId, nonchosenPatchId): | |
| result = np.copy(self.exampleMap) | |
| #mod patch ID | |
| chosenPatchId = chosenPatchId[0] % self.total_patches_count | |
| if len(nonchosenPatchId)>0: | |
| nonchosenPatchId = nonchosenPatchId % self.total_patches_count | |
| #exlcude chosen from nonchosen | |
| nonchosenPatchId = np.delete(nonchosenPatchId, np.where(nonchosenPatchId == chosenPatchId)) | |
| #highlight non chosen candidates | |
| c = [0.25, 0.9 ,0.45] | |
| self.highlightPatches(result, nonchosenPatchId, color=c, highlight_width = 4, alpha = 0.5) | |
| #hightlight chosen | |
| c = [1.0, 0.25, 0.15] | |
| self.highlightPatches(result, [chosenPatchId], color=c, highlight_width = 4, alpha = 1) | |
| return result | |
| def highlightPatches(self, writeResult, patchesIDs, color, highlight_width = 2, solid = False, alpha = 0.1): | |
| searchWindow = self.patchSize + 2*self.overlapSize | |
| #number of possible steps | |
| row_steps = floor((np.shape(writeResult)[0] - searchWindow) / self.windowStep) + 1 | |
| col_steps = floor((np.shape(writeResult)[1] - searchWindow) / self.windowStep) + 1 | |
| for i in range(len(patchesIDs)): | |
| chosenPatchId = patchesIDs[i] | |
| #patch Id to step | |
| patch_row = floor(chosenPatchId / col_steps) | |
| patch_col = chosenPatchId - patch_row * col_steps | |
| #highlight chosen patch (below are boundaries of the example patch) | |
| row_start = self.windowStep* patch_row | |
| row_end = self.windowStep * patch_row + searchWindow | |
| col_start = self.windowStep * patch_col | |
| col_end = self.windowStep * patch_col + searchWindow | |
| if not solid: | |
| w = highlight_width | |
| overlap = np.copy(writeResult[row_start:row_start+w, col_start:col_end]) | |
| writeResult[row_start:row_start+w, col_start:col_end] = overlap * (1-alpha) + (np.zeros(np.shape(overlap))+color) * alpha #top | |
| overlap = np.copy(writeResult[row_end-w:row_end, col_start:col_end]) | |
| writeResult[row_end-w:row_end, col_start:col_end] = overlap * (1-alpha) + (np.zeros(np.shape(overlap))+color) * alpha #bot | |
| overlap = np.copy( writeResult[row_start:row_end, col_start:col_start+w]) | |
| writeResult[row_start:row_end, col_start:col_start+w] = overlap * (1-alpha) + (np.zeros(np.shape(overlap))+color) * alpha #left | |
| overlap = np.copy(writeResult[row_start:row_end, col_end-w:col_end]) | |
| writeResult[row_start:row_end, col_end-w:col_end] = overlap * (1-alpha) + (np.zeros(np.shape(overlap))+color) * alpha #end | |
| else: | |
| a = alpha | |
| writeResult[row_start:row_end, col_start:col_end] = writeResult[row_start:row_end, col_start:col_end] * (1-a) + (np.zeros(np.shape(writeResult[row_start:row_end, col_start:col_end]))+color) * a | |
| def resize(self, imgArray, targetSize): | |
| img = Image.fromarray(np.uint8(imgArray*255)) | |
| img = img.resize((targetSize[0], targetSize[1]), resample=0, box=None) | |
| return np.array(img)/255 | |
| def findMostSimilarPatches(self, overlapArea_Top, overlapArea_Left, coord, in_k=5): | |
| #check which KD tree we need to use | |
| if (overlapArea_Top is not None) and (overlapArea_Left is not None): | |
| combined = self.getCombinedOverlap(overlapArea_Top.reshape(-1), overlapArea_Left.reshape(-1)) | |
| dist, ind = self.kdtree_combined.query([combined], k=in_k) | |
| elif overlapArea_Top is not None: | |
| dist, ind = self.kdtree_topOverlap.query([overlapArea_Top.reshape(-1)], k=in_k) | |
| elif overlapArea_Left is not None: | |
| dist, ind = self.kdtree_leftOverlap.query([overlapArea_Left.reshape(-1)], k=in_k) | |
| else: | |
| raise Exception("ERROR: no valid overlap area is passed to -findMostSimilarPatch-") | |
| dist = dist[0] | |
| ind = ind[0] | |
| return dist, ind | |
| #disallow visually similar blocks to be placed next to each other | |
| def checkForMirrors(self, dist, ind, coord, thres = 3): | |
| remove_i = [] | |
| #do I have a top or left neighbour | |
| if coord[0]-1>-1: | |
| top_neigh = int(self.idMap[coord[0]-1, coord[1]]) | |
| for i in range(len(ind)): | |
| if (abs(ind[i]%self.total_patches_count - top_neigh%self.total_patches_count) < thres): | |
| remove_i.append(i) | |
| if coord[1]-1>-1: | |
| left_neigh = int(self.idMap[coord[0], coord[1]-1]) | |
| for i in range(len(ind)): | |
| if (abs(ind[i]%self.total_patches_count - left_neigh%self.total_patches_count) < thres): | |
| remove_i.append(i) | |
| dist = np.delete(dist, remove_i) | |
| ind = np.delete(ind, remove_i) | |
| return dist, ind | |
| def distances2probability(self, distances, PARM_truncation, PARM_attenuation): | |
| probabilities = 1 - distances / np.max(distances) | |
| probabilities *= (probabilities > PARM_truncation) | |
| probabilities = pow(probabilities, PARM_attenuation) #attenuate the values | |
| #check if we didn't truncate everything! | |
| if np.sum(probabilities) == 0: | |
| #then just revert it | |
| probabilities = 1 - distances / np.max(distances) | |
| probabilities *= (probabilities > PARM_truncation*np.max(probabilities)) # truncate the values (we want top truncate%) | |
| probabilities = pow(probabilities, PARM_attenuation) | |
| probabilities /= np.sum(probabilities) #normalize so they add up to one | |
| return probabilities | |
| def getOverlapAreaTop(self, coord): | |
| #do I have a top neighbour | |
| if coord[0]-1>-1: | |
| canvasPatch = self.patchCoord2canvasPatch(coord) | |
| return canvasPatch[0:self.overlapSize, :, :] | |
| else: | |
| return None | |
| def getOverlapAreaLeft(self, coord): | |
| #do I have a left neighbour | |
| if coord[1]-1>-1: | |
| canvasPatch = self.patchCoord2canvasPatch(coord) | |
| return canvasPatch[:, 0:self.overlapSize, :] | |
| else: | |
| return None | |
| def initKDtrees(self): | |
| #prepate overlap patches | |
| topOverlap = self.examplePatches[:, 0:self.overlapSize, :, :] | |
| leftOverlap = self.examplePatches[:, :, 0:self.overlapSize, :] | |
| shape_top = np.shape(topOverlap) | |
| shape_left = np.shape(leftOverlap) | |
| flatten_top = topOverlap.reshape(shape_top[0], -1) | |
| flatten_left = leftOverlap.reshape(shape_left[0], -1) | |
| flatten_combined = self.getCombinedOverlap(flatten_top, flatten_left) | |
| tree_top = KDTree(flatten_top) | |
| tree_left = KDTree(flatten_left) | |
| tree_combined = KDTree(flatten_combined) | |
| return tree_top, tree_left, tree_combined | |
| #the corner of 2 overlaps is counted double | |
| def getCombinedOverlap(self, top, left): | |
| shape = np.shape(top) | |
| if len(shape) > 1: | |
| combined = np.zeros((shape[0], shape[1]*2)) | |
| combined[0:shape[0], 0:shape[1]] = top | |
| combined[0:shape[0], shape[1]:shape[1]*2] = left | |
| else: | |
| combined = np.zeros((shape[0]*2)) | |
| combined[0:shape[0]] = top | |
| combined[shape[0]:shape[0]*2] = left | |
| return combined | |
| def initFirstPatch(self): | |
| #grab a random block | |
| patchId = randint(0, np.shape(self.examplePatches)[0]) | |
| #mark out fill map | |
| self.filledMap[0, 0] = 1 | |
| self.idMap[0, 0] = patchId % self.total_patches_count | |
| #update canvas | |
| self.updateCanvas(patchId, 0, 0, False, False) | |
| #visualize | |
| # self.visualize([0,0], [patchId], []) | |
| def prepareExamplePatches(self): | |
| searchKernelSize = self.patchSize + 2 * self.overlapSize | |
| result = view_as_windows(self.exampleMap, [searchKernelSize, searchKernelSize, 3] , self.windowStep) | |
| shape = np.shape(result) | |
| result = result.reshape(shape[0]*shape[1], searchKernelSize, searchKernelSize, 3) | |
| self.total_patches_count = shape[0]*shape[1] | |
| if self.mirror_hor: | |
| #flip along horizonal axis | |
| hor_result = np.zeros(np.shape(result)) | |
| for i in range(self.total_patches_count): | |
| hor_result[i] = result[i][::-1, :, :] | |
| result = np.concatenate((result, hor_result)) | |
| if self.mirror_vert: | |
| vert_result = np.zeros((shape[0]*shape[1], searchKernelSize, searchKernelSize, 3)) | |
| for i in range(self.total_patches_count): | |
| vert_result[i] = result[i][:, ::-1, :] | |
| result = np.concatenate((result, vert_result)) | |
| return result | |
| def initCanvas(self): | |
| #check whether the outputSize adheres to patch+overlap size | |
| num_patches_X = ceil((self.outputSize[0]-self.overlapSize)/(self.patchSize+self.overlapSize)) | |
| num_patches_Y = ceil((self.outputSize[1]-self.overlapSize)/(self.patchSize+self.overlapSize)) | |
| #calc needed output image size | |
| required_size_X = num_patches_X*self.patchSize + (num_patches_X+1)*self.overlapSize | |
| required_size_Y = num_patches_Y*self.patchSize + (num_patches_X+1)*self.overlapSize | |
| #create empty canvas | |
| canvas = np.zeros((required_size_X, required_size_Y, 3)) | |
| filledMap = np.zeros((num_patches_X, num_patches_Y)) #map showing which patches have been resolved | |
| idMap = np.zeros((num_patches_X, num_patches_Y)) - 1 #stores patches id | |
| print("modified output size: ", np.shape(canvas)) | |
| print("number of patches: ", np.shape(filledMap)[0]) | |
| return canvas, filledMap, idMap | |
| def idCoordTo2DCoord(self, idCoord, imgSize): | |
| row = int(floor(idCoord / imgSize[0])) | |
| col = int(idCoord - row * imgSize[1]) | |
| return [row, col] | |
| def updateCanvas(self, inputPatchId, coord_X, coord_Y, blendTop = False, blendLeft = False): | |
| #translate Patch coordinate into Canvas coordinate | |
| x_range = self.patchCoord2canvasCoord(coord_X) | |
| y_range = self.patchCoord2canvasCoord(coord_Y) | |
| examplePatch = self.examplePatches[inputPatchId] | |
| if blendLeft: | |
| canvasOverlap = self.canvas[x_range[0]:x_range[1], y_range[0]:y_range[0]+self.overlapSize] | |
| examplePatchOverlap = np.copy(examplePatch[0][:, 0:self.overlapSize]) | |
| examplePatch[0][:, 0:self.overlapSize] = self.linearBlendOverlaps(canvasOverlap, examplePatchOverlap, 'left') | |
| if blendTop: | |
| canvasOverlap = self.canvas[x_range[0]:x_range[0]+self.overlapSize, y_range[0]:y_range[1]] | |
| examplePatchOverlap = np.copy(examplePatch[0][0:self.overlapSize, :]) | |
| examplePatch[0][0:self.overlapSize, :] = self.linearBlendOverlaps(canvasOverlap, examplePatchOverlap, 'top') | |
| self.canvas[x_range[0]:x_range[1], y_range[0]:y_range[1]] = examplePatch | |
| def linearBlendOverlaps(self, canvasOverlap, examplePatchOverlap, mode): | |
| if mode == 'left': | |
| mask = np.repeat(np.arange(self.overlapSize)[np.newaxis, :], np.shape(canvasOverlap)[0], axis=0) / self.overlapSize | |
| elif mode == 'top': | |
| mask = np.repeat(np.arange(self.overlapSize)[:, np.newaxis], np.shape(canvasOverlap)[1], axis=1) / self.overlapSize | |
| mask = np.repeat(mask[:, :, np.newaxis], 3, axis=2) #cast to 3d array | |
| return canvasOverlap * (1 - mask) + examplePatchOverlap * mask | |
| #def minimumBoundaryError(self, canvasOverlap, examplePatchOverlap, mode) | |
| def patchCoord2canvasCoord(self, coord): | |
| return [(self.patchSize+self.overlapSize)*coord, (self.patchSize+self.overlapSize)*(coord+1) + self.overlapSize] | |
| def patchCoord2canvasPatch(self, coord): | |
| x_range = self.patchCoord2canvasCoord(coord[0]) | |
| y_range = self.patchCoord2canvasCoord(coord[1]) | |
| return np.copy(self.canvas[x_range[0]:x_range[1], y_range[0]:y_range[1]]) | |
| def loadExampleMap(self, exampleMapPath): | |
| exampleMap = io.imread(exampleMapPath) #returns an MxNx3 array | |
| exampleMap = exampleMap / 255.0 #normalize | |
| #make sure it is 3channel RGB | |
| if (np.shape(exampleMap)[-1] > 3): | |
| exampleMap = exampleMap[:,:,:3] #remove Alpha Channel | |
| elif (len(np.shape(exampleMap)) == 2): | |
| exampleMap = np.repeat(exampleMap[np.newaxis, :, :], 3, axis=0) #convert from Grayscale to RGB | |
| return exampleMap | |