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 main():
args = parse_args()
cfg = Config.fromfile(args.config)
img_tensor = get_img_tensor(args.input, args.use_cuda)
# global attribute predictor will not use landmarks
# just set a default value
landmark_tensor = torch.zeros(8)
print(cfg.model)
model = build_predictor(cfg.model)
load_checkpoint(model, args.checkpoint, map_location='cpu')
print('model loaded from {}'.format(args.checkpoint))
if args.use_cuda:
model.cuda()
landmark_tensor = landmark_tensor.cuda()
model.eval()
# predict probabilities for each attribute
attr_prob, cate_prob = model(
img_tensor, attr=None, landmark=landmark_tensor, return_loss=False)
attr_predictor = AttrPredictor(cfg.data.test)
cate_predictor = CatePredictor(cfg.data.test)
attr_predictor.show_prediction(attr_prob)
cate_predictor.show_prediction(cate_prob)
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 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)
cfg.load_from = args.checkpoint
# init distributed env first
distributed = False
# init logger
logger = get_root_logger(cfg.log_level)
logger.info('Distributed test: {}'.format(distributed))
# create model
model = build_fashion_recommender(cfg.model)
load_checkpoint(model, cfg.load_from, map_location='cpu')
print('load checkpoint from: {}'.format(cfg.load_from))
if args.use_cuda:
model.cuda()
model.eval()
# prepare input data
img_tensors = []
item_ids = []
for dirpath, dirname, fns in os.walk(args.input_dir):
for imgname in fns:
item_ids.append(imgname.split('.')[0])
tensor = get_img_tensor(
os.path.join(dirpath, imgname), args.use_cuda)
img_tensors.append(tensor)
img_tensors = torch.cat(img_tensors)
# test
embeds = []
with torch.no_grad():
embed = model(img_tensors, return_loss=False)
embeds.append(embed.data.cpu())
embeds = torch.cat(embeds)
try:
metric = model.module.triplet_net.metric_branch
except Exception:
metric = None
# get compatibility score, so far only support images from polyvore
dataset = build_dataset(cfg.data.test)
score = dataset.get_single_compatibility_score(embeds, item_ids, metric,
args.use_cuda)
print("Compatibility score: {:.3f}".format(score))
def retrieve():
file = request.files.get('image')
img_tensor = get_img_tensor(file, True)
query_feat = model_ret(img_tensor, landmark=None, return_loss=False)
query_feat = query_feat.data.cpu().numpy()
gallery_set = build_dataset(cfg_ret.data.gallery)
gallery_embeds = _process_embeds(gallery_set, model_ret, cfg_ret)
retriever = ClothesRetriever(cfg_ret.data.gallery.img_file, cfg_ret.data_root,
cfg_ret.data.gallery.img_path)
result = retriever.show_retrieved_images(query_feat, gallery_embeds)
resultDict = {}
resultDict['paths'] = result
return jsonify(resultDict)
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)
img_tensor = get_img_tensor(args.input, args.use_cuda)
cfg.model.pretrained = None
model = build_predictor(cfg.model)
load_checkpoint(model, args.checkpoint, map_location='cpu')
if args.use_cuda:
model.cuda()
model.eval()
# predict probabilities for each attribute
attr_prob = model(img_tensor, attr=None, landmark=None, return_loss=False)
attr_predictor = AttrPredictor(cfg.data.test)
attr_predictor.show_prediction(attr_prob)
def main():
args = parse_args()
cfg = Config.fromfile(args.config)
cfg.model.pretrained = None
model = build_predictor(cfg.model)
load_checkpoint(model, args.checkpoint, map_location='cpu')
if args.use_cuda:
model.cuda()
model.eval()
attr_predictor = AttrPredictor(cfg.data.test)
files=os.listdir(args.input)
for i in tqdm(range(len(files))) :
filename=files[i]
#print(filename)
filename=args.input+'/'+filename
img_tensor = get_img_tensor(filename, args.use_cuda)
attr_prob = model(img_tensor, attr=None, landmark=None, return_loss=False)
attr_predictor.show_prediction(attr_prob,filename)
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 annotate():
file = request.files.get('image')
img_tensor = get_img_tensor(file, False)
# predict probabilities for each attribute
fine_attr_prob, cate_prob = model_fine(
img_tensor, attr=None, landmark=landmark_tensor, return_loss=False)
coarse_attr_prob = model_coarse(
img_tensor, attr=None, landmark=landmark_tensor, return_loss=False)
coarse_attrs = jsonResult(
coarse_attr_prob, coarse_attr_predictor.attr_idx2name)
cats = jsonResult(cate_prob, cate_predictor.cate_idx2name)
colors = colorgram.extract(file, 5)
color_names = [get_color_name(x.rgb)[1] for x in colors]
resultDict = {}
resultDict['attributes'] = coarse_attrs
resultDict['categories'] = cats
resultDict['colors'] = color_names
return jsonify(resultDict)
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")
def responce():
print('MMFashion')
if (app.debug):
print("flask.request.method : ", flask.request.method)
print("flask.request.headers \n: ", flask.request.headers)
if (flask.request.headers["User-Agent"].split("/")[0]
in "python-requests"):
json_data = json.loads(flask.request.json)
else:
json_data = flask.request.get_json()
img_path = args.input
global model
if args.use_cuda:
model = model.cuda()
#------------------------------------------
# ブラウザから送信された画像データの変換
#------------------------------------------
img_cv = conv_base64_to_cv(json_data["img_base64"])
cv2.imwrite(img_path, img_cv)
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()
retrieved_paths = retriever.show_retrieved_images(query_feat,
gallery_embeds)
retrieved_imgs = []
for retrieved_path in retrieved_paths:
retrieved_path = retrieved_path.replace('data/In-shop/', '', 1)
retrieved_img = Image.open(retrieved_path)
retrieved_img_base64 = conv_pillow_to_base64(retrieved_img)
retrieved_imgs.append(retrieved_img_base64)
#------------------------------------------
# json 形式のレスポンスメッセージを作成
#------------------------------------------
#torch.cuda.empty_cache()
http_status_code = 200
response = flask.jsonify({
'status':
'OK',
'origin_img':
conv_pillow_to_base64(Image.open(img_path)),
'retrieved_imgs':
retrieved_imgs
})
# Access-Control-Allow-Origin
response.headers.add('Access-Control-Allow-Origin', '*')
response.headers.add('Access-Control-Allow-Headers',
'Content-Type,Authorization')
if (app.debug):
print("response.headers : \n", response.headers)
# release gpu
del img_tensor, query_feat, retrieved_paths
if args.use_cuda:
model = model.cpu()
torch.cuda.empty_cache()
return response, http_status_code
def visual_search(img, models, dt_results, det_filename):
args = parse_args()
# Placeholder for results
results = []
categories = []
# Parse the input
model_rt = models['model_rt']
model_lm = models['model_lm']
gallery_embeds = models['gallery_embeds']
retriever = models['retriever']
seed = 0
torch.manual_seed(seed)
random.seed(seed)
np.random.seed(seed)
if torch.cuda.is_available():
torch.cuda.manual_seed_all(seed)
# 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)
class_id = 15
# Generate landmark from query image tensor
_, pred_lm = model_lm(img_tensor, return_loss=False)
# Retrieval
query_feat = model_rt(img_tensor, landmark=pred_lm, return_loss=False)
query_feat = query_feat.data.cpu().numpy()
result = retriever.show_retrieved_images(img, query_feat,
gallery_embeds, class_id,
args.visualize)
results.append(result)
categories.append("Clothes")
return results, categories
# Visualize detection results
visualize.visualize_detection(dt_results, img, det_filename,
args.visualize)
# Iterate clothes retrieval for every detected ROI
for i, roi in enumerate(dt_results['rois']):
# Cropped input image based on the detected ROI
img_crop = img[roi[0]:roi[2], roi[1]:roi[3], :]
# Visualize cropped image
class_id = dt_results["class_ids"][i]
if (class_id <= 6):
categories.append('Top')
elif (class_id >= 10):
categories.append('One-Piece')
else:
categories.append('Bottom')
# 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)
# Retrieval
query_feat = model_rt(img_tensor, landmark=pred_lm, return_loss=False)
query_feat = query_feat.data.cpu().numpy()
result = retriever.show_retrieved_images(img_crop, query_feat,
gallery_embeds, class_id,
args.visualize)
results.append(result)
return results, categories
def get_vector(image_path):
img_tensor = get_img_tensor(image_path, False)
query_feat = model(img_tensor, landmark=None, return_loss=False)
query_feat = query_feat.data.cpu().numpy()
return query_feat
