dynamic_utils.py

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