how to get the same catogary but different objects in the same batch

[dengchengxifrank]

Dear xie,thanks for your wonderful work.About the code,I wonder how to get the same catogary but different objects in the same batch.For example,the input is torch.Size([batch_size, n_views, img_c, img_h, img_w]),but I wonder as for the data from the different batch,is it from the same catogray?And how can I change the code to get the result as torch.Size([batch_size, n_category, n_views, img_c, img_h, img_w])? thanks a lot!

[hzxie]

Just modify the code below.

Pix2Vox/utils/data_loaders.py

Lines 31 to 89 in 5f75fbd

	class ShapeNetDataset(torch.utils.data.dataset.Dataset):
	"""ShapeNetDataset class used for PyTorch DataLoader"""
	def __init__(self, dataset_type, file_list, n_views_rendering, transforms=None):
	self.dataset_type = dataset_type
	self.file_list = file_list
	self.transforms = transforms
	self.n_views_rendering = n_views_rendering
	def __len__(self):
	return len(self.file_list)
	def __getitem__(self, idx):
	taxonomy_name, sample_name, rendering_images, volume = self.get_datum(idx)
	if self.transforms:
	rendering_images = self.transforms(rendering_images)
	return taxonomy_name, sample_name, rendering_images, volume
	def set_n_views_rendering(self, n_views_rendering):
	self.n_views_rendering = n_views_rendering
	def get_datum(self, idx):
	taxonomy_name = self.file_list[idx]['taxonomy_name']
	sample_name = self.file_list[idx]['sample_name']
	rendering_image_paths = self.file_list[idx]['rendering_images']
	volume_path = self.file_list[idx]['volume']
	# Get data of rendering images
	if self.dataset_type == DatasetType.TRAIN:
	selected_rendering_image_paths = [
	rendering_image_paths[i]
	for i in random.sample(range(len(rendering_image_paths)), self.n_views_rendering)
	]
	else:
	selected_rendering_image_paths = [rendering_image_paths[i] for i in range(self.n_views_rendering)]
	rendering_images = []
	for image_path in selected_rendering_image_paths:
	rendering_image = cv2.imread(image_path, cv2.IMREAD_UNCHANGED).astype(np.float32) / 255.
	if len(rendering_image.shape) < 3:
	print('[FATAL] %s It seems that there is something wrong with the image file %s' %
	(dt.now(), image_path))
	sys.exit(2)
	rendering_images.append(rendering_image)
	# Get data of volume
	_, suffix = os.path.splitext(volume_path)
	if suffix == '.mat':
	volume = scipy.io.loadmat(volume_path)
	volume = volume['Volume'].astype(np.float32)
	elif suffix == '.binvox':
	with open(volume_path, 'rb') as f:
	volume = utils.binvox_rw.read_as_3d_array(f)
	volume = volume.data.astype(np.float32)
	return taxonomy_name, sample_name, np.asarray(rendering_images), volume