File size: 3,742 Bytes
e5765b1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 |
# 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
|