Body Segmentation

# Generated silhouette pairs replace the excluded third-party archive.
from pathlib import Path
from PIL import Image
import numpy as np
import os
import matplotlib.pyplot as plt

dataset_root = Path("course_body_segmentation")
image_root = dataset_root / "images"
mask_root = dataset_root / "masks"
image_root.mkdir(parents=True, exist_ok=True)
mask_root.mkdir(parents=True, exist_ok=True)
yy, xx = np.mgrid[:128, :96]
for index in range(8):
    center_x = 48 + (index % 3 - 1) * 5
    head = (xx - center_x) ** 2 + (yy - 24) ** 2 < 10 ** 2
    torso = (np.abs(xx - center_x) < 13) & (yy >= 34) & (yy < 88)
    legs = ((np.abs(xx - (center_x - 7)) < 6) | (np.abs(xx - (center_x + 7)) < 6)) & (yy >= 82) & (yy < 122)
    mask = (head | torso | legs).astype(np.uint8)
    image = np.stack((
        35 + 150 * mask,
        55 + ((xx * 2 + index * 13) % 120),
        75 + ((yy + index * 17) % 100),
    ), axis=-1).astype(np.uint8)
    filename = f"person_{index:02d}.png"
    Image.fromarray(image).save(image_root / filename)
    Image.fromarray(mask * 255).save(mask_root / filename)

dataset_path = str(dataset_root)
print(f"Prepared {len(list(image_root.glob('*.png')))} generated image/mask pairs")

img_path = os.path.join(dataset_path,'images')
mask_path = os.path.join(dataset_path,'masks')

fnames = os.listdir(img_path)

def load_image(img_name):
    img = plt.imread(os.path.join(img_path,img_name))
    mask = plt.imread(os.path.join(mask_path,img_name))
    return img,mask

img, mask = load_image(fnames[5])

fig,ax = plt.subplots(1,2,figsize=(10,5))
ax[0].imshow(img)
ax[1].imshow(mask)
ax[0].axis('off')
ax[1].axis('off')