Unet-OCR-2.0/postprocess.py

import numpy as np
from PIL import Image
from matplotlib import pyplot as plt
import matplotlib.image as mpimg
import cv2
import imageio
import os
from os import listdir
from os.path import isfile, join
import argparse
import shutil

args = argparse.ArgumentParser()
args.add_argument("-i", "--input", required=True, help= "absolute path to mask")
args.add_argument("-e", "--expected", required=True, help= "absolute path to image")
args.add_argument("-o", "--output", required=True, help="path to save image")
args = vars(args.parse_args())

masked_file=args["input"]
original_file=args["expected"]
outdir=args["output"]

img_pred = cv2.imread(masked_file)

row = img_pred.shape[0]
col = img_pred.shape[1]

img_exp = cv2.imread(original_file)
new_img = np.ones(img_exp.shape)*255
new_img=new_img.astype(np.uint8)

for r in range(row):
  for c in range(col):
    if img_pred[r,c][0] == 0 and img_pred[r,c][1] == 0 and img_pred[r,c][2] == 0:
      new_img[r,c][0] = img_exp[r,c][0]
      new_img[r,c][1] = img_exp[r,c][1]
      new_img[r,c][2] = img_exp[r,c][2]

newimg = masked_file+"_pos"
imageio.imwrite(newimg, new_img, format="png")

shutil.move(newimg,outdir)