191 lines
6.4 KiB
Python
191 lines
6.4 KiB
Python
import random
|
|
from datetime import datetime
|
|
import os
|
|
import json
|
|
import pandas as pd
|
|
import matplotlib.pyplot as plt
|
|
from matplotlib.ticker import FuncFormatter
|
|
from matplotlib.patches import FancyBboxPatch
|
|
|
|
|
|
class YcbObjects:
|
|
def __init__(self, load_path, mod_orn=None, mod_stiffness=None, exclude=None):
|
|
self.load_path = load_path
|
|
self.mod_orn = mod_orn
|
|
self.mod_stiffness = mod_stiffness
|
|
with open(load_path + '/obj_list.txt') as f:
|
|
lines = f.readlines()
|
|
self.obj_names = [line.rstrip('\n') for line in lines]
|
|
if exclude is not None:
|
|
for obj_name in exclude:
|
|
self.obj_names.remove(obj_name)
|
|
|
|
def shuffle_objects(self):
|
|
random.shuffle(self.obj_names)
|
|
|
|
def get_obj_path(self, obj_name):
|
|
return f'{self.load_path}/Ycb{obj_name}/model.urdf'
|
|
|
|
def check_mod_orn(self, obj_name):
|
|
if self.mod_orn is not None and obj_name in self.mod_orn:
|
|
return True
|
|
return False
|
|
|
|
def check_mod_stiffness(self, obj_name):
|
|
if self.mod_stiffness is not None and obj_name in self.mod_stiffness:
|
|
return True
|
|
return False
|
|
|
|
def get_obj_info(self, obj_name):
|
|
return self.get_obj_path(obj_name), self.check_mod_orn(obj_name), self.check_mod_stiffness(obj_name)
|
|
|
|
def get_n_first_obj_info(self, n):
|
|
info = []
|
|
for obj_name in self.obj_names[:n]:
|
|
info.append(self.get_obj_info(obj_name))
|
|
return info
|
|
|
|
|
|
class PackPileData:
|
|
|
|
def __init__(self, num_obj, trials, save_path, scenario):
|
|
self.num_obj = num_obj
|
|
self.trials = trials
|
|
self.save_path = save_path
|
|
|
|
if not os.path.exists(save_path):
|
|
os.mkdir(save_path)
|
|
now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
|
|
self.save_dir = f'{save_path}/{now}_{scenario}'
|
|
os.mkdir(self.save_dir)
|
|
|
|
self.tries = 0
|
|
self.succes_grasp = 0
|
|
self.succes_target = 0
|
|
|
|
def add_try(self):
|
|
self.tries += 1
|
|
|
|
def add_succes_target(self):
|
|
self.succes_target += 1
|
|
|
|
def add_succes_grasp(self):
|
|
self.succes_grasp += 1
|
|
|
|
def summarize(self):
|
|
grasp_acc = self.succes_grasp / self.tries
|
|
target_acc = self.succes_target / self.tries
|
|
perc_obj_cleared = self.succes_target / (self.trials * self.num_obj)
|
|
|
|
with open(f'{self.save_dir}/summary.txt', 'w') as f:
|
|
f.write(
|
|
f'Stats for {self.num_obj} objects out of {self.trials} trials\n')
|
|
f.write(
|
|
f'Target acc={target_acc:.3f} ({self.succes_target}/{self.tries})\n')
|
|
f.write(
|
|
f'Grasp acc={grasp_acc:.3f} ({self.succes_grasp}/{self.tries})\n')
|
|
f.write(
|
|
f'Percentage objects cleared={perc_obj_cleared} ({self.succes_target}/{(self.trials * self.num_obj)})\n')
|
|
|
|
|
|
class IsolatedObjData:
|
|
|
|
def __init__(self, obj_names, trials, save_path):
|
|
self.obj_names = obj_names
|
|
self.trials = trials
|
|
self.succes_target = dict.fromkeys(obj_names, 0)
|
|
self.succes_grasp = dict.fromkeys(obj_names, 0)
|
|
self.tries = dict.fromkeys(obj_names, 0)
|
|
|
|
if not os.path.exists(save_path):
|
|
os.mkdir(save_path)
|
|
now = datetime.now().strftime('%Y-%m-%d %H:%M:%S')
|
|
self.save_dir = f'{save_path}/{now}_iso_obj'
|
|
os.mkdir(self.save_dir)
|
|
|
|
def add_succes_target(self, obj_name):
|
|
self.succes_target[obj_name] += 1
|
|
|
|
def add_succes_grasp(self, obj_name):
|
|
self.succes_grasp[obj_name] += 1
|
|
|
|
def add_try(self, obj_name):
|
|
self.tries[obj_name] += 1
|
|
|
|
def write_json(self):
|
|
data_tries = json.dumps(self.tries)
|
|
data_target = json.dumps(self.succes_target)
|
|
data_grasp = json.dumps(self.succes_grasp)
|
|
f = open(self.save_dir+'/data_tries.json', 'w')
|
|
f.write(data_tries)
|
|
f.close()
|
|
f = open(self.save_dir+'/data_target.json', 'w')
|
|
f.write(data_target)
|
|
f.close()
|
|
f = open(self.save_dir+'/data_grasp.json', 'w')
|
|
f.write(data_grasp)
|
|
f.close()
|
|
|
|
|
|
def plot(path, tries, target, grasp, trials):
|
|
succes_rate = dict.fromkeys(tries.keys())
|
|
for obj in succes_rate.keys():
|
|
t = tries[obj]
|
|
if t == 0:
|
|
t = 1
|
|
acc_target = target[obj] / t
|
|
acc_grasp = grasp[obj] / t
|
|
succes_rate[obj] = (acc_target, acc_grasp)
|
|
|
|
plt.rc('axes', titlesize=13) # fontsize of the axes title
|
|
plt.rc('axes', labelsize=12) # fontsize of the x and y labels
|
|
|
|
df = pd.DataFrame(succes_rate).T
|
|
df.columns = ['Target', 'Grasped']
|
|
df = df.sort_values(by='Target', ascending=True)
|
|
ax = df.plot(kind='bar', color=['#88CCEE', '#CC6677'])
|
|
plt.xlabel('Object name')
|
|
plt.ylabel('Succes rate (%)')
|
|
plt.title(
|
|
f'Succes rate of objects grasped and placed in target | {trials} runs')
|
|
plt.grid(color='#95a5a6', linestyle='-', linewidth=1, axis='y', alpha=0.5)
|
|
ax.set_axisbelow(True)
|
|
ax.yaxis.set_major_formatter(
|
|
FuncFormatter(lambda y, _: '{:.0%}'.format(y)))
|
|
plt.setp(ax.xaxis.get_majorticklabels(), rotation=45,
|
|
ha="right", rotation_mode="anchor")
|
|
plt.locator_params(axis="y", nbins=11)
|
|
plt.legend(loc='lower right')
|
|
plt.subplots_adjust(bottom=0.28)
|
|
|
|
plt.savefig(path+'/plot.png')
|
|
|
|
|
|
def write_summary(path, tries, target, grasp):
|
|
with open(path+'/summary.txt', 'w') as f:
|
|
total_tries = sum(tries.values())
|
|
total_target = sum(target.values())
|
|
total_grasp = sum(grasp.values())
|
|
f.write('Total:\n')
|
|
f.write(
|
|
f'Target acc={total_target/total_tries:.3f} ({total_target}/{total_tries}) Grasp acc={total_grasp/total_tries:.3f} ({total_grasp}/{total_tries})\n')
|
|
f.write('\n')
|
|
f.write("Accuracy per object:\n")
|
|
for obj in tries.keys():
|
|
n_tries = tries[obj]
|
|
n_t = target[obj]
|
|
n_g = grasp[obj]
|
|
f.write(
|
|
f'{obj}: Target acc={n_t/n_tries:.3f} ({n_t}/{n_tries}) Grasp acc={n_g/n_tries:.3f} ({n_g}/{n_tries})\n')
|
|
|
|
|
|
def summarize(path, trials):
|
|
with open(path+'/data_tries.json') as data:
|
|
tries = json.load(data)
|
|
with open(path+'/data_target.json') as data:
|
|
target = json.load(data)
|
|
with open(path+'/data_grasp.json') as data:
|
|
grasp = json.load(data)
|
|
plot(path, tries, target, grasp, trials)
|
|
write_summary(path, tries, target, grasp)
|