in envs/thor.py [0:0]
def get_target_obj(self, obj_property, overlap_thresh=0.3):
objId_to_obj = {obj['objectId']:obj for obj in self.state.metadata['objects'] if obj['visible'] and obj['objectId']!=self.inv_obj}
instance_segs = self.state.instance_segmentation_frame # (300, 300, 3)
color_to_count = Image.fromarray(instance_segs, 'RGB').getcolors()
color_to_count = dict({pix:cnt for cnt,pix in color_to_count})
color_to_objId = self.state.color_to_object_id
active_px = instance_segs[self.center[0]:self.center[1], self.center[0]:self.center[1]] # (B, B, 3)
S = active_px.shape[0]
instance_counter = collections.defaultdict(list)
for i, j in itertools.product(range(S), range(S)):
color = tuple(active_px[i, j])
if color not in color_to_objId or color_to_objId[color] not in objId_to_obj:
continue
instance_counter[color].append(np.abs(i-S//2) + np.abs(j-S//2))
instance_counter = [{'color':color, 'N':len(scores), 'objectId':color_to_objId[color], 'dist':np.mean(scores), 'p1':len(scores)/S**2, 'p2':len(scores)/color_to_count[color]} for color, scores in instance_counter.items()]
# either >K% of the object is inside the box, OR K% of the pixels belong to that object
all_targets = [inst for inst in instance_counter if inst['p1']>overlap_thresh or inst['p2']>overlap_thresh]
all_targets = sorted(all_targets, key=lambda x: x['dist'])
act_targets = [candidate for candidate in all_targets if obj_property(objId_to_obj[candidate['objectId']])]
targets = {'objectId':None, 'obj':None, 'center_objectId':None, 'center_obj':None, 'int_target':None}
if len(all_targets)>0:
objId = all_targets[0]['objectId']
targets.update({'center_objectId':objId, 'center_obj':objId_to_obj[objId], 'int_target':'center_obj'})
if len(act_targets)>0:
objId = act_targets[0]['objectId']
targets.update({'objectId':objId, 'obj':objId_to_obj[objId], 'int_target':'obj'})
if targets['int_target'] is None:
targets['int_target'] = 'dummy'
return targets