def download_dataset(self, dataset_dir, dataset_url): """Downloads and extracts dataset. Args: dataset_dir (str): dataset directory. dataset_url (str): url to download dataset. """ if osp.exists(dataset_dir): return if dataset_url is None: raise RuntimeError('{} dataset needs to be manually ' 'prepared, please follow the ' 'document to prepare this dataset'.format( self.__class__.__name__)) print('Creating directory "{}"'.format(dataset_dir)) mkdir_if_missing(dataset_dir) fpath = osp.join(dataset_dir, osp.basename(dataset_url)) print('Downloading {} dataset to "{}"'.format(self.__class__.__name__, dataset_dir)) download_url(dataset_url, fpath) print('Extracting "{}"'.format(fpath)) try: tar = tarfile.open(fpath) tar.extractall(path=dataset_dir) tar.close() except: zip_ref = zipfile.ZipFile(fpath, 'r') zip_ref.extractall(dataset_dir) zip_ref.close() print('{} dataset is ready'.format(self.__class__.__name__))
def preprocess_split(self): # This function is a bit complex and ugly, what it does is # 1. extract data from cuhk-03.mat and save as png images # 2. create 20 classic splits (Li et al. CVPR'14) # 3. create new split (Zhong et al. CVPR'17) if osp.exists(self.imgs_labeled_dir) \ and osp.exists(self.imgs_detected_dir) \ and osp.exists(self.split_classic_det_json_path) \ and osp.exists(self.split_classic_lab_json_path) \ and osp.exists(self.split_new_det_json_path) \ and osp.exists(self.split_new_lab_json_path): return import h5py from scipy.misc import imsave from scipy.io import loadmat mkdir_if_missing(self.imgs_detected_dir) mkdir_if_missing(self.imgs_labeled_dir) print('Extract image data from "{}" and save as png'.format( self.raw_mat_path)) mat = h5py.File(self.raw_mat_path, 'r') def _deref(ref): return mat[ref][:].T def _process_images(img_refs, campid, pid, save_dir): img_paths = [] # Note: some persons only have images for one view for imgid, img_ref in enumerate(img_refs): img = _deref(img_ref) if img.size == 0 or img.ndim < 3: continue # skip empty cell # images are saved with the following format, index-1 (ensure uniqueness) # campid: index of camera pair (1-5) # pid: index of person in 'campid'-th camera pair # viewid: index of view, {1, 2} # imgid: index of image, (1-10) viewid = 1 if imgid < 5 else 2 img_name = '{:01d}_{:03d}_{:01d}_{:02d}.png'.format( campid + 1, pid + 1, viewid, imgid + 1) img_path = osp.join(save_dir, img_name) if not osp.isfile(img_path): imsave(img_path, img) img_paths.append(img_path) return img_paths def _extract_img(image_type): print('Processing {} images ...'.format(image_type)) meta_data = [] imgs_dir = self.imgs_detected_dir if image_type == 'detected' else self.imgs_labeled_dir for campid, camp_ref in enumerate(mat[image_type][0]): camp = _deref(camp_ref) num_pids = camp.shape[0] for pid in range(num_pids): img_paths = _process_images(camp[pid, :], campid, pid, imgs_dir) assert len( img_paths) > 0, 'campid{}-pid{} has no images'.format( campid, pid) meta_data.append((campid + 1, pid + 1, img_paths)) print('- done camera pair {} with {} identities'.format( campid + 1, num_pids)) return meta_data meta_detected = _extract_img('detected') meta_labeled = _extract_img('labeled') def _extract_classic_split(meta_data, test_split): train, test = [], [] num_train_pids, num_test_pids = 0, 0 num_train_imgs, num_test_imgs = 0, 0 for i, (campid, pid, img_paths) in enumerate(meta_data): if [campid, pid] in test_split: for img_path in img_paths: camid = int(osp.basename(img_path).split('_') [2]) - 1 # make it 0-based test.append((img_path, num_test_pids, camid)) num_test_pids += 1 num_test_imgs += len(img_paths) else: for img_path in img_paths: camid = int(osp.basename(img_path).split('_') [2]) - 1 # make it 0-based train.append((img_path, num_train_pids, camid)) num_train_pids += 1 num_train_imgs += len(img_paths) return train, num_train_pids, num_train_imgs, test, num_test_pids, num_test_imgs print('Creating classic splits (# = 20) ...') splits_classic_det, splits_classic_lab = [], [] for split_ref in mat['testsets'][0]: test_split = _deref(split_ref).tolist() # create split for detected images train, num_train_pids, num_train_imgs, test, num_test_pids, num_test_imgs = \ _extract_classic_split(meta_detected, test_split) splits_classic_det.append({ 'train': train, 'query': test, 'gallery': test, 'num_train_pids': num_train_pids, 'num_train_imgs': num_train_imgs, 'num_query_pids': num_test_pids, 'num_query_imgs': num_test_imgs, 'num_gallery_pids': num_test_pids, 'num_gallery_imgs': num_test_imgs }) # create split for labeled images train, num_train_pids, num_train_imgs, test, num_test_pids, num_test_imgs = \ _extract_classic_split(meta_labeled, test_split) splits_classic_lab.append({ 'train': train, 'query': test, 'gallery': test, 'num_train_pids': num_train_pids, 'num_train_imgs': num_train_imgs, 'num_query_pids': num_test_pids, 'num_query_imgs': num_test_imgs, 'num_gallery_pids': num_test_pids, 'num_gallery_imgs': num_test_imgs }) write_json(splits_classic_det, self.split_classic_det_json_path) write_json(splits_classic_lab, self.split_classic_lab_json_path) def _extract_set(filelist, pids, pid2label, idxs, img_dir, relabel): tmp_set = [] unique_pids = set() for idx in idxs: img_name = filelist[idx][0] camid = int(img_name.split('_')[2]) - 1 # make it 0-based pid = pids[idx] if relabel: pid = pid2label[pid] img_path = osp.join(img_dir, img_name) tmp_set.append((img_path, int(pid), camid)) unique_pids.add(pid) return tmp_set, len(unique_pids), len(idxs) def _extract_new_split(split_dict, img_dir): train_idxs = split_dict['train_idx'].flatten() - 1 # index-0 pids = split_dict['labels'].flatten() train_pids = set(pids[train_idxs]) pid2label = {pid: label for label, pid in enumerate(train_pids)} query_idxs = split_dict['query_idx'].flatten() - 1 gallery_idxs = split_dict['gallery_idx'].flatten() - 1 filelist = split_dict['filelist'].flatten() train_info = _extract_set(filelist, pids, pid2label, train_idxs, img_dir, relabel=True) query_info = _extract_set(filelist, pids, pid2label, query_idxs, img_dir, relabel=False) gallery_info = _extract_set(filelist, pids, pid2label, gallery_idxs, img_dir, relabel=False) return train_info, query_info, gallery_info print('Creating new split for detected images (767/700) ...') train_info, query_info, gallery_info = _extract_new_split( loadmat(self.split_new_det_mat_path), self.imgs_detected_dir) split = [{ 'train': train_info[0], 'query': query_info[0], 'gallery': gallery_info[0], 'num_train_pids': train_info[1], 'num_train_imgs': train_info[2], 'num_query_pids': query_info[1], 'num_query_imgs': query_info[2], 'num_gallery_pids': gallery_info[1], 'num_gallery_imgs': gallery_info[2] }] write_json(split, self.split_new_det_json_path) print('Creating new split for labeled images (767/700) ...') train_info, query_info, gallery_info = _extract_new_split( loadmat(self.split_new_lab_mat_path), self.imgs_labeled_dir) split = [{ 'train': train_info[0], 'query': query_info[0], 'gallery': gallery_info[0], 'num_train_pids': train_info[1], 'num_train_imgs': train_info[2], 'num_query_pids': query_info[1], 'num_query_imgs': query_info[2], 'num_gallery_pids': gallery_info[1], 'num_gallery_imgs': gallery_info[2] }] write_json(split, self.split_new_lab_json_path)
def visactmap(self, testloader, save_dir, width, height, print_freq): """Visualizes CNN activation maps to see where the CNN focuses on to extract features. This function takes as input the query images of target datasets Reference: - Zagoruyko and Komodakis. Paying more attention to attention: Improving the performance of convolutional neural networks via attention transfer. ICLR, 2017 - Zhou et al. Omni-Scale Feature Learning for Person Re-Identification. ICCV, 2019. """ self.model.eval() imagenet_mean = [0.485, 0.456, 0.406] imagenet_std = [0.229, 0.224, 0.225] for target in list(testloader.keys()): queryloader = testloader[target]['query'] # original images and activation maps are saved individually actmap_dir = osp.join(save_dir, 'actmap_' + target) mkdir_if_missing(actmap_dir) print('Visualizing activation maps for {} ...'.format(target)) for batch_idx, data in enumerate(queryloader): imgs, paths = data[0], data[3] if self.use_gpu: imgs = imgs.cuda() # forward to get convolutional feature maps try: outputs = self.model(imgs, return_featuremaps=True) except TypeError: raise TypeError('forward() got unexpected keyword argument "return_featuremaps". ' \ 'Please add return_featuremaps as an input argument to forward(). When ' \ 'return_featuremaps=True, return feature maps only.') if outputs.dim() != 4: raise ValueError('The model output is supposed to have ' \ 'shape of (b, c, h, w), i.e. 4 dimensions, but got {} dimensions. ' 'Please make sure you set the model output at eval mode ' 'to be the last convolutional feature maps'.format(outputs.dim())) # compute activation maps outputs = (outputs**2).sum(1) b, h, w = outputs.size() outputs = outputs.view(b, h * w) outputs = F.normalize(outputs, p=2, dim=1) outputs = outputs.view(b, h, w) if self.use_gpu: imgs, outputs = imgs.cpu(), outputs.cpu() for j in range(outputs.size(0)): # get image name path = paths[j] imname = osp.basename(osp.splitext(path)[0]) # RGB image img = imgs[j, ...] for t, m, s in zip(img, imagenet_mean, imagenet_std): t.mul_(s).add_(m).clamp_(0, 1) img_np = np.uint8(np.floor(img.numpy() * 255)) img_np = img_np.transpose( (1, 2, 0)) # (c, h, w) -> (h, w, c) # activation map am = outputs[j, ...].numpy() am = cv2.resize(am, (width, height)) am = 255 * (am - np.max(am)) / (np.max(am) - np.min(am) + 1e-12) am = np.uint8(np.floor(am)) am = cv2.applyColorMap(am, cv2.COLORMAP_JET) # overlapped overlapped = img_np * 0.3 + am * 0.7 overlapped[overlapped > 255] = 255 overlapped = overlapped.astype(np.uint8) # save images in a single figure (add white spacing between images) # from left to right: original image, activation map, overlapped image grid_img = 255 * np.ones( (height, 3 * width + 2 * GRID_SPACING, 3), dtype=np.uint8) grid_img[:, :width, :] = img_np[:, :, ::-1] grid_img[:, width + GRID_SPACING:2 * width + GRID_SPACING, :] = am grid_img[:, 2 * width + 2 * GRID_SPACING:, :] = overlapped cv2.imwrite(osp.join(actmap_dir, imname + '.jpg'), grid_img) if (batch_idx + 1) % print_freq == 0: print('- done batch {}/{}'.format(batch_idx + 1, len(queryloader)))
def visactmap(model, test_loader, save_dir, width, height, use_gpu, img_mean=None, img_std=None): if img_mean is None or img_std is None: # use imagenet mean and std img_mean = IMAGENET_MEAN img_std = IMAGENET_STD model.eval() for target in list(test_loader.keys()): data_loader = test_loader[target]['query'] # only process query images # original images and activation maps are saved individually actmap_dir = osp.join(save_dir, 'actmap_' + target) mkdir_if_missing(actmap_dir) print('Visualizing activation maps for {} ...'.format(target)) for batch_idx, data in enumerate(data_loader): imgs, paths = data[0], data[3] if use_gpu: imgs = imgs.cuda() # forward to get convolutional feature maps try: outputs = model(imgs) outputs = outputs[3]['after'][0] except TypeError: raise TypeError( 'forward() got unexpected keyword argument "return_featuremaps". ' 'Please add return_featuremaps as an input argument to forward(). When ' 'return_featuremaps=True, return feature maps only.') if outputs.dim() != 4: raise ValueError( 'The model output is supposed to have ' 'shape of (b, c, h, w), i.e. 4 dimensions, but got {} dimensions. ' 'Please make sure you set the model output at eval mode ' 'to be the last convolutional feature maps'.format( outputs.dim())) # compute activation maps outputs = (outputs**2).sum(1) b, h, w = outputs.size() outputs = outputs.view(b, h * w) outputs = F.normalize(outputs, p=2, dim=1) outputs = outputs.view(b, h, w) if use_gpu: imgs, outputs = imgs.cpu(), outputs.cpu() for j in range(outputs.size(0)): # get image name p = paths[j] imname = p.split('/') imname = imname[-1] imdir = p.split('/') imdir = imdir[-2] # RGB image img = imgs[j, ...] for t, m, s in zip(img, img_mean, img_std): t.mul_(s).add_(m).clamp_(0, 1) img_np = np.uint8(np.floor(img.numpy() * 255)) img_np = img_np.transpose((1, 2, 0)) # (c, h, w) -> (h, w, c) # activation map am = outputs[j, ...].numpy() am = cv2.resize(am, (width, height)) am = 255 * (am - np.min(am)) / (np.max(am) - np.min(am) + 1e-12) am = np.uint8(np.floor(am)) am = cv2.applyColorMap(am, cv2.COLORMAP_JET) # overlapped overlapped = img_np * 0.3 + am * 0.7 overlapped[overlapped > 255] = 255 overlapped = overlapped.astype(np.uint8) # save images in a single figure (add white spacing between images) # from left to right: original image, activation map, overlapped image grid_img = 255 * np.ones( (height, 3 * width + 2 * GRID_SPACING, 3), dtype=np.uint8) grid_img[:, :width, :] = img_np[:, :, ::-1] grid_img[:, width + GRID_SPACING:2 * width + GRID_SPACING, :] = am grid_img[:, 2 * width + 2 * GRID_SPACING:, :] = overlapped if not osp.exists(osp.join(actmap_dir, imdir)): mkdir_if_missing(osp.join(actmap_dir, imdir)) cv2.imwrite(osp.join(actmap_dir, imdir, imname), grid_img) if (batch_idx + 1) % 10 == 0: print('- done batch {}/{}'.format(batch_idx + 1, len(data_loader)))