# 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 | |