| import os |
| from pathlib import Path |
|
|
| import cv2 |
| import torch |
| from model import BiSeNet |
| from PIL import Image |
| from torch.utils.data import Dataset |
| from torchvision import transforms |
| from tqdm import tqdm |
|
|
| |
|
|
|
|
| class MaskDataset(Dataset): |
| def __init__(self, img_root, mask_root): |
| img_dir = Path(img_root) |
| self.to_tensor_normalize = transforms.Compose( |
| [ |
| transforms.ToTensor(), |
| transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]), |
| ] |
| ) |
| self.img_files = list(img_dir.glob(f"**/*.jpg")) |
| self.img_files.sort() |
| self.mask_files = [os.path.join(mask_root, os.path.relpath(img_path, img_root)) for img_path in self.img_files] |
|
|
| def __len__(self): |
| return len(self.mask_files) |
|
|
| def __getitem__(self, index): |
| img = Image.open(self.img_files[index]).convert("RGB") |
| return {"img": self.to_tensor_normalize(img), "mask_path": self.mask_files[index]} |
|
|
|
|
| class MaskDataLoader: |
| def __init__(self): |
| """Initialize this class""" |
| self.dataset = MaskDataset(img_root="/data/dataset/face_1k/alignHQ", mask_root="/data/dataset/face_1k/mask") |
|
|
| self.dataloader = torch.utils.data.DataLoader( |
| self.dataset, batch_size=8, shuffle=True, num_workers=8, drop_last=False |
| ) |
|
|
| def __len__(self): |
| """Return the number of data in the dataset""" |
| return len(self.dataset) / 8 |
|
|
| def __iter__(self): |
| """Return a batch of data""" |
| for data in self.dataloader: |
| yield data |
|
|
|
|
| if __name__ == "__main__": |
| dataloader = MaskDataLoader() |
| bisenet_path = "/data/useful_ckpt/face_parsing/parsing_model_79999_iter.pth" |
| bisenet = BiSeNet(n_classes=19) |
| device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu") |
| bisenet.to(device) |
| state_dict = torch.load(bisenet_path, map_location=device) |
| bisenet.load_state_dict(state_dict) |
| bisenet.eval() |
|
|
| for data in tqdm(dataloader): |
| mask, ignore_ids = bisenet.get_mask(data["img"].to(device), 256) |
| mask = (mask * 255).to(torch.uint8).cpu().numpy().transpose(0, 2, 3, 1).repeat(3, 3) |
|
|
| for i in range(mask.shape[0]): |
| if ignore_ids[i]: |
| continue |
| path = data["mask_path"][i] |
| dirname = os.path.dirname(path) |
| os.makedirs(dirname, exist_ok=True) |
| cv2.imwrite(path, mask[i]) |
|
|
