def main():
args = get_arguments()
num_classes = dataset_settings[args.dataset]['num_classes']
input_size = dataset_settings[args.dataset]['input_size']
label = dataset_settings[args.dataset]['label']
model = network(num_classes=num_classes, pretrained=None).cuda()
model = nn.DataParallel(model)
state_dict = torch.load(args.restore_weight)
model.load_state_dict(state_dict)
model.eval()
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.406, 0.456, 0.485], std=[0.225, 0.224, 0.229])
])
dataset = SCHPDataset(root=args.input, input_size=input_size, transform=transform)
dataloader = DataLoader(dataset)
if not os.path.exists(args.output):
os.makedirs(args.output)
palette = get_palette(num_classes)
with torch.no_grad():
for idx, batch in enumerate(dataloader):
image, meta = batch
img_name = meta['name'][0]
c = meta['center'].numpy()[0]
s = meta['scale'].numpy()[0]
w = meta['width'].numpy()[0]
h = meta['height'].numpy()[0]
output = model(image.cuda())
upsample = torch.nn.Upsample(size=input_size, mode='bilinear', align_corners=True)
upsample_output = upsample(output)
upsample_output = upsample_output.squeeze()
upsample_output = upsample_output.permute(1, 2, 0) #CHW -> HWC
logits_result = transform_logits(upsample_output.data.cpu().numpy(), c, s, w, h, input_size=input_size)
parsing_result = np.argmax(logits_result, axis=2)
parsing_result_path = os.path.join(args.output, img_name[:-4]+'.png')
# plt.imshow(parsing_result)
# plt.show()
output_img = Image.fromarray(np.asarray(parsing_result, dtype=np.uint8))
output_img.putpalette(palette)
output_img.save(parsing_result_path)
if args.logits:
#logits_result_path = os.path.join(args.output, img_name[:-4] + '.npy')
np_result_path = os.path.join(args.output, img_name[:-4] + '.npy')
#np.save(logits_result_path, logits_result)
np.save(np_result_path, parsing_result)
return
def valid(model, valloader, input_size, num_samples, gpus):
model.eval()
parsing_preds = np.zeros((num_samples, input_size[0], input_size[1]),
dtype=np.uint8)
scales = np.zeros((num_samples, 2), dtype=np.float32)
centers = np.zeros((num_samples, 2), dtype=np.int32)
widths = np.zeros((num_samples), dtype=np.int32)
heights = np.zeros((num_samples), dtype=np.int32)
idx = 0
interp = torch.nn.Upsample(size=(input_size[0], input_size[1]), mode='bilinear', align_corners=True)
with torch.no_grad():
for index, batch in enumerate(valloader):
image, _, _, meta, heatmap_params = batch
num_images = image.size(0)
if index % 10 == 0:
print('%d processd' % (index * num_images))
c = meta['center'].numpy()
s = meta['scale'].numpy()
w = meta['width'].numpy()
h = meta['height'].numpy()
scales[idx:idx + num_images, :] = s[:, :]
centers[idx:idx + num_images, :] = c[:, :]
widths[idx:idx + num_images] = w
heights[idx:idx + num_images] = h
outputs = model(image.cuda(), heatmap_params)
if gpus > 1:
for output in outputs:
parsing = output[0][-1]
nums = len(parsing)
print(nums)
print(parsing.shape)
parsing = interp(parsing).data.cpu().numpy()
parsing = parsing.transpose(0, 2, 3, 1) # NCHW NHWC
for i_ in range(nums):
parsing_preds[idx+i_] = transform_logits(parsing.data.cpu().numpy(), centers[idx+i_], scales[idx+i_], heights[idx+i_], widths[idx+i_], input_size=input_size)
parsing_preds[idx+i_] = np.asarray(np.argmax(parsing_preds[idx+i_] , axis=2), dtype=np.uint8)
idx += nums
else:
parsing = outputs[0][-1]
parsing = interp(parsing).data.cpu().numpy()
parsing = parsing.transpose(0, 2, 3, 1) # NCHW NHWC
parsing_preds[idx:idx + num_images, :, :] = np.asarray(np.argmax(parsing, axis=3), dtype=np.uint8)
idx += num_images
parsing_preds = parsing_preds[:num_samples, :, :]
return parsing_preds, scales, centers
def main():
args = get_arguments()
num_classes = dataset_settings[args.dataset]['num_classes']
input_size = dataset_settings[args.dataset]['input_size']
label = dataset_settings[args.dataset]['label']
model = network(num_classes=num_classes, pretrained=None).cuda()
model = nn.DataParallel(model)
state_dict = torch.load(args.restore_weight)
model.load_state_dict(state_dict)
model.eval()
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.406, 0.456, 0.485],
std=[0.225, 0.224, 0.229])
])
INPUT_DIR = os.path.join(args.root, args.input)
dataset = SCHPDataset(root=INPUT_DIR,
input_size=input_size,
transform=transform)
dataloader = DataLoader(dataset)
#palette = get_palette(num_classes)
#mask_palette = get_palette_mask(num_classes)
palette = get_palette(20)
mask_palette = get_palette_mask(20)
full_mask_palette = get_palette_fullP_mask(20)
with torch.no_grad():
for idx, batch in enumerate(dataloader):
image, meta = batch
img_name = meta['name'][0]
c = meta['center'].numpy()[0]
s = meta['scale'].numpy()[0]
w = meta['width'].numpy()[0]
h = meta['height'].numpy()[0]
output = model(image.cuda())
upsample = torch.nn.Upsample(size=input_size,
mode='bilinear',
align_corners=True)
upsample_output = upsample(output)
upsample_output = upsample_output.squeeze()
upsample_output = upsample_output.permute(1, 2, 0) #CHW -> HWC
logits_result = transform_logits(
upsample_output.data.cpu().numpy(),
c,
s,
w,
h,
input_size=input_size)
parsing_result = np.argmax(logits_result, axis=2)
OUTDIR_GRAY = os.path.join(args.root, args.output)
OUTDIR_VIS = os.path.join(args.root, args.output_vis)
#OUTDIR_COLOR=args.root + '/train_color/'
#OUTDIR_WB=args.root + '/train_img_WB/'
#OUTDIR_EDGE=args.root + '/train_edge/'
os.makedirs(OUTDIR_GRAY, exist_ok=True)
os.makedirs(OUTDIR_VIS, exist_ok=True)
#os.makedirs(OUTDIR_COLOR, exist_ok=True)
#os.makedirs(OUTDIR_EDGE, exist_ok=True)
# os.makedirs(OUTDIR_WB, exist_ok=True)
prefix = img_name.split('-')[0]
parsing_result_path = os.path.join(OUTDIR_VIS,
img_name[:-4] + '.png')
parsing_result_path_gray = os.path.join(OUTDIR_GRAY,
img_name[:-4] + '.png')
#cloth_mask_path = os.path.join(OUTDIR_EDGE,img_name[:-4]+'.png')
#white_mask_path = os.path.join(OUTDIR_WB,img_name[:-4]+'.jpg')
#cloth_path = os.path.join(OUTDIR_COLOR,img_name[:-4]+'.jpg')
output_img = Image.fromarray(
np.asarray(parsing_result, dtype=np.uint8))
#original mask all class
original_mask = output_img
original_mask = map_pixel(original_mask, img_name)
original_mask.save(parsing_result_path_gray)
original_mask.putpalette(palette)
original_mask.save(parsing_result_path)
# if not 'front' in img_name: continue
# #new mask with only upper cloth
# new_mask = output_img
# new_mask.putpalette(mask_palette)
# RGB_mask = new_mask.convert('RGB')
# L_mask = new_mask.convert('L')
# new_mask_full = output_img
# new_mask_full.putpalette(full_mask_palette)
# RGB_mask_full = new_mask_full.convert('RGB')
# L_mask_full = new_mask_full.convert('L')
# original_image_path = os.path.join(INPUT_DIR, img_name)
# original_image = Image.open(original_image_path)
## #original_save_path = os.path.join(OUTDIR ,img_name[:-4]+'.png')
## #original_image.save(original_save_path)
# masked_image = ImageChops.multiply(original_image, RGB_mask)
# masked_image_full = ImageChops.multiply(original_image, RGB_mask_full)
## #reverse_mask = ImageOps.invert(RGB_mask)
## #white = Image.new("RGB", original_image.size, "white")
## #background = ImageChops.multiply(white, reverse_mask)
## #masked_image.paste(white, (0,0), reverse_mask)
# bg_image_full = Image.new("RGBA", masked_image_full.size,(255,255,255,255))
# white_masked_image = bg_image_full.paste(masked_image_full.convert('RGBA'),(0,0),L_mask_full)
# bg_image_full = bg_image_full.convert("RGB")
# bg_image_full.save(white_mask_path)
# if not '-p-' in img_name: continue
# bg_image = Image.new("RGBA", masked_image.size,(255,255,255,255))
# white_masked_image = bg_image.paste(masked_image.convert('RGBA'),(0,0),L_mask)
# bg_image = bg_image.convert("RGB")
# L_mask.save(cloth_mask_path)
# bg_image.save(cloth_path)
if args.logits:
logits_result_path = os.path.join(args.root,
img_name[:-4] + '.npy')
np.save(logits_result_path, logits_result)
return
def main():
args = get_arguments()
num_classes = dataset_settings[args.dataset]['num_classes']
input_size = dataset_settings[args.dataset]['input_size']
label = dataset_settings[args.dataset]['label']
model = network(num_classes=num_classes, pretrained=None).cuda()
model = nn.DataParallel(model)
state_dict = torch.load(args.restore_weight)
model.load_state_dict(state_dict)
model.eval()
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.406, 0.456, 0.485],
std=[0.225, 0.224, 0.229])
])
dataset = SCHPDataset(root=args.input,
input_size=input_size,
transform=transform)
dataloader = DataLoader(dataset)
palette = get_palette(num_classes)
print(f"Found {len(dataloader)} files")
with torch.no_grad():
for idx, batch in tqdm(enumerate(dataloader), total=len(dataloader)):
image, meta = batch
img_name = meta['name'][0]
img_path = meta['path'][0]
c = meta['center'].numpy()[0]
s = meta['scale'].numpy()[0]
w = meta['width'].numpy()[0]
h = meta['height'].numpy()[0]
output = model(image.cuda())
upsample = torch.nn.Upsample(size=input_size,
mode='bilinear',
align_corners=True)
upsample_output = upsample(output)
upsample_output = upsample_output.squeeze()
upsample_output = upsample_output.permute(1, 2, 0) #CHW -> HWC
logits_result = transform_logits(
upsample_output.data.cpu().numpy(),
c,
s,
w,
h,
input_size=input_size)
parsing_result = np.argmax(logits_result, axis=2)
img_subpath = img_path.replace(args.input, "").lstrip("/")
parsing_result_path = os.path.join(
args.output, img_subpath[:-4] + args.postfix_filename + '.png')
os.makedirs(os.path.dirname(parsing_result_path), exist_ok=True)
output_img = Image.fromarray(
np.asarray(parsing_result, dtype=np.uint8))
output_img.putpalette(palette) # colors the labels it seems
output_img.save(parsing_result_path)
if args.logits:
fname = img_name[:-4]
logits_result_path = os.path.join(
args.output,
img_subpath[:-4] + args.postfix_filename + '.npy')
if args.argmax_logits:
logits_result_path += "c" # c for compressed
result = parsing_result
else:
result = logits_result
np.save(logits_result_path, result)
return
def main():
args = get_arguments()
num_classes = dataset_settings[args.dataset]['num_classes']
input_size = dataset_settings[args.dataset]['input_size']
label = dataset_settings[args.dataset]['label']
model = network(num_classes=num_classes, pretrained=None).cuda()
model = nn.DataParallel(model)
state_dict = torch.load(args.restore_weight)
model.load_state_dict(state_dict)
model.eval()
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.406, 0.456, 0.485],
std=[0.225, 0.224, 0.229])
])
subdirs = [x[1] for x in os.walk(args.input)]
for subdir in subdirs[0]:
DIR = args.input + '/' + subdir
dataset = SCHPDataset(root=DIR,
input_size=input_size,
transform=transform)
dataloader = DataLoader(dataset)
PARSING_DIR = args.input + '/test_label'
os.makedirs(PARSING_DIR, exist_ok=True)
CLOTHING_DIR = args.input + '/test_color'
os.makedirs(CLOTHING_DIR, exist_ok=True)
CLOTHMASK_DIR = args.input + '/test_edge'
os.makedirs(CLOTHMASK_DIR, exist_ok=True)
palette = get_palette(num_classes)
mask_palette = get_palette_mask(num_classes)
with torch.no_grad():
for idx, batch in enumerate(dataloader):
image, meta = batch
img_name = meta['name'][0]
c = meta['center'].numpy()[0]
s = meta['scale'].numpy()[0]
w = meta['width'].numpy()[0]
h = meta['height'].numpy()[0]
output = model(image.cuda())
upsample = torch.nn.Upsample(size=input_size,
mode='bilinear',
align_corners=True)
upsample_output = upsample(output)
upsample_output = upsample_output.squeeze()
upsample_output = upsample_output.permute(1, 2, 0) #CHW -> HWC
logits_result = transform_logits(
upsample_output.data.cpu().numpy(),
c,
s,
w,
h,
input_size=input_size)
parsing_result = np.argmax(logits_result, axis=2)
parsing_result_path = os.path.join(PARSING_DIR,
img_name[:-4] + '.png')
cloth_mask_path = os.path.join(CLOTHMASK_DIR,
img_name[:-4] + '.jpg')
cloth_path = os.path.join(CLOTHING_DIR, img_name[:-4] + '.png')
output_img = Image.fromarray(
np.asarray(parsing_result, dtype=np.uint8))
#original mask all class
original_mask = output_img
original_mask.save(parsing_result_path)
# original_mask.putpalette(palette)
# original_mask.save(parsing_result_path)
if not '-p-' in img_name: continue
#new mask with only upper cloth
new_mask = output_img
new_mask.putpalette(mask_palette)
RGB_mask = new_mask.convert('RGB')
L_mask = new_mask.convert('L')
original_image_path = os.path.join(DIR, img_name)
print('original_image_path is ', original_image_path)
original_image = Image.open(original_image_path)
#original_save_path = os.path.join(OUTDIR ,img_name[:-4]+'.png')
#original_image.save(original_save_path)
masked_image = ImageChops.multiply(original_image, RGB_mask)
bg_image = Image.new("RGBA", masked_image.size,
(255, 255, 255, 255))
white_masked_image = bg_image.paste(
masked_image.convert('RGBA'), (0, 0), L_mask)
#reverse_mask = ImageOps.invert(RGB_mask)
#white = Image.new("RGB", original_image.size, "white")
#background = ImageChops.multiply(white, reverse_mask)
#masked_image.paste(white, (0,0), reverse_mask)
L_mask.save(cloth_mask_path)
bg_image.save(cloth_path)
if args.logits:
logits_result_path = os.path.join(args.output,
img_name[:-4] + '.npy')
np.save(logits_result_path, logits_result)
return
def main():
args = get_arguments()
os.environ["CUDA_VISIBLE_DEVICES"] = args.gpu
gpus = [int(i) for i in args.gpu.split(',')]
num_classes = dataset_settings[args.dataset]['num_classes']
input_size = dataset_settings[args.dataset]['input_size']
label = dataset_settings[args.dataset]['label']
model = network(num_classes=num_classes,
pretrained=None,
with_my_bn=args.with_my_bn,
batch_size=args.batch_size).cuda()
model = nn.DataParallel(model)
state_dict = torch.load(args.restore_weight)
model.load_state_dict(state_dict)
model.eval()
transform = transforms.Compose([
transforms.ToTensor(),
transforms.Normalize(mean=[0.406, 0.456, 0.485],
std=[0.225, 0.224, 0.229])
])
lip_dataset = LIPDataSet(args.data_dir,
'val',
crop_size=input_size,
transform=transform,
rotation_factor=0,
flip_prob=-0.1,
scale_factor=0)
# dataset = SCHPDataset(root=args.input, input_size=input_size, transform=transform)
dataloader = DataLoader(lip_dataset,
batch_size=args.batch_size * len(gpus),
drop_last=True)
if not os.path.exists(args.output):
os.makedirs(args.output)
palette = get_palette(num_classes)
with torch.no_grad():
for idx, batch in enumerate(dataloader):
image, meta, heatmaps = batch
img_names = meta['name']
cs = meta['center'].numpy()
ss = meta['scale'].numpy()
ws = meta['width'].numpy()
hs = meta['height'].numpy()
print(image.shape)
output = model(image.cuda(), heatmaps)
print('output shape', len(output), len(output[0]),
len(output[0][1]))
output = output[0][1]
print(output.shape)
upsample = torch.nn.Upsample(size=input_size,
mode='bilinear',
align_corners=True)
upsample_outputs = upsample(output)
upsample_outputs = upsample_outputs.squeeze()
# print('upsample_output ', upsample_output.shape)
for upsample_output, img_name, c, s, w, h in zip(
upsample_outputs, img_names, cs, ss, ws, hs):
upsample_output = upsample_output.permute(1, 2, 0) #CHW -> HWC
logits_result = transform_logits(
upsample_output.data.cpu().numpy(),
c,
s,
w,
h,
input_size=input_size)
parsing_result = np.argmax(logits_result, axis=2)
parsing_result_path = os.path.join(args.output,
img_name[:-4] + '.png')
output_img = Image.fromarray(
np.asarray(parsing_result, dtype=np.uint8))
output_img.putpalette(palette)
output_img.save(parsing_result_path)
if args.logits:
logits_result_path = os.path.join(args.output,
img_name[:-4] + '.npy')
np.save(logits_result_path, logits_result)
return
