Spaces:
Running
Running
| 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 | |