File size: 1,762 Bytes
8da8f47 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 |
import os
from PIL import Image
import numpy as np
import argparse
if __name__ == "__main__":
parser = argparse.ArgumentParser()
parser.add_argument('--threshold', default=0.2, type=float, help='Path to option YMAL file.')
args = parser.parse_args()
input_folder = 'results/test_age-set'
output_folder = 'results/mask'
for filename in os.listdir(input_folder):
if filename.endswith('_0_0_LRGT.png'):
digits = filename.split('_')[0]
if digits.isdigit():
digits = int(digits)
if digits >= 0 and digits <= 1000:
input_path_LRGT = os.path.join(input_folder, filename)
input_path_SR_h = os.path.join(input_folder, filename).replace('LRGT', 'SR_h')
image_LRGT = Image.open(input_path_LRGT).convert("RGB")
image_SR_h = Image.open(input_path_SR_h).convert("RGB")
w, h = image_SR_h.size
image_LRGT = image_LRGT.resize((w, h))
array_LRGT = np.array(image_LRGT) / 255.
array_SR_h = np.array(image_SR_h) / 255.
residual = np.abs(array_LRGT - array_SR_h)
threshold = args.threshold
mask = np.where(residual > threshold, 1, 0)
os.makedirs(output_folder, exist_ok=True)
output_path = os.path.join(output_folder, str(digits+1).zfill(4)+'.png')
mask = np.sum(mask, axis=2)
mask_image = Image.fromarray((mask * 255).astype(np.uint8))
mask_image.save(output_path) |