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")