Inpaint an Image

inpaint.py

import cv2
import numpy as np

def inpaint(input_file_path, output_path):
  !rm -r /content/input
  !mkdir /content/input
  !cp $input_file_path /content/input
  parts = input_file_path.split("/")
  mask_name = "/content/input/" + parts[-1]
  mask_name = mask_name[:-4]+ "_mask.png"
  image = cv2.imread(input_file_path)

  gray1 = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
  cv2.imwrite("/content/gray1.png", gray1)

  gray2 = gray1.copy()

  h, w = gray2.shape[:2]
  cv2.floodFill(gray2, None, seedPoint=(0, 0), newVal=128)
  cv2.floodFill(gray2, None, seedPoint=(w-1, 0), newVal=128)
  cv2.floodFill(gray2, None, seedPoint=(0, h-1), newVal=128)
  cv2.floodFill(gray2, None, seedPoint=(w-1, h-1), newVal=128)
  cv2.imwrite("/content/gray2.png", gray2)

  diff1 = cv2.subtract(gray1, gray2)
  diff2 = cv2.subtract(gray2, gray1)
  diff = cv2.add(diff1, diff2)
  cv2.imwrite("/content/diff.png", diff)

  t, thresh  = cv2.threshold(diff, 64, 255, cv2.THRESH_BINARY)
  cv2.imwrite("/content/thresh.png", thresh)

  kernel = np.ones((31,31),np.uint8)
  mask = cv2.dilate(thresh, kernel, iterations = 1)
  cv2.imwrite(mask_name, mask)

  average = np.mean(mask)

  if average > 1:
    !export TORCH_HOME=$(pwd) && export PYTHONPATH=.; python bin/predict.py model.path=/content/lama/big-lama indir=/content/input outdir=$output_path > /dev/null 2>&1
    print(input_file_path, round(average, 3), "filled")
  else:
    !cp $input_file_path $output_path
    print(input_file_path, round(average, 3), "copied")