Spaces:
Runtime error
Runtime error
File size: 1,784 Bytes
a0bcaae bb0f5a9 a0bcaae bb0f5a9 a0bcaae bb0f5a9 a0bcaae bb0f5a9 a0bcaae bb0f5a9 a0bcaae bb0f5a9 a0bcaae bb0f5a9 a0bcaae bb0f5a9 a0bcaae bb0f5a9 |
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 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 |
# Copyright (c) SenseTime Research. All rights reserved.
import os
import click
import cv2
import numpy as np
def bg_white(seg, raw, blur_level=3, gaussian=81):
seg = cv2.blur(seg, (blur_level, blur_level))
empty = np.ones_like(seg)
seg_bg = (empty - seg) * 255
seg_bg = cv2.GaussianBlur(seg_bg, (gaussian, gaussian), 0)
background_mask = cv2.cvtColor(
255 - cv2.cvtColor(seg, cv2.COLOR_BGR2GRAY), cv2.COLOR_GRAY2BGR)
masked_fg = (raw * (1 / 255)) * (seg * (1 / 255))
masked_bg = (seg_bg * (1 / 255)) * (background_mask * (1 / 255))
frame = np.uint8(cv2.add(masked_bg, masked_fg)*255)
return frame
"""
To turn background into white.
Examples:
\b
python bg_white.py --raw_img_dir=./SHHQ-1.0/no_segment/ --raw_seg_dir=./SHHQ-1.0/segments/ \\
--outdir=./SHHQ-1.0/bg_white/
"""
@click.command()
@click.pass_context
@click.option('--raw_img_dir', default="./SHHQ-1.0/no_segment/", help='folder of raw image', required=True)
@click.option('--raw_seg_dir', default='./SHHQ-1.0/segments/', help='folder of segmentation masks', required=True)
@click.option('--outdir', help='Where to save the output images', default="./SHHQ-1.0/bg_white/", type=str, required=True, metavar='DIR')
def main(
ctx: click.Context,
raw_img_dir: str,
raw_seg_dir: str,
outdir: str):
os.makedirs(outdir, exist_ok=True)
files = os.listdir(raw_img_dir)
for file in files:
print(file)
raw = cv2.imread(os.path.join(raw_img_dir, file))
seg = cv2.imread(os.path.join(raw_seg_dir, file))
assert raw is not None
assert seg is not None
white_frame = bg_white(seg, raw)
cv2.imwrite(os.path.join(outdir, file), white_frame)
if __name__ == "__main__":
main()
|