def main():
seed = 0
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
args = parse_args()
cfg = Config.fromfile(args.config)
img_tensor, w, h = get_img_tensor(args.input, args.use_cuda, get_size=True)
# build model and load checkpoint
model = build_landmark_detector(cfg.model)
print('model built')
load_checkpoint(model, args.checkpoint)
print('load checkpoint from: {}'.format(args.checkpoint))
if args.use_cuda:
model.cuda()
# detect landmark
model.eval()
pred_vis, pred_lm = model(img_tensor, return_loss=False)
pred_lm = pred_lm.data.cpu().numpy()
vis_lms = []
for i, vis in enumerate(pred_vis):
if vis >= 0.5:
print('detected landmark {} {}'.format(pred_lm[i][0] * (w / 224.),
pred_lm[i][1] * (h / 224.)))
vis_lms.append(pred_lm[i])
draw_landmarks(args.input, vis_lms)
def calc_landmark(input_img_path):
seed = 0
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
cfg = Config.fromfile("configs/landmark_detect/landmark_detect_resnet.py")
img_tensor, w, h = get_img_tensor(input_img_path, True, get_size=True)
# build model and load checkpoint
model = build_landmark_detector(cfg.model)
print('model built')
load_checkpoint(model, "checkpoint/Landmark/epoch_50.pth")
print(
'load checkpoint from: {}'.format("checkpoint/Landmark/epoch_50.pth"))
model.cuda()
# detect landmark
model.eval()
pred_vis, pred_lm = model(img_tensor, return_loss=False)
pred_lm = pred_lm.data.cpu().numpy()
vis_lms = []
for i, vis in enumerate(pred_vis):
if vis >= 0.5:
print('detected landmark {} {}'.format(pred_lm[i][0] * (w / 224.),
pred_lm[i][1] * (h / 224.)))
vis_lms.append(pred_lm[i])
return vis_lms
def _init_models():
args = parse_args()
# Build retrieval model and load checkpoint
cfg = mmcv.Config.fromfile(args.config_retrieval)
model_rt = build_retriever(cfg.model)
load_checkpoint(model_rt, args.checkpoint_retrieval)
print('load retriever checkpoint from {}'.format(
args.checkpoint_retrieval))
# Build landmark detection model and load checkpoint
cfg_lm = mmcv.Config.fromfile(args.config_landmark)
model_lm = build_landmark_detector(cfg_lm.model)
load_checkpoint(model_lm, args.checkpoint_landmark)
print('load landmark detector checkpoint from: {}'.format(
args.checkpoint_landmark))
if args.use_cuda:
model_rt.cuda()
model_lm.cuda()
model_rt.eval()
model_lm.eval()
# Build database for retrieval
gallery_list = np.load(args.image_list)
gallery_embeds = _process_embeds(args.image_embeddings)
retriever = ClothesRetriever(gallery_list, [args.topk])
print('build database for retrieval')
# Return retrieval, landmark, and detection model, database for retrieval and retriever
return model_rt, model_lm, gallery_embeds, retriever
def add_landmarks(config, queried_images_path, checkpoint, use_cuda=True):
"""
Function to save predicted for queried image landmarks to json file with the same name as an image.
Receives config for model, path with images, path to downloaded checkpoint.
"""
seed = 0
torch.manual_seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
cfg = Config.fromfile(config)
# build model and load checkpoint
model = build_landmark_detector(cfg.model)
print('model built')
load_checkpoint(model, checkpoint)
print('load checkpoint from: {}'.format(checkpoint))
if use_cuda:
model.cuda()
# detect landmark
model.eval()
images = glob.glob(queried_images_path + "/*.jpg")
for image in images:
image_id = image.split('/')[-1].split('.')[0]
img_tensor, w, h = get_img_tensor(image, use_cuda, get_size=True)
pred_vis, pred_lm = model(img_tensor, return_loss=False)
pred_lm = pred_lm.data.cpu().numpy()
vis_lms = {'landmarks': []}
for i, vis in enumerate(pred_vis):
if vis >= 0.5:
print('detected landmarks: {} {} for image {}'.format(
pred_lm[i][0] * (w / 224.), pred_lm[i][1] * (h / 224.),
image_id))
vis_lms['landmarks'].append(pred_lm[i].tolist())
if vis_lms['landmarks']:
json.dump(
vis_lms,
open(queried_images_path + '/' + str(image_id) + '.json',
'w'))
return queried_images_path
def main():
args = parse_args()
cfg = Config.fromfile(args.config)
if args.work_dir is not None:
cfg.work_dir = args.work_dir
if args.resume_from is not None:
cfg.resume_from = args.resume_from
# init distributed env first
if args.launcher == 'none':
distributed = False
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
# init logger
logger = get_root_logger(cfg.log_level)
logger.info('Distributed training: {}'.format(distributed))
# set random seeds
if args.seed is not None:
logger.info('Set random seed to {}'.format(args.seed))
set_random_seed(args.seed)
# build model
model = build_landmark_detector(cfg.model)
print('model built')
if cfg.init_weights_from:
model = init_weights_from(cfg.init_weights_from, model)
# data loader
dataset = get_dataset(cfg.data.train)
# train
train_landmark_detector(
model,
dataset,
cfg,
distributed=distributed,
validate=args.validate,
logger=logger)
def main():
args = parse_args()
cfg = Config.fromfile(args.config)
if args.work_dir is not None:
cfg.work_dir = args.work_dir
# init distributed env first
if args.launcher == 'none':
distributed = False
else:
distributed = True
init_dist(args.launcher, **cfg.dist_params)
if args.checkpoint is not None:
cfg.load_from = args.checkpoint
# init logger
logger = get_root_logger(cfg.log_level)
logger.info('Distributed test: {}'.format(distributed))
# data loader
test_dataset = get_dataset(cfg.data.test)
print('dataset loaded')
# build model and load checkpoint
model = build_landmark_detector(cfg.model)
print('model built')
checkpoint = load_checkpoint(model, cfg.load_from, map_location='cpu')
print('load checkpoint from: {}'.format(cfg.load_from))
# test
test_landmark_detector(
model,
test_dataset,
cfg,
distributed=distributed,
validate=args.validate,
logger=logger)
def main():
seed = 0
torch.manual_seed(seed)
random.seed(seed)
np.random.seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
args = parse_args()
# Build retrieval model and load checkpoint
cfg = mmcv.Config.fromfile(args.config_retrieval)
model = build_retriever(cfg.model)
load_checkpoint(model, args.checkpoint_retrieval)
print('load retriever checkpoint from {}'.format(
args.checkpoint_retrieval))
# Build landmark detection model and load checkpoint
cfg_lm = mmcv.Config.fromfile(args.config_landmark)
model_lm = build_landmark_detector(cfg_lm.model)
load_checkpoint(model_lm, args.checkpoint_landmark)
print('load landmark detector checkpoint from: {}'.format(
args.checkpoint_landmark))
if args.use_cuda:
model.cuda()
model_lm.cuda()
model.eval()
model_lm.eval()
# Build model for detection
config = DetectionConfig()
model_dt = MaskRCNN(mode="inference", config=config, model_dir="None")
model_dt.load_weights(args.detection_weights, by_name=True)
# Import image list to be processed
img_list = np.load(args.img_list)
# Initiate empty list to store generated embeddings
embeds = []
# filename of the output
filename = os.path.join(args.save_dir, args.output_fn)
iteration = 0
for idx, img_path in enumerate(img_list):
#if idx <= 3900:
# continue
# Import image
img_path = './static/images' + img_path[1:]
img = skimage.io.imread(img_path)
# Extract class name
class_name = img_path.split("/")[5]
selected_class = None
if (len(img.shape) == 2) or (img.shape[2] == 1):
img = cv2.cvtColor(img, cv2.COLOR_GRAY2RGB)
# Perform detection
dt_results = model_dt.detect([img], verbose=1)[0]
# If no detected clothing
if len(dt_results['rois']) == 0:
# Create tensor from query image
img_roi = Image.fromarray(img)
img_tensor = get_img_tensor(img_roi, args.use_cuda)
# Generate landmark from image tensor
_, pred_lm = model_lm(img_tensor, return_loss=False)
# Generate embedding
embed = model(img_tensor, landmark=pred_lm, return_loss=False)
selected_roi = np.array([0, 0, 0, 0])
elif ((class_name == "Atasan") or
(class_name == "Baju_Hamil") or
(class_name == "Baju_Tidur___Pakaian_Dalam") or
(class_name == "Blazer") or
(class_name == "Dress_copy") or
(class_name == "Hoodies___Sweatshirts") or
(class_name == "Jaket___Coat") or
(class_name == "Kaos_Kaki___Stocking") or
(class_name == "KnitWear___Cardigan") or
(class_name == "Playsuits___Jumpsuits") or
(class_name == "Plus_Size") or
(class_name == "Swimwear___Beachwear")) and\
((1 not in dt_results["class_ids"]) and
(2 not in dt_results["class_ids"]) and
(3 not in dt_results["class_ids"]) and
(4 not in dt_results["class_ids"]) and
(5 not in dt_results["class_ids"]) and
(6 not in dt_results["class_ids"]) and
(10 not in dt_results["class_ids"]) and
(11 not in dt_results["class_ids"]) and
(12 not in dt_results["class_ids"]) and
(13 not in dt_results["class_ids"])):
# Create tensor from query image
img_roi = Image.fromarray(img)
img_tensor = get_img_tensor(img_roi, args.use_cuda)
# Generate landmark from image tensor
_, pred_lm = model_lm(img_tensor, return_loss=False)
# Generate embedding
embed = model(img_tensor, landmark=pred_lm, return_loss=False)
selected_roi = np.array([0, 0, 0, 0])
elif ((class_name == "Celana_Pendek") or
(class_name == "Pants___Leggings") or
(class_name == "Jeans") or
(class_name == "Rok")) and\
((7 not in dt_results["class_ids"]) and
(8 not in dt_results["class_ids"]) and
(9 not in dt_results["class_ids"])):
# Create tensor from query image
img_roi = Image.fromarray(img)
img_tensor = get_img_tensor(img_roi, args.use_cuda)
# Generate landmark from image tensor
_, pred_lm = model_lm(img_tensor, return_loss=False)
# Generate embedding
embed = model(img_tensor, landmark=pred_lm, return_loss=False)
selected_roi = np.array([0, 0, 0, 0])
else:
max_area = 0
selected_roi = 0
selected_class = 0
# Iterate clothes retrieval for every detected ROI
for i, roi in enumerate(dt_results['rois']):
class_id = int(dt_results["class_ids"][i])
if ((class_name == "Atasan") or
(class_name == "Baju_Hamil") or
(class_name == "Baju_Tidur___Pakaian_Dalam") or
(class_name == "Blazer") or
(class_name == "Dress_copy") or
(class_name == "Hoodies___Sweatshirts") or
(class_name == "Jaket___Coat") or
(class_name == "Kaos_Kaki___Stocking") or
(class_name == "KnitWear___Cardigan") or
(class_name == "Playsuits___Jumpsuits") or
(class_name == "Plus_Size") or
(class_name == "Swimwear___Beachwear")) and\
(class_id != 7 and class_id != 8 and class_id != 9):
area = (roi[2] - roi[0]) * (roi[3] - roi[1])
if area < max_area:
continue
max_area = area
selected_roi = roi
selected_class = class_id
# Cropped input image based on the detected ROI
img_crop = img[roi[0]:roi[2], roi[1]:roi[3], :]
# Create tensor from query image
img_roi = Image.fromarray(img_crop)
img_tensor = get_img_tensor(img_roi, args.use_cuda)
# Generate landmark from query image tensor
_, pred_lm = model_lm(img_tensor, return_loss=False)
# Generate embedding
embed = model(img_tensor,
landmark=pred_lm,
return_loss=False)
elif ((class_name == "Celana_Pendek") or
(class_name == "Pants___Leggings") or
(class_name == "Jeans") or
(class_name == "Rok")) and\
(class_id == 7 or class_id == 8 or class_id == 9):
area = (roi[2] - roi[0]) * (roi[3] - roi[1])
if area < max_area:
continue
max_area = area
selected_roi = roi
selected_class = class_id
# Cropped input image based on the detected ROI
img_crop = img[roi[0]:roi[2], roi[1]:roi[3], :]
# Create tensor from query image
img_roi = Image.fromarray(img_crop)
img_tensor = get_img_tensor(img_roi, args.use_cuda)
# Generate landmark from query image tensor
_, pred_lm = model_lm(img_tensor, return_loss=False)
# Generate embedding
embed = model(img_tensor,
landmark=pred_lm,
return_loss=False)
print(selected_roi)
start_point = (selected_roi[1], selected_roi[0])
end_point = (selected_roi[3], selected_roi[2])
img = cv2.rectangle(img, start_point, end_point, (51, 51, 255), 2)
#img = cv2.putText(img, selected_class, (0, 0), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1, cv2.LINE_AA)
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
img = cv2.resize(img, (400, int(img.shape[0] * 400 / img.shape[1])),
interpolation=cv2.INTER_AREA)
print(class_name, selected_class)
cv2.imshow("detection results", img)
cv2.waitKey(1000)
embeds.append(embed.detach().numpy())
print(str(idx) + ". Finish generating embedding for " + img_path)
if (idx != 0) and (idx % 100 == 0):
iteration += 1
np.save(filename + '_' + str(iteration), embeds)
print('Embeddings {} are saved'.format(iteration))
embeds = []
iteration += 1
np.save(filename + '_' + str(iteration), embeds)
print('Embeddings {} are saved'.format(iteration))
embeds = []
embed_combine = []
for i in range(iteration):
embeds = np.load(filename + '_' + str(i + 1) + '.npy')
for embed in embeds:
embed_combine.append(embed)
np.save(filename, embed_combine)
print("All embeddings have been combined")
