You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

143 lines
5.9 KiB

import numpy as np
from maskedImage import MaskedImage
from nnf import Nnf
class Inpaint:
def __init__(self,input,mask,radius):
self.initial = MaskedImage(input,mask)
self.radius = radius
self.pyramid = [self.initial]
source = self.initial
while source.width>radius and source.height>radius:
source = source.downsample()
self.pyramid.append(source)
maxLevel = len(self.pyramid)
for level in range(maxLevel-1,0,-1):
print(level)
source = self.pyramid[level]
if (level == maxLevel-1):
target = source.copy()
for y in range(target.height):
for x in range(target.width):
target.mask[y,x] = False
self.nnfSourceToTarget = Nnf(source,target,radius)
self.nnfSourceToTarget.randomize()
self.nnfTargetToSource = Nnf(target,source,radius)
self.nnfTargetToSource.randomize()
else:
newNnf = Nnf(source,target,radius)
newNnf.initializeFromNnf(self.nnfSourceToTarget)
self.nnfSourceToTarget = newNnf
newNnfRev = Nnf(target,source,radius)
newNnfRev.initializeFromNnf(self.nnfTargetToSource)
self.nnfTargetToSource = newNnfRev
target = self.ExpectationMaximization(level)
return target.image
def ExpectationMaximization(self,level):
iterEM = 1+2*level
iterNnf = min(7,1+level)
source = self.nnfSourceToTarget.input
target = self.nnfTargetToSource.output
newTarget = None
for emloop in range(1,iterEM+1):
if (newTarget != None):
self.nnfSourceToTarget.output = newTarget
self.nnfTargetToSource.input = newTarget
target = newTarget
newTarget = None
for y in range(source.height):
for x in range(source.width):
if not source.containsMask(x,y,self.radius):
self.nnfSourceToTarget.field[y,x,0] = x
self.nnfSourceToTarget.field[y,x,1] = y
self.nnfSourceToTarget.field[y,x,2] = 0
for y in range(target.height):
for x in range(target.width):
if not source.containsMask(x,y,self.radius):
self.nnfTargetToSource.field[y,x,0] = x
self.nnfTargetToSource.field[y,x,1] = y
self.nnfTargetToSource.field[y,x,2] = 0
self.nnfSourceToTarget.minimize(iterNnf)
self.nnfTargetToSource.minimize(iterNnf)
upscaled = False
if level>=1 and emloop == iterEM:
newSource = self.pyramid[level-1]
newTarget = target.upscale(newSource.height,newSource.width)
upscaled = True
else:
newSource = self.pyramid[level]
newTarget = target.copy()
upscaled = False
vote = np.zeros((newTarget.width, newTarget.height, 4))
self.ExpectationStep(self.nnfSourceToTarget,True,vote,newSource,upscaled)
self.ExpectationStep(self.nnfTargetToSource,False,vote,newSource,upscaled)
self.MaximizationStep(newTarget, vote)
result = MaskedImage.resize(newTarget.image,self.initial.width,self.initial.height)
from demo import Demo
Demo.display(result)
return newTarget
def ExpectationStep(self,nnf,sourceToTarget, vote, source, upscale):
for y in range(nnf.input.height):
for x in range(nnf.input.width):
xp = nnf.field[y,x,0]
yp = nnf.field[y,x,1]
dp = nnf.field[y,x,2]
w = MaskedImage.similarity[dp]
for dy in range(-nnf.patchSize,nnf.patchSize):
for dx in range(-nnf.patchSize,nnf.patchSize):
if sourceToTarget:
xs = x+dx
ys = y+dy
xt = xp+dx
yt = yp+dy
else:
xs = xp+dx
ys = yp+dy
xt = x+dx
yt = y+dy
if not 0<=xs<nnf.input.width:
continue
if not 0<=ys<nnf.input.height:
continue
if not 0<=xt<nnf.input.width:
continue
if not 0<=yt<nnf.input.height:
continue
if upscale:
self.weightedCopy(source,2*xs,2*ys,vote,2*xt,2*yt,w)
self.weightedCopy(source,2*xs+1,2*ys,vote,2*xt+1,2*yt,w)
self.weightedCopy(source,2*xs,2*ys+1,vote,2*xt,2*yt+1,w)
self.weightedCopy(source,2*xs+1,2*ys+1,vote,2*xt+1,2*yt+1,w)
else:
self.weightedCopy(source,xs,ys,vote,xt,yt,w)
def weightedCopy(self,src,xs,ys,vote,xd,yd,w):
if src.mask[ys,xs]:
return
vote[xd,yd,0] += w*src.image[ys,xs,0]
vote[xd,yd,1] += w*src.image[ys,xs,1]
vote[xd,yd,2] += w*src.image[ys,xs,2]
vote[xd,yd,3] += w
def MaximizationStep(self,target,vote):
for y in range(target.height):
for x in range(target.width):
if vote[x,y,3]>0:
r = int(vote[x,y,0]/vote[x,y,3])
g = int(vote[x,y,1]/vote[x,y,3])
b = int(vote[x,y,2]/vote[x,y,3])
target.image[y,x,0] = r
target.image[y,x,1] = g
target.image[y,x,2] = b
target.mask[y,x] = False