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
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
|
|
|
|
|