|
import torch |
|
import numpy as np |
|
|
|
import os |
|
|
|
from skimage import io |
|
import cv2 |
|
xyz_from_rgb = np.array( |
|
[[0.412453, 0.357580, 0.180423], [0.212671, 0.715160, 0.072169], [0.019334, 0.119193, 0.950227]] |
|
) |
|
rgb_from_xyz = np.array( |
|
[[3.24048134, -0.96925495, 0.05564664], [-1.53715152, 1.87599, -0.20404134], [-0.49853633, 0.04155593, 1.05731107]] |
|
) |
|
|
|
|
|
def tensor_lab2rgb(input): |
|
""" |
|
n * 3* h *w |
|
""" |
|
input_trans = input.transpose(1, 2).transpose(2, 3) |
|
L, a, b = input_trans[:, :, :, 0:1], input_trans[:, :, :, 1:2], input_trans[:, :, :, 2:] |
|
y = (L + 16.0) / 116.0 |
|
x = (a / 500.0) + y |
|
z = y - (b / 200.0) |
|
|
|
neg_mask = z.data < 0 |
|
z[neg_mask] = 0 |
|
xyz = torch.cat((x, y, z), dim=3) |
|
|
|
mask = xyz.data > 0.2068966 |
|
mask_xyz = xyz.clone() |
|
mask_xyz[mask] = torch.pow(xyz[mask], 3.0) |
|
mask_xyz[~mask] = (xyz[~mask] - 16.0 / 116.0) / 7.787 |
|
mask_xyz[:, :, :, 0] = mask_xyz[:, :, :, 0] * 0.95047 |
|
mask_xyz[:, :, :, 2] = mask_xyz[:, :, :, 2] * 1.08883 |
|
|
|
rgb_trans = torch.mm(mask_xyz.view(-1, 3), torch.from_numpy(rgb_from_xyz).type_as(xyz)).view( |
|
input.size(0), input.size(2), input.size(3), 3 |
|
) |
|
rgb = rgb_trans.transpose(2, 3).transpose(1, 2) |
|
|
|
mask = rgb > 0.0031308 |
|
mask_rgb = rgb.clone() |
|
mask_rgb[mask] = 1.055 * torch.pow(rgb[mask], 1 / 2.4) - 0.055 |
|
mask_rgb[~mask] = rgb[~mask] * 12.92 |
|
|
|
neg_mask = mask_rgb.data < 0 |
|
large_mask = mask_rgb.data > 1 |
|
mask_rgb[neg_mask] = 0 |
|
mask_rgb[large_mask] = 1 |
|
return mask_rgb |
|
|
|
def get_files(img_dir): |
|
imgs, masks, xmls = list_files(img_dir) |
|
return imgs, masks, xmls |
|
|
|
|
|
def list_files(in_path): |
|
img_files = [] |
|
mask_files = [] |
|
gt_files = [] |
|
for (dirpath, dirnames, filenames) in os.walk(in_path): |
|
for file in filenames: |
|
filename, ext = os.path.splitext(file) |
|
ext = str.lower(ext) |
|
if ext == '.jpg' or ext == '.jpeg' or ext == '.gif' or ext == '.png' or ext == '.pgm': |
|
img_files.append(os.path.join(dirpath, file)) |
|
elif ext == '.bmp': |
|
mask_files.append(os.path.join(dirpath, file)) |
|
elif ext == '.xml' or ext == '.gt' or ext == '.txt': |
|
gt_files.append(os.path.join(dirpath, file)) |
|
elif ext == '.zip': |
|
continue |
|
return img_files, mask_files, gt_files |
|
|
|
|
|
def load_image(img_file): |
|
img = io.imread(img_file) |
|
if img.shape[0] == 2: |
|
img = img[0] |
|
if len(img.shape) == 2: |
|
img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB) |
|
if img.shape[2] == 4: |
|
img = img[:, :, :3] |
|
img = np.array(img) |
|
|
|
return img |
