repo_description: A simple interface for editing natural photos with generative neural networks.
repo_topics: computer-vision convolutional-neural-networks deep-learning gans interfaces machine-learning
code: ### Neural Photo Editor
  # A Brock, 2016
  
  ### Imports
  
  from Tkinter import * # Note that I dislike the * on the Tkinter import, but all the tutorials seem to do that so I stuck with it.
  from tkColorChooser import askcolor # This produces an OS-dependent color selector. I like the windows one best, and can't stand the linux one.
  from collections import OrderedDict
  from PIL import Image, ImageTk
  import numpy as np
  import scipy.misc
  
  
  from API import IAN
  ### Step 1: Create theano functions
  
  # Initialize model
  model = IAN(config_path = 'IAN_simple.py', dnn = True)
  
  ### Prepare GUI functions
  print('Compiling remaining functions')
  
  # Create master
  master = Tk()
  master.title( "Neural Photo Editor" )
  
  # RGB interpreter convenience function
  def rgb(r,g,b):
      return '#%02x%02x%02x' % (r,g,b)
      
  # Convert RGB to bi-directional RB scale.
  def rb(i):
      # return rgb(int(i*int(i>0)),0, -int(i*int(i<0)))
      return rgb(255+max(int(i*int(i<0)),-255),255-min(abs(int(i)),255), 255-min(int(i*int(i>0)),255))
  
  # Convenience functions to go from [0,255] to [-1,1] and [-1,1] to [0,255]    
  def to_tanh(input):
      return 2.0*(input/255.0)-1.0
   
  def from_tanh(input):
      return 255.0*(input+1)/2.0  
      
  # Ground truth image
  GIM=np.asarray(np.load('CelebAValid.npz')['arr_0'][420])
  
  # Image for modification
  IM = GIM
  
  # Reconstruction
  RECON = IM
  
  # Error between reconstruction and current image
  ERROR = np.zeros(np.shape(IM),dtype=np.float32)
  
  # Change between Recon and Current
  DELTA = np.zeros(np.shape(IM),dtype=np.float32)
  
  # User-Painted Mask, currently not implemented.
  USER_MASK=np.mean(DELTA,axis=0)
  
  # Are we operating on a photo or a sample?
  SAMPLE_FLAG=0
  
  
  ### Latent Canvas Variables
  # Latent Square dimensions
  dim = [10,10]
  
  # Squared Latent Array
  Z = np.zeros((dim[0],dim[1]),dtype=np.float32)
  
  # Pixel-wise resolution for latent canvas
  res = 16
  
  # Array that holds the actual latent canvas
  r = np.zeros((res*dim[0],res*dim[1]),dtype=np.float32)
  
  # Painted rectangles for free-form latent painting
  painted_rects = []
  
  # Actual latent rectangles
  rects = np.zeros((dim[0],dim[1]),dtype=int)
  
  ### Output Display Variables
  
  # RGB paintbrush array
  myRGB = np.zeros((1,3,64,64),dtype=np.float32);
  
  # Canvas width and height
  canvas_width = 400
  canvas_height = 400
  
  # border width
  bd =2 
  # Brush color
  color = IntVar() 
  color.set(0)
  
  # Brush size
  d = IntVar() 
  d.set(12)
  
  # Selected Color
  mycol = (0,0,0)
  
  # Function to update display
  def update_photo(data=None,widget=None):
      global Z
      if data is None: # By default, assume we're updating with the current value of Z
          data = np.repeat(np.repeat(np.uint8(from_tanh(model.sample_at(np.float32([Z.flatten()]))[0])),4,1),4,2)
      else:
          data = np.repeat(np.repeat(np.uint8(data),4,1),4,2)
      
      if widget is None:
          widget = output
      # Reshape image to canvas
      mshape = (4*64,4*64,1)
      im = Image.fromarray(np.concatenate([np.reshape(data[0],mshape),np.reshape(data[1],mshape),np.reshape(data[2],mshape)],axis=2),mode='RGB')
      
      # Make sure photo is an object of the current widget so the garbage collector doesn't wreck it
      widget.photo = ImageTk.PhotoImage(image=im)
      widget.create_image(0,0,image=widget.photo,anchor=NW)
      widget.tag_raise(pixel_rect)
  
  # Function to update the latent canvas.    
  def update_canvas(widget=None):
      global r, Z, res, rects, painted_rects
      if widget is None:
          widget = w
      # Update display values
      r = np.repeat(np.repeat(Z,r.shape[0]//Z.shape[0],0),r.shape[1]//Z.shape[1],1)
      
      # If we're letting freeform painting happen, delete the painted rectangles
      for p in painted_rects:
          w.delete(p)
      painted_rects = []
      
      for i in range(Z.shape[0]):
          for j in range(Z.shape[1]):
              w.itemconfig(int(rects[i,j]),fill = rb(255*Z[i,j]),outline = rb(255*Z[i,j]))
  
  # Function to move the paintbrush
  def move_mouse( event ):
      global output
      # using a rectangle width equivalent to d/4 (so 1-16)
      
      # First, get location and extent of local patch
      x,y = event.x//4,event.y//4
      brush_width = ((d.get()//4)+1)
      
      # if x is near the left corner, then the minimum x is dependent on how close it is to the left
      xmin = max(min(x-brush_width//2,64 - brush_width),0) # This 64 may need to change if the canvas size changes
      xmax = xmin+brush_width
      
      ymin = max(min(y-brush_width//2,64 - brush_width),0) # This 64 may need to change if the canvas size changes
      ymax = ymin+brush_width
      
      # update output canvas
      output.coords(pixel_rect,4*xmin,4*ymin,4*xmax,4*ymax)
      output.tag_raise(pixel_rect)
      output.itemconfig(pixel_rect,outline = rgb(mycol[0],mycol[1],mycol[2]))
  
  ### Optional functions for the Neural Painter
  
  # Localized Gaussian Smoothing Kernel
  # Use this if you want changes to MASK to be more localized to the brush location in soe sense
  def gk(c1,r1,c2,r2):
      # First, create X and Y arrays indicating distance to the boundaries of the paintbrush
      # In this current context, im is the ordinal number of pixels (64 typically)
      sigma = 0.3
      im = 64
      x = np.repeat([np.concatenate([np.mgrid[-c1:0],np.zeros(c2-c1),np.mgrid[1:1+im-c2]])],im,axis=0)
      y = np.repeat(np.vstack(np.concatenate([np.mgrid[-r1:0],np.zeros(r2-r1),np.mgrid[1:1+im-r2]])),im,axis=1)
      g = np.exp(-(x**2/float(im)+y**2/float(im))/(2*sigma**2))
      return np.repeat([g],3,axis=0) # remove the 3 if you want to apply this to mask rather than an RGB channel  
  # This function reduces the likelihood of a change based on how close each individual pixel is to a maximal value.
  # Consider conditioning this based on the gK value and the requested color. I.E. instead of just a flat distance from 128,
  # have it be a difference from the expected color at a given location. This could also be used to "weight" the image towards staying the same.
  def upperlim(image):
      h=1
      return (1.0/((1.0/h)*np.abs(image-128)+1))
      
  # Similar to upperlim, this function changes the value of the correction term if it's going to move pixels too close to a maximal value
  def dampen(input,correct):
      # The closer input+correct is to -1 or 1, the further it is from 0. 
      # We're okay with almost all values (i.e. between 0 and 0.8) but as we approach 1 we want to slow the change
      thresh = 0.75
      m = (input+correct)>thresh
      return -input*m+correct*(1-m)+thresh*m   
      
  ### Neural Painter Function     
  def paint( event ):
      global Z, output, myRGB, IM, ERROR, RECON, USER_MASK, SAMPLE_FLAG
      
      # Move the paintbrush
      move_mouse(event)
      
      # Define a gradient descent step-size
      weight = 0.05
      
      # Get paintbrush location
      [x1,y1,x2,y2] = [coordinate//4 for coordinate in output.coords(pixel_rect)]
      
      # Get dIM/dZ that minimizes the difference between IM and RGB in the domain of the paintbrush
      temp = np.asarray(model.imgradRGB(x1,y1,x2,y2,np.float32(to_tanh(myRGB)),np.float32([Z.flatten()]))[0])
      grad = temp.reshape((10,10))*(1+(x2-x1))
      
      # Update Z
      Z -=weight*grad
      
      # If operating on a sample, update sample
      if  SAMPLE_FLAG:
          update_canvas(w)
          update_photo(None,output)
      # Else, update photo
      else:
          # Difference between current image and reconstruction
          DELTA = model.sample_at(np.float32([Z.flatten()]))[0]-to_tanh(np.float32(RECON))
          
          # Not-Yet-Implemented User Mask feature
          # USER_MASK[y1:y2,x1:x2]+=0.05
          
          # Get MASK
          MASK=scipy.ndimage.filters.gaussian_filter(np.min([np.mean(np.abs(DELTA),axis=0),np.ones((64,64))],axis=0),0.7)
          
          # Optionally dampen D
          # D = dampen(to_tanh(np.float32(RECON)),MASK*DELTA+(1-MASK)*ERROR)
          
          # Update image
          D = MASK*DELTA+(1-MASK)*ERROR
          IM = np.uint8(from_tanh(to_tanh(RECON)+D))
          
          # Pass updates
          update_canvas(w)
          update_photo(IM,output)     
          
  # Load an image and infer/reconstruct from it. Update this with a function to load your own images if you want to edit
  # non-celebA photos.
  def infer():
      global Z,w,GIM,IM,ERROR,RECON,DELTA,USER_MASK,SAMPLE_FLAG
      val = myentry.get()
      try:
          val = int(val)
          GIM = np.asarray(np.load('CelebAValid.npz')['arr_0'][val])
          IM = GIM
      except ValueError:
          print "No input"
          val = 420
          GIM = np.asarray(np.load('CelebAValid.npz')['arr_0'][val])
          IM = GIM
      # myentry.delete(0, END) # Optionally, clear entry after typing it in   
      
      # Reset Delta
      DELTA = np.zeros(np.shape(IM),dtype=np.float32)
      
      # Infer and reshape latents. This can be done without an intermediate variable if desired
      s = model.encode_images(np.asarray([to_tanh(IM)],dtype=np.float32))
      Z = np.reshape(s[0],np.shape(Z))
      
      # Get reconstruction
      RECON = np.uint8(from_tanh(model.sample_at(np.float32([Z.flatten()]))[0]))
      
      # Get error
      ERROR = to_tanh(np.float32(IM)) - to_tanh(np.float32(RECON))
      
      # Reset user mask
      USER_MASK*=0
      
      # Clear the sample flag
      SAMPLE_FLAG=0
      
      # Update photo
      update_photo(IM,output)
      update_canvas(w) 
  
  # Paint directly into the latent space
  def paint_latents( event ):
     global r, Z, output,painted_rects,MASK,USER_MASK,RECON 
  
     # Get extent of latent paintbrush
     x1, y1 = ( event.x - d.get() ), ( event.y - d.get() )
     x2, y2 = ( event.x + d.get() ), ( event.y + d.get() )
     
     selected_widget = event.widget
     
     # Paint in latent space and update Z
     painted_rects.append(event.widget.create_rectangle( x1, y1, x2, y2, fill = rb(color.get()),outline = rb(color.get()) ))
     r[max((y1-bd),0):min((y2-bd),r.shape[0]),max((x1-bd),0):min((x2-bd),r.shape[1])] = color.get()/255.0;
     Z = np.asarray([np.mean(o) for v in [np.hsplit(h,Z.shape[0])\
                                                           for h in np.vsplit((r),Z.shape[1])]\
                                                           for o in v]).reshape(Z.shape[0],Z.shape[1])
     if SAMPLE_FLAG:
          update_photo(None,output)
          update_canvas(w) # Remove this if you wish to see a more free-form paintbrush
     else:     
          DELTA = model.sample_at(np.float32([Z.flatten()]))[0]-to_tanh(np.float32(RECON))
          MASK=scipy.ndimage.filters.gaussian_filter(np.min([np.mean(np.abs(DELTA),axis=0),np.ones((64,64))],axis=0),0.7)
          # D = dampen(to_tanh(np.float32(RECON)),MASK*DELTA+(1-MASK)*ERROR)
          D = MASK*DELTA+(1-MASK)*ERROR
          IM = np.uint8(from_tanh(to_tanh(RECON)+D))
          update_canvas(w) # Remove this if you wish to see a more free-form paintbrush
          update_photo(IM,output)  
  
  # Scroll to lighten or darken an image patch
  def scroll( event ):
      global r,Z,output
      # Optional alternate method to get a single X Y point
      # x,y = np.floor( ( event.x - (output.winfo_rootx() - master.winfo_rootx()) ) / 4), np.floor( ( event.y - (output.winfo_rooty() - master.winfo_rooty()) ) / 4)
      weight = 0.1
      [x1,y1,x2,y2] = [coordinate//4 for coordinate in output.coords(pixel_rect)]
      grad = np.reshape(model.imgrad(x1,y1,x2,y2,np.float32([Z.flatten()]))[0],Z.shape)*(1+(x2-x1))
      Z+=np.sign(event.delta)*weight*grad
      update_canvas(w)
      update_photo(None,output)
  
  # Samples in the latent space 
  def sample():
      global  Z, output,RECON,IM,ERROR,SAMPLE_FLAG
      Z = np.random.randn(Z.shape[0],Z.shape[1])
      # Z = np.random.uniform(low=-1.0,high=1.0,size=(Z.shape[0],Z.shape[1])) # Optionally get uniform sample
      
      # Update reconstruction and error
      RECON = np.uint8(from_tanh(model.sample_at(np.float32([Z.flatten()]))[0]))
      ERROR = to_tanh(np.float32(IM)) - to_tanh(np.float32(RECON))
      update_canvas(w)
      SAMPLE_FLAG=1
      update_photo(None,output)
  
  # Reset to ground-truth image
  def Reset():
      global GIM,IM,Z, DELTA,RECON,ERROR,USER_MASK,SAMPLE_FLAG
      IM  = GIM
      Z = np.reshape(model.encode_images(np.asarray([to_tanh(IM)],dtype=np.float32))[0],np.shape(Z))
      DELTA = np.zeros(np.shape(IM),dtype=np.float32)
      RECON = np.uint8(from_tanh(model.sample_at(np.float32([Z.flatten()]))[0]))
      ERROR = to_tanh(np.float32(IM)) - to_tanh(np.float32(RECON))
      USER_MASK*=0
      SAMPLE_FLAG=0
      update_canvas(w)
      update_photo(IM,output)
  
  def UpdateGIM():
      global GIM,IM
      GIM = IM
      Reset()# Recalc the latent space for the new ground-truth image.
      
  # Change brush size
  def update_brush(event):
      brush.create_rectangle(0,0,25,25,fill=rgb(255,255,255),outline=rgb(255,255,255))
      brush.create_rectangle( int(12.5-d.get()/4.0), int(12.5-d.get()/4.0), int(12.5+d.get()/4.0), int(12.5+d.get()/4.0), fill = rb(color.get()),outline = rb(color.get()) )
  
  # assign color picker values to myRGB
  def getColor():
      global myRGB, mycol
      col = askcolor(mycol)
      if col[0] is None:
          return # Dont change color if Cancel pressed.
      mycol = col[0]
      for i in xrange(3): myRGB[0,i,:,:] = mycol[i]; # assign
      
  # Optional function to "lock" latents so that gradients are always evaluated with respect to the locked Z
  # def lock():
      # global Z,locked, Zlock, lockbutton
      # lockbutton.config(relief='raised' if locked else 'sunken')
      # Zlock = Z if not locked else Zlock
      # locked = not locked
  # lockbutton = Button(f, text="Lock", command=lock,relief='raised')
  # lockbutton.pack(side=LEFT)
      
  ### Prepare GUI
  master.bind("<MouseWheel>",scroll)
  
  # Prepare drawing canvas
  f=Frame(master)
  f.pack(side=TOP)
  output = Canvas(f,name='output',width=64*4,height=64*4)
  output.bind('<Motion>',move_mouse)
  output.bind('<B1-Motion>', paint )
  pixel_rect = output.create_rectangle(0,0,4,4,outline = 'yellow')
  output.pack()   
  
  # Prepare latent canvas
  f = Frame(master,width=res*dim[0],height=dim[1]*10)
  f.pack(side=TOP)
  w = Canvas(f,name='canvas', width=res*dim[0],height=res*dim[1])
  w.bind( "<B1-Motion>", paint_latents )
  # Produce painted rectangles
  for i in range(Z.shape[0]):
      for j in range(Z.shape[1]):
          rects[i,j] = w.create_rectangle( j*res, i*res, (j+1)*res, (i+1)*res, fill = rb(255*Z[i,j]),outline = rb(255*Z[i,j]) )
  # w.create_rectangle( 0,0,res*dim[0],res*dim[1], fill = rgb(255,255,255),outline=rgb(255,255,255)) # Optionally Initialize canvas to white 
  w.pack()
  
  
  # Color gradient    
  gradient = Canvas(master, width=400, height=20)
  gradient.pack(side=TOP)
  # gradient.grid(row=i+1)
  for j in range(-200,200):
      gradient.create_rectangle(j*255/200+200,0,j*255/200+201,20,fill = rb(j*255/200),outline=rb(j*255/200))
  # Color scale slider
  f= Frame(master)
  Scale(master, from_=-255, to=255,length=canvas_width, variable = color,orient=HORIZONTAL,showvalue=0,command=update_brush).pack(side=TOP)
  
  # Buttons and brushes
  Button(f, text="Sample", command=sample).pack(side=LEFT)
  Button(f, text="Reset", command=Reset).pack(side=LEFT)
  Button(f, text="Update", command=UpdateGIM).pack(side=LEFT)
  brush = Canvas(f,width=25,height=25)
  Scale(f, from_=0, to=64,length=100,width=25, variable = d,orient=HORIZONTAL,showvalue=0,command=update_brush).pack(side=LEFT) # Brush diameter scale
  brush.pack(side=LEFT)
  inferbutton = Button(f, text="Infer", command=infer)
  inferbutton.pack(side=LEFT)
  colorbutton=Button(f,text='Col',command=getColor)
  colorbutton.pack(side=LEFT) 
  myentry = Entry()
  myentry.pack(side=LEFT)
  f.pack(side=TOP)
  
  
  print('Running')  
  # Reset and infer to kick it off
  Reset()
  infer() 
  mainloop()
  
  
  
  