How to generate blackwhite mask image?
derrick-xwp opened this issue · comments
Use after effects, or if you want to use python try https://github.com/A1029384756/Automatic-Rotoscoper
I used rembg. It's basically just a nicely-polished u2net for segmasking: https://github.com/danielgatis/rembg
The below script will take square 1024x1024 single-person images as input from the directory at "in_path", and result in both images you need (the person standing in front of a black background, as well as the white segmask image that corresponds to it).
import os
from rembg.bg import remove
import numpy as np
import io
import cv2
from PIL import Image
import sys
def main(in_path, out_path):
count = 0
fails = []
for f in os.listdir(in_path):
try:
count+=1
if not os.path.exists(f'{out_path}/{f}'):
print(f'{in_path}/{f}')
img1 = cv2.imread(f'{in_path}/{f}')
#replace 1024 on the next 2 lines with the resolution of the desired output
f1 = 1024 / img1.shape[1]
f2 = 1024 / img1.shape[0]
fsz = min(f1, f2) # resizing factor
dim = (int(img1.shape[1] * fsz), int(img1.shape[0] * fsz))
fl = cv2.resize(img1, dim)
is_success, fl4 = cv2.imencode(".jpg", fl)
fl3 = fl4.tobytes()
result = remove(fl3)
im5 = Image.open(io.BytesIO(result)).convert("RGBA")
background = Image.new('RGBA', im5.size, (0,0,0))
im52 = Image.alpha_composite(background, im5)
img52 = im52.convert("RGB")
im52 = im52.save(f'{out_path}/{f}')
im = cv2.imread(f'{out_path}/{f}', cv2.IMREAD_UNCHANGED)
mask = cv2.threshold(im, 0, 255, cv2.THRESH_BINARY)[1]
cv2.imwrite(f'{out_path}/mask_{f}', mask)
print(f'count:{count} file: {f}')
else:
print(f'{f} exists, count: {count}')
except Exception as e:
print(sys.exc_info()[2])
print(e)
fails.append(f)
failnames = np.array([fails])
np.savetxt('failed_images.txt', failnames, delimiter='\n', fmt="%s")
if __name__ == '__main__':
in_path = '/2TB-Datasets/human_generator/tomask_tool/tomask'
out_path = '/2TB-Datasets/human_generator/tomask_tool/masks'
main(in_path, out_path)
example input image that you would put in 'tomask' (your input image doesn't have to have a white backgound, rembg should remove any background, just make sure there is only one person in the image):
the example results that land in 'masks' per the above script:
(ignore the poor quality person image, it's an output from an unfinished StyleGAN3 training)
Try VanceAI. It’s an online photo enhancer that has 20 features, including an AI background Remover. With this free web tool, you can remove the background from an image without any hassle. Whether it be an object or a person, VanceAI will never let you down.