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)
model = build_retriever(cfg.model)
load_checkpoint(model, args.checkpoint)
print('load checkpoint from {}'.format(args.checkpoint))
if args.use_cuda:
model.cuda()
model.eval()
img_tensor = get_img_tensor(args.input, args.use_cuda)
query_feat = model(img_tensor, landmark=None, return_loss=False)
query_feat = query_feat.data.cpu().numpy()
gallery_set = build_dataset(cfg.data.gallery)
gallery_embeds = _process_embeds(gallery_set, model, cfg)
retriever = ClothesRetriever(cfg.data.gallery.img_file, cfg.data_root,
cfg.data.gallery.img_path)
retriever.show_retrieved_images(query_feat, gallery_embeds)
def main():
args = parse_args()
cfg = Config.fromfile(args.config)
img = cv2.imread(args.input, -1)
img = cv2.resize(img, (224, 224))
img_tensor = img_to_tensor(img, squeeze=True, cuda=args.use_cuda)
model = build_retriever(cfg.model)
checkpoint = load_checkpoint(model, args.checkpoint, map_location='cpu')
if args.use_cuda:
model.cuda()
model.eval()
query_embed = model(img_tensor, landmark=None, return_loss=False)
query_embed = query_embed.data.cpu().numpy()
gallery_set = build_dataset(cfg.data.gallery)
gallery_embeds = _process_embeds(gallery_set, model, cfg)
retriever = ClothesRetriever(cfg.data.gallery.img_file, cfg.data_root,
cfg.data.gallery.img_path)
retriever.show_retrieved_images(query_embed, gallery_embeds)
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 extract_features(image_set, cfg, save_feature_dir):
model = build_retriever(cfg.model)
print('model built')
model = MMDataParallel(model, device_ids=cfg.gpus.test).cuda()
model.eval()
embeds = _process_embeds(image_set, model, cfg)
if not os.path.exists(save_feature_dir):
os.makedirs(save_feature_dir)
save_path = os.path.join(save_feature_dir, 'extracted_features.mat')
sio.savemat(save_path, {'embeds': embeds})
print('extracted features saved to : %s' % save_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
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 predictor to extract embeddings
model = build_retriever(cfg.model)
print('model built')
if cfg.init_weights_from is not None:
model = init_weights_from(cfg.init_weights_from, model)
print('Initialize model weights from {}'.format(cfg.init_weights_from))
# data loader
dataset = build_dataset(cfg.data.train)
print('dataset loaded')
# train
train_retriever(
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
cfg.data.query.find_three = False
cfg.data.gallery.find_three = False
query_set, gallery_set = build_dataset(cfg.data.query), build_dataset(
cfg.data.gallery)
print('dataset loaded')
# build model and load checkpoint
model = build_retriever(cfg.model)
print('model built')
load_checkpoint(model, cfg.load_from)
print('load checkpoint from: {}'.format(cfg.load_from))
# test
test_retriever(model,
query_set,
gallery_set,
cfg,
distributed=distributed,
validate=args.validate,
logger=logger)
def main():
seed = 0
torch.manual_seed(seed)
args = parse_args()
if args.use_cuda and torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
cfg = Config.fromfile(args.config)
model = build_retriever(cfg.model)
load_checkpoint(model, args.checkpoint, map_location=torch.device('cuda:0'))
print('load checkpoint from {}'.format(args.checkpoint))
if args.use_cuda:
model.cuda()
model.eval()
print('Model evaled')
img_tensor = get_img_tensor(args.input, args.use_cuda)
print('Image tensor got.')
query_feat = model(img_tensor, landmark=None, return_loss=False)
print('Query feat 1')
query_feat = query_feat.data.cpu().numpy()
print('Query feat 2')
gallery_set = build_dataset(cfg.data.gallery)
print('Gallery set')
gallery_embeds = _process_embeds(gallery_set, model, cfg)
print('Gallery embeds')
retriever = ClothesRetriever(cfg.data.gallery.img_file, cfg.data_root,
cfg.data.gallery.img_path)
print('Retriever')
results = retriever.show_retrieved_images(query_feat, gallery_embeds)
print('Show retriever')
for result in results:
print(result)
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")
from __future__ import division
import argparse
from mmcv import Config
from mmfashion.apis import (get_root_logger, init_dist, set_random_seed,
train_retriever, test_retriever)
from mmfashion.datasets import build_dataset
from mmfashion.models import build_retriever
from mmfashion.utils import init_weights_from
# Load the config from the custom file
# cfg_fname = 'configs/retriever_in_shop/global_retriever_vgg_loss_id_triplet.py' # Triplet network
cfg_fname = 'configs/retriever_in_shop/global_retriever_vgg_loss_id.py' # Plain siamese
cfg = Config.fromfile(cfg_fname)
# Data loader
train_set = build_dataset(cfg.data.train)
query_set = build_dataset(cfg.data.query)
gallery_set = build_dataset(cfg.data.gallery)
print('datasets loaded')
# Build model and load checkpoint
model = build_retriever(cfg.model)
print('model built')
# Train
train_retriever(model, train_set, cfg, distributed=False)
# Test
test_retriever(model, query_set, gallery_set, cfg, distributed=False)
cfg_coarse = Config.fromfile(
'./configs/attribute_predict_coarse/global_predictor_resnet_attr.py')
cfg_ret = Config.fromfile('configs/retriever_in_shop/global_retriever_vgg_loss_id.py')
# global attribute predictor will not use landmarks
# just set a default value
landmark_tensor = torch.zeros(8)
model_fine = build_predictor(cfg_fine.model)
load_checkpoint(model_fine, './checkpoint/vgg16_fine_global.pth',
map_location='cpu')
model_coarse = build_predictor(cfg_coarse.model)
load_checkpoint(model_coarse, './checkpoint/resnet_coarse_global.pth',
map_location='cpu')
model_ret = build_retriever(cfg_ret.model).cuda()
load_checkpoint(model_ret, 'checkpoint/Retrieve/vgg/global/epoch_100.pth', map_location=torch.device('cuda:0'))
print('Models loaded.')
model_fine.eval()
model_coarse.eval()
model_ret.eval()
cate_predictor = CatePredictor(cfg_fine.data.test)
coarse_attr_predictor = AttrPredictor(cfg_coarse.data.test)
app.run(host="0.0.0.0", port=80)
