def __getitem__()

in DataLoaders/MIMEandPlan_DataLoader.py [0:0]

• 42 lines of code
• 2 McCabe index (conditional complexity)

	def __getitem__(self, index, split=None, return_plan_run=None):
	# def __getitem__(self, inputs):
		'''
		# Returns Joint Angles as: 
		# List of length Number_Timesteps, with each element of the list a dictionary containing the sequence of joint angles. 
		# Assumes index is within range [0,len(filelist)-1]
		'''

		if split=="train":
			file = self.train_filelist[index]
		elif split=="val":
			file = self.val_filelist[index]
		elif split=="test":
			file = self.test_filelist[index]
		elif split is None: 
			file = self.filelist[index]
			
		left_gripper = np.loadtxt(os.path.join(os.path.split(file)[0],'left_gripper.txt'))
		right_gripper = np.loadtxt(os.path.join(os.path.split(file)[0],'right_gripper.txt'))

		orig_left_traj = np.load(osp.join(osp.split(file)[0], 'Left_EE.npy'))
		orig_right_traj = np.load(osp.join(osp.split(file)[0], 'Right_EE.npy'))        

		joint_angle_trajectory = []

		folder = "New_Plans"
		if return_plan_run is not None:
			ee_plan = np.load(os.path.join(os.path.split(file)[0],"{0}/Run{1}_EE_Plan.npy".format(folder,return_plan_run)))
			ja_plan = np.load(os.path.join(os.path.split(file)[0],"{0}/Run{1}_JA_Plan.npy".format(folder,return_plan_run)))

		# Open file. 
		with open(file, 'r') as file:
			lines = file.readlines()
			for line in lines:
				dict_element = eval(line.rstrip('\n'))
				if len(dict_element.keys()) == len(self.joint_names):
					# some files have extra lines with gripper keys e.g. MIME_jointangles/4/12405Nov19/joint_angles.txt
					array_element = np.array([dict_element[joint] for joint in self.joint_names])
					joint_angle_trajectory.append(array_element)

		joint_angle_trajectory = np.array(joint_angle_trajectory)

		n_samples = len(orig_left_traj) // self.ds_freq

		elem = {}
		elem['joint_angle_trajectory'] = resample(joint_angle_trajectory, n_samples)
		elem['left_trajectory'] = resample(orig_left_traj, n_samples)
		elem['right_trajectory'] = resample(orig_right_traj, n_samples)
		elem['left_gripper'] = resample(left_gripper, n_samples)
		elem['right_gripper'] = resample(right_gripper, n_samples)
		elem['path_prefix'] = os.path.split(self.filelist[index])[0]
		elem['JA_Plan'] = ja_plan
		elem['EE_Plan'] = ee_plan

		return elem