| # import os | |
| # import cv2 | |
| # import numpy as np | |
| # import torch | |
| # import albumentations as albu | |
| # from iglovikov_helper_functions.utils.image_utils import load_rgb, pad, unpad | |
| # from iglovikov_helper_functions.dl.pytorch.utils import tensor_from_rgb_image | |
| # from cloths_segmentation.pre_trained_models import create_model | |
| # from tqdm import tqdm | |
| # model = create_model("Unet_2020-10-30") | |
| # model.to("cuda") | |
| # model.eval() | |
| # input_dir = "../../image" | |
| # output_dir = "../../seg_masks" | |
| # os.makedirs(output_dir, exist_ok=True) | |
| # for image_filename in tqdm(os.listdir(input_dir), colour="green"): | |
| # image_path = os.path.join(input_dir, image_filename) | |
| # image = cv2.imread(image_path) | |
| # image_2_extract = image | |
| # image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| # transform = albu.Compose([albu.Normalize(p=1)], p=1) | |
| # padded_image, pads = pad(image, factor=32, border=cv2.BORDER_CONSTANT) | |
| # x = transform(image=padded_image)["image"] | |
| # x = torch.unsqueeze(tensor_from_rgb_image(x), 0).to("cuda") | |
| # with torch.no_grad(): | |
| # prediction = model(x)[0][0] | |
| # mask = (prediction > 0).cpu().numpy().astype(np.uint8) | |
| # mask = unpad(mask, pads) | |
| # rmask = (cv2.cvtColor(mask, cv2.COLOR_BGR2RGB) * 255).astype(np.uint8) | |
| # mask2 = np.where((rmask < 255), 0, 1).astype('uint8') | |
| # image_2_extract = image_2_extract * mask2[:, :, 1, np.newaxis] | |
| # tmp = cv2.cvtColor(image_2_extract, cv2.COLOR_BGR2GRAY) | |
| # _, alpha = cv2.threshold(tmp, 0, 255, cv2.THRESH_BINARY) | |
| # b, g, r = cv2.split(image_2_extract) | |
| # rgba = [b, g, r, alpha] | |
| # dst = cv2.merge(rgba, 4) | |
| # output_image_path = os.path.join(output_dir, image_filename.replace(".jpg", ".png")) | |
| # cv2.imwrite(output_image_path, dst) | |
| # # break | |
| import os | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| import albumentations as albu | |
| from iglovikov_helper_functions.utils.image_utils import load_rgb, pad, unpad | |
| from iglovikov_helper_functions.dl.pytorch.utils import tensor_from_rgb_image | |
| from cloths_segmentation.pre_trained_models import create_model | |
| from tqdm import tqdm | |
| # Create the model and wrap it with DataParallel | |
| model = create_model("Unet_2020-10-30") | |
| model = torch.nn.DataParallel(model) | |
| # Move the model to CUDA devices | |
| model.to("cuda") | |
| model.eval() | |
| input_dir = "../../image" | |
| output_dir = "../../seg_masks" | |
| os.makedirs(output_dir, exist_ok=True) | |
| for image_filename in tqdm(os.listdir(input_dir), colour="green"): | |
| image_path = os.path.join(input_dir, image_filename) | |
| image = cv2.imread(image_path) | |
| image_2_extract = image | |
| image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) | |
| transform = albu.Compose([albu.Normalize(p=1)], p=1) | |
| padded_image, pads = pad(image, factor=32, border=cv2.BORDER_CONSTANT) | |
| x = transform(image=padded_image)["image"] | |
| x = torch.unsqueeze(tensor_from_rgb_image(x), 0).to("cuda") | |
| with torch.no_grad(): | |
| # Use DataParallel to perform inference on all 4 GPUs | |
| prediction = model(x)[0][0] | |
| mask = (prediction > 0).cpu().numpy().astype(np.uint8) | |
| mask = unpad(mask, pads) | |
| rmask = (cv2.cvtColor(mask, cv2.COLOR_BGR2RGB) * 255).astype(np.uint8) | |
| mask2 = np.where((rmask < 255), 0, 1).astype('uint8') | |
| image_2_extract = image_2_extract * mask2[:, :, 1, np.newaxis] | |
| tmp = cv2.cvtColor(image_2_extract, cv2.COLOR_BGR2GRAY) | |
| _, alpha = cv2.threshold(tmp, 0, 255, cv2.THRESH_BINARY) | |
| b, g, r = cv2.split(image_2_extract) | |
| rgba = [b, g, r, alpha] | |
| dst = cv2.merge(rgba, 4) | |
| output_image_path = os.path.join(output_dir, image_filename.replace(".jpg", ".png")) | |
| cv2.imwrite(output_image_path, dst) | |
| # break | |