import numpy as np import random import copy PRIVATE_OBJECT_LIST = ["AlarmClock", "CellPhone", "Cloth", "CreditCard", "HandTower", "KeyChain", "Laptop", "Pillow", "TeddyBear", "Towel", "Watch"] COMMONLY_USED_OBJECT_LIST = ['Book', 'CD', 'Newspaper'] FILTER_TYPE = [ "AlarmClock", "BaseballBat", "BasketBall", "Bowl", "Book" "CellPhone", "Cloth", "CreditCard", 'Cup', 'DishSponge', 'Fork', 'KeyChain', "Laptop", 'Knife', "Mug", "Newspaper", 'Pen', 'Pencil', 'Pillow', 'Plate', 'Spoon', 'SprayBottle', 'Vase', 'Watch', ] PICKUPABLE_OBJECT_LIST = ['AlarmClock', 'AluminumFoil', 'Apple', 'AppleSliced', 'BaseballBat', 'BasketBall', 'Book', 'Boots', 'Bottle', 'Bowl', 'Box', 'Bread', 'BreadSliced', 'ButterKnife', 'Candle', 'CD', 'CellPhone', 'Cloth', 'CreditCard', 'Cup', 'DishSponge', 'Dumbbell', 'Egg', 'EggCracked', 'Fork', 'HandTowel', 'Kettle', 'KeyChain', 'Knife', 'Ladle', 'Laptop', 'Lettuce', 'LettuceSliced', 'Mug', 'Newspaper', 'Pan', 'PaperTowelRoll', 'Pen', 'Pencil', 'PepperShaker', 'Pillow', 'Plate', 'Plunger', 'Pot', 'Potato', 'PotatoSliced', 'RemoteControl', 'SaltShaker', 'ScrubBrush', 'SoapBar', 'SoapBottle', 'Spatula', 'Spoon', 'SprayBottle', 'Statue', 'TableTopDecor', 'TeddyBear', 'TennisRacket', 'TissueBox', 'ToiletPaper', 'Tomato', 'TomatoSliced', 'Towel', 'Vase', 'Watch', 'WateringCan', 'WineBottle'] RECEPTACLE_LIST = ['ArmChair', 'Bathtub', 'BathtubBasin', 'Bed', 'Bowl', 'Box', 'Cabinet', 'Chair', 'CoffeeMachine', 'CoffeeTable', 'CounterTop', 'Cup', 'Desk', 'DiningTable', 'Drawer', 'Dresser', 'Fridge', 'GarbageCan', 'HandTowelHolder', 'LaundryHamper', 'Microwave', 'Mug', 'Ottoman', 'Pan', 'Plate', 'Pot', 'Safe', 'Shelf', 'SideTable', 'Sink', 'SinkBasin', 'Sofa', 'StoveBurner', 'Toaster', 'Toilet', 'ToiletPaperHanger', 'TowelHolder', 'TVStand'] def get_obj_type(obj_id): return obj_id.split('|')[0] def get_obj_belong_type(obj_id): obj_type = get_obj_type(obj_id) belong_type = "public" if obj_type in PRIVATE_OBJECT_LIST: belong_type = "private" elif obj_type in COMMONLY_USED_OBJECT_LIST: belong_type = "part_public" return belong_type class Sampler(): def __init__(self): self.schedules = [] self.activity_base = {} self.character = [] self.character_specific = {} self.scene_layout = {} self.object_prior = {} def clear_character_specific(self): self.character_specific = {} for c in self.character: self.character_specific[c] = {} self.character_specific[c]['room'] = {} self.character_specific[c]['related_object'] = [] def set_activity(self, character_info, activity_base, schedules): self.activity_base = activity_base self.character_info = character_info self.character = [c['name'] for c in self.character_info] self.schedules = schedules self.clear_character_specific() def set_scene_layout(self, scene_layout): self.scene_layout = scene_layout self.clear_character_specific() def get_character_group(self): result = [[c] for c in self.character] return result def object_allocation(self): rooms = self.scene_layout['room'] bedrooms = [room for room in rooms if room['roomType'] == 'Bedroom'] bathrooms = [room for room in rooms if room['roomType'] == 'Bathroom'] character_group = self.get_character_group() bedroom_stack = copy.deepcopy(bedrooms) bathroom_stack = copy.deepcopy(bathrooms) for cg in character_group: if len(bedroom_stack) > 0: cd_bedroom = bedroom_stack[-1] bedroom_stack.pop(-1) else: cd_bedroom = bedrooms[0] if len(bathroom_stack) > 0: cd_bathroom = bathroom_stack[-1] bathroom_stack.pop(-1) else: cd_bathroom = bathrooms[0] for c in cg: self.character_specific[c]['room'] = { 'Bedroom':cd_bedroom['roomId'], 'Bathroom':cd_bathroom['roomId'] } #object pickup_objects = self.scene_layout['pickupable object'] for pobj in pickup_objects: pobj_type = get_obj_type(pobj) #private if pobj_type in PRIVATE_OBJECT_LIST: for character in self.character_specific: private_objects = self.character_specific[character]['related_object'] exist = False for o in private_objects: if get_obj_type(o) == pobj_type: exist = True break if not exist: self.character_specific[character]['related_object'].append(pobj) break #common used (not all characters) elif pobj_type in COMMONLY_USED_OBJECT_LIST: if len(self.character) >= 2: selected_characters = random.sample(self.character, 2) else: selected_characters = self.character for sc in selected_characters: self.character_specific[sc]['related_object'].append(pobj) def sample_schedules(self, id=0): character = self.character[id] schedules = self.schedules[id] activity_base = self.activity_base[id] results = {} for day in schedules: day_schedule = [] for activity in schedules[day]: activity_name = activity['activity'] #character = activity['character'] if activity_name not in activity_base: continue activity_content = activity_base[activity_name] #sample room rooms = list(activity_content.keys()) room_probs = np.array([activity_content[room]['room_prob'] for room in activity_content]) if not room_probs.sum() == 1: room_probs = room_probs / room_probs.sum() choose_ids = [i for i in range(len(rooms))] choose_id = np.random.choice(choose_ids, p=room_probs.ravel()) room = rooms[choose_id] #room ground to scene final_rooms = [] #卧室和厕所找分配 if (room == 'Bedroom' or room == 'Bathroom'): final_room_id = self.character_specific[character]['room'][room] for r in self.scene_layout['room']: if r['roomId'] == final_room_id: final_room = r break else: for scene_room in self.scene_layout['room']: if scene_room['roomType'] == room: final_rooms.append(scene_room) #final_room = final_rooms[0] final_room = np.random.choice(final_rooms) #sample object object_effect = activity_content[room]['object_effect'] sample_objects = [] for obj in object_effect: object_prob = float(object_effect[obj]['object_prob']) random_num = random.random() if random_num < object_prob: sample_objects.append(obj) #sample receptacle sample_receptacles = [] for obj in sample_objects: #除去房间中没有的receptacle，并重新归一化 possible_receptacles_init = object_effect[obj]['receptacles'] possible_receptacles = [] for pr in possible_receptacles_init: for child in final_room['children']: if pr[0] == child.split('|')[0] and pr[0]!='Fridge': possible_receptacles.append(pr) break if len(possible_receptacles) == 0: continue prob_sum = 0 for pr in possible_receptacles: prob_sum += pr[1] for i in range(len(possible_receptacles)): possible_receptacles[i][1] = possible_receptacles[i][1] / prob_sum receptacle_probs = np.array([recep[1] for recep in possible_receptacles]) choose_ids = [i for i in range(len(receptacle_probs))] choose_id = np.random.choice(choose_ids, p=receptacle_probs.ravel()) receptacle = possible_receptacles[choose_id][0] sample_receptacles.append([obj, receptacle]) final_sample = [] for type_sample in sample_receptacles: obj_type = type_sample[0] receptacle_type = type_sample[1] final_obj = "" final_recep = "" related_objs = self.character_specific[character]['related_object'] if obj_type in PRIVATE_OBJECT_LIST: for ro in related_objs: if get_obj_type(ro) == obj_type: final_obj = ro break else: exist = False for ro in related_objs: if get_obj_type(ro) == obj_type: exist = True final_obj = ro break if not exist: #从全部找 cand_final_obj = [] for pick_obj in self.scene_layout['pickupable object']: if pick_obj.split("|")[0] == obj_type: #final_obj = pick_obj cand_final_obj.append(pick_obj) if len(cand_final_obj) > 0: final_obj = np.random.choice(cand_final_obj) for recep in final_room['children']: if recep.split('|')[0] == receptacle_type: final_recep = recep break if len(final_obj)>0 and len(final_recep) > 0: final_sample.append({ 'object':final_obj, 'receptacle':final_recep }) day_schedule.append({ 'activity':activity_name, 'character':character, 'start_time':activity['start_time'], 'end_time':activity['end_time'], 'room':final_room['roomId'], 'room_type':final_room['roomType'], 'content':final_sample, }) results[day] = day_schedule return results def sample(self): self.object_allocation() sample_schedules = {} keys = list(self.schedules[0].keys()) for key in keys: sample_schedules[key] = [] for i in range(len(self.character)): schedule = self.sample_schedules(i) for key in keys: sample_schedules[key].extend(schedule[key]) def compare_time(item): return item['start_time'] for key in keys: sample_schedules[key] = sorted(sample_schedules[key], key=compare_time) return sample_schedules, None, None def get_layout(house_info): objects = house_info['objects'] rooms = house_info['rooms'] room_info = [] receptacle_info = [] pickupable_object_info = [] for room in rooms: room_info.append({ 'roomId':room['id'], 'roomType':room['roomType'], 'roomPolygon':room['floorPolygon'], 'children':[], }) def isIn(obj, room): pos = obj['position'] room_polygon = room['roomPolygon'] polygon_n = len(room_polygon) j = polygon_n - 1 res = False for i in range(polygon_n): if ((room_polygon[i]['z']>pos['z'])!=(room_polygon[j]['z']>pos['z']) and \ pos['x']< (room_polygon[j]['x'] - room_polygon[i]['x']) * (pos['z'] - room_polygon[i]['z'])/(room_polygon[j]['z'] - room_polygon[i]['z']) + room_polygon[i]['x'] ): res = not res j = i return res def get_type(oid): return oid.split('|')[0] def get_receptacle_info(obj): rinfo = { 'id':obj['id'], } if 'children' in obj: rc = [] for child in obj['children']: rc.append(child['id']) if get_type(child['id']) in PICKUPABLE_OBJECT_LIST: pickupable_object_info.append(child['id']) if get_type(child['id']) in RECEPTACLE_LIST: get_receptacle_info(child) if len(rc) >0: rinfo['children'] = rc receptacle_info.append(rinfo) for obj in objects: for room in room_info: if isIn(obj, room): room['children'].append(obj['id']) #print("{} in {}".format(obj['id'], room['roomType'])) break if get_type(obj['id']) in RECEPTACLE_LIST: get_receptacle_info(obj) if get_type(obj['id']) in PICKUPABLE_OBJECT_LIST: pickupable_object_info.append(obj['id']) for r in room_info: del r['roomPolygon'] layout_info = { 'room':room_info, 'receptacle': receptacle_info, 'pickupable object' : pickupable_object_info } return layout_info