''' InpaintingMethod3B.py


This is method 3 B.
It differs from 3 (A) in that it averages only "good" pixels
during the relaxation.  Basically it propagates out from the edges
and does not average in values from any invalid pixels.
An invalid pixel is either (a) a barrier pixel or (b) an as
yet unpainted pixel.

During each iteration, it updates a per-pixel variable (stored in
the Green component of the combined image) that is the number of
VALID neighbors.  This equals 4 - #invalidNeighbors.


--------




At each pixel within the inpainting region, compute the an average 
of the pixel values of its "valid" neighbors.  Here we use the weights of
0.25 at each of N,E,W,S.

Method 3 adds "barriers" within the inpainting region.

A barrier pixel is indicated in the input image by the presence of
a pure red pixel.   To separate these, use the Mask together with
the "damaged with barriers" image.  Red pixels inside the inpainting
regions are barrier pixels.

The method allows barrier pixels to take values like other pixels
do within the inpainting region.  However, they are not allowed
to contribute to the values of other pixels.

Pixel may have any number between 0 and 3 of its neighbors as barrier
pixels, because if it had 4 then it could not be inpainted at all.
We need four separate rules for updating pixels, depending on their
numbers of barrier pixels.  One rule divides the sum of values by 4,
another by 3, another by 2 and another by 1.  Actually the same rule
will work, provided we have a separate image containing the number of
barrier pixels in the neighborhood.
 
In order to use PixelMath with effectively 3 source images rather than 2,
we can encode the Mask and the number of barrier pixels into the
red and green components of one image.

So the algorithm needs to compute the number of barrier pixels at
each location as a first step.


Open the following windows.

damagedWithBarriers: Image to be fixed.

mask: Mask image identifying the pixels to be inpainted.


Barriers(x,y)=if DamagedWithBarriers(x,y)=RGB(255,0,0) then 1 else 0.

counts(x,y) = S1(x-1,y)+S1(x,y-1)+S(x+1,y)+S1(x,y+1)

CombinedMaskAndBarriers(x,y)=RGB(if S1(x,y)=255 then 1 else 0,  NumBarriers(x,y), 0)

NonBarrierPixels  (current image, except 0s where barrier pixels are)

current: Source for update iteration

new: Destination for iterations.

diff: difference image (absolute values)

accum: image for accumulating sums of differences.



'''

damagedWB = pmOpenImage(0, '../Inpainting/damagedWithBarriers3.png')

mask = pmOpenImage(0, '../Inpainting/Mask2.png')
w = pmGetImageWidth(mask)
h = pmGetImageHeight(mask)

''' Clean up the input image within the inpainting areas.  It is supposed to be pure black, but
maybe during Writing by the GIMP, some pixels got set a little greater.  They need to be exactly 0
because that's the cue used to determine whether they are still unpained.'''
pmSetSource1(damagedWB)
pmSetSource2(mask)
pmSetDestination(damagedWB)
 # If within inpaint area and not barrier pixel then set to 0
pmSetFormula("If Red2(x,y)>10 then if Red1(x,y) < 10 then 0 else 255 else s1(x,y)")
pmCompute()


barriers= pmNewImage(0, "Barriers", w, h, 0, 0, 0)
barcnt  = pmNewImage(0, "Count of Barrier pixel neighbors (static)", w, h, 0, 0, 0)
counts  = pmNewImage(0, "Counts of valid pixels in neighborhood (dynamic)", w, h, 0, 0, 0)
valids  = pmNewImage(0, "1 if pixel can be used, 0 otherwise (dynamic)", w, h, 0, 0, 0)

combined= pmNewImage(0, "Combined: mask, counts, barriers", w, h, 0, 0, 0)

current = pmNewImage(0, "Current version", w, h, 0, 0, 0)
updated = pmNewImage(0, "updated version", w, h, 0, 0, 0)
diff    = pmNewImage(0, "Difference of last two", w, h, 0, 0, 0)
temp    = pmNewImage(0, "Half-combined (temporary)", w, h, 0, 0, 0)
accum   = pmNewImage(0, "Accumulator (temporary)", w, h, 0, 0, 0)
noBarr  = pmNewImage(0, "Current but 0 at barriers", w, h, 0, 0, 0)

pmPositionWindow(updated, 700, 300, 2*w, 2*h)
pmZoom(updated, w/2, h/2)

pmHideProgress()

# Isolate the barrier pixels
pmSetSource1(damagedWB)
pmSetSource2(mask)
pmSetDestination(barriers)
pmSetFormula("If Red1(x,y)>250 and Red2(x,y)>1 then 1 else 0")
pmCompute()

# Find the counts of barrier in each neighborhood.
pmSetSource1(barriers)
pmSetSource2(current)
pmSetDestination(barcnt)
pmSetFormula("if x>0 then S1(x-1,y) else 0)+(if y>0 then S1(x,y-1) else 0)+(if x<xmax then S1(x+1,y) else 0)+(if y<ymax then S1(x,y+1) else 0")
pmCompute()


# Now get ready for the iteration.
# Put the damaged image (without barrier values) into current version of image.

pmSetSource1(damagedWB)
pmSetSource2(combined)
pmSetDestination(current)
  # the max term in the following formula accelerates the inpainting by filling
  # from left to right and bottom to top with colors from the borders.
  # This can reduce the number of diffusion iterations.
pmSetFormula("if red2(x,y)=0 then S1(x,y) else max(if blue2(x-1,y)=0 then s1(x-y,y) else 128, if blue2(x,y-1)=0 then s1(x,y-1) else 128)")
pmCompute()

def clearAccums():
  pmSetRedAccumulator(0); pmSetGreenAccumulator(0); pmSetBlueAccumulator(0)

n = 150
for i in range(n):

  # Compute noBarr (i.e., current image with 0s at barrier positions)
  pmSetSource1(current)
  pmSetSource2(barriers)
  pmSetDestination(noBarr)
  pmSetFormula("if S2(x,y)=0 then S1(x,y) else 0")
  pmCompute()

  # Find the "valid" pixels that can be used to propagate values.
  # To be valid, it must have a non-black value and not be a barrier
  # First, simply compute a characteristic function for valid pixels generally.
  #pmSetSource2(barriers)
  pmSetDestination(valids)
  pmSetFormula("if (Red1(x,y)+Green1(x,y)+Blue1(x,y))>0 and S2(x,y)=0 then 1 else 0")
  pmCompute()

  # Count the number of valid neighbors
  pmSetSource2(valids)
  pmSetDestination(counts)
  pmSetFormula("(if x>0 then S2(x-1,y) else 0)+(if y>0 then S2(x,y-1) else 0)+(if x<xmax then S2(x+1,y) else 0)+(if y<ymax then S2(x,y+1) else 0)")
  pmCompute()

  # Now combine the mask and valid counts into one image.
  pmSetSource1(mask)
  pmSetSource2(counts)
  pmSetDestination(temp)
  # Red=1 means in inpainting region. Green gives count of barrier neighbors.
  pmSetFormula("RGB(if Red1(x,y)>1 then 1 else 0, S2(x,y), 0)")
  pmCompute()
  # Make the blue component hold the barrier characteristic fn.
  pmSetSource1(temp)
  pmSetSource2(barriers)
  pmSetDestination(combined)
  pmSetFormula("(S1(x,y)+RGB(0,0,S2(x,y)))")
  pmCompute()


  # perform relaxation:
  pmSetSource1(noBarr)
  pmSetSource2(combined)
  pmSetDestination(updated)
  # If pixel is in inpainting region, then average its valid neighbors. Else take orig. value.
  pmSetFormula("if Red2(x,y)>0 then if Green2(x,y)>0 then (1/(Green2(x,y)))*((if x>0 then S1(x-1,y) else 0)+(if y>0 then S1(x,y-1) else 0)+(if x<xmax then S1(x+1,y) else 0)+(if y<ymax then S1(x,y+1) else 0)) else 0 else S1(x,y)")
  pmCompute()

  # compute difference:
  pmSetSource1(current)
  pmSetSource2(updated)
  pmSetDestination(diff)
  pmSetFormula("abs(s1(x,y)-s2(x,y))")
  pmCompute()

  # find the total of the differences
  clearAccums()
  pmSetSource1(diff)
  pmSetDestination(accum)
  pmSetFormula('putinac(getac+S1(x,y))')
  pmCompute()
  # Here are the totals, possible quite large:

  red = pmReadRedAccumulator()
  green = pmReadGreenAccumulator()
  blue = pmReadBlueAccumulator()
  totalDifference = red+green+blue
  print "Iteration "+str(i)+", Total of differences: "+str(totalDifference)

  if totalDifference < 3: break

  # prepare for next iteration:
  pmSetSource1(updated)
  #pmSetSource2(mask)
  pmSetDestination(current)
  pmSetFormula('S1(x,y)')
  pmCompute()