alters-mit / sticky_mitten_avatar Goto Github PK

from pathlib import Path
from typing import List, Union
from tdw.tdw_utils import TDWUtils
from tdw.output_data import Images, CameraMatrices
#from sticky_mitten_avatar.avatars import Arm
from sticky_mitten_avatar import StickyMittenAvatarController
from sticky_mitten_avatar.util import get_data


import numpy as np

class PutObjectInContainer(StickyMittenAvatarController):
    """
    1. Create a sticky mitten avatar, a jug, and a container. Add an overhead camera for image capture.
    2. The avatar picks up the jug.
    3. The avatar goes to the container.
    4. The avatar puts the jug in the container.

    Save an image per frame.
    """

    def __init__(self, output_dir, port: int = 1071, launch_build: bool = True):
        """
        :param port: The port number.
        :param launch_build: If True, automatically launch the build.
        """

        super().__init__(port=port, launch_build=launch_build, audio=False, \
                id_pass=False, demo=False, grayscale_depth=False)
        self.output_dir = Path(output_dir)
        if not self.output_dir.exists():
            self.output_dir.mkdir(parents=True)
        self.output_dir = str(self.output_dir.resolve())
        # Save images every frame, if possible.
        self.frame_count = 0
        self.o_id = 0
        self.bowl_id = 1

    def _get_scene_init_commands(self, scene: str = None, layout: int = None) -> List[dict]:
        commands = super()._get_scene_init_commands()
        '''# Add a jug.
        self.o_id, jug_commands = self._add_object("jug05",
                                                   position={"x": -0.2, "y": 0, "z": 0.285},
                                                   scale={"x": 0.8, "y": 0.8, "z": 0.8})
        commands.extend(jug_commands)
        # Add a container.
        bowl_position = {"x": 1.2, "y": 0, "z": 0.25}
        self.bowl_id, bowl_commands = self._add_object("serving_bowl",
                                                       position=bowl_position,
                                                       rotation={"x": 0, "y": 30, "z": 0},
                                                       scale={"x": 1.3, "y": 1, "z": 1.3})'''
        
        table_position = {"x": 2.4, "y": 0, "z": 4.25}                                        
        self.table_id, table_commands = self._add_object("small_table_green_marble",
                                                        position = table_position,
                                                        rotation={"x": 0, "y": 0, "z": 0},
                                                        scale={"x": 1.3, "y": 1, "z": 1.3})
        commands.extend(table_commands)
        return commands

        
        
    def save_infomation(self, output_directory='.'):
        frame_count = self.frame._frame_count
        self.frame.save_images(output_directory=output_directory)
        np.save(f'{output_directory}/depth_{frame_count}.npy', self.frame.get_depth_values())
        np.save(f'{output_directory}/{frame_count}.npy', self.frame.camera_matrix)
        
        '''import pickle
        info = {}
        info['pos'] = self.frame.'''
    
    def run(self) -> None:
        """
        Run a single trial. Save images per frame.
        """
        import time
        start = time.time()
        self.init_scene()
        #print(self.occupancy_map.shape)
        #np.save('2amap.npy', self.occupancy_map)
        output_directory = 'new'
        self.save_infomation(output_directory)
        
        #self.move_forward_by(distance=)
        
        
        self.turn_to(target=self.table_id)
        self.save_infomation(output_directory)
        self.move_forward_by(distance=-1.5)
        self.save_infomation(output_directory)
        self.turn_to(target=self.table_id)
        self.save_infomation(output_directory)
        

        self.end()


if __name__ == "__main__":
    port = 10073
    remote = False
    if remote:
        from utils import create_tdw, kill_tdw
        docker_id = create_tdw(port=port)    
        try:
            PutObjectInContainer(port=port, output_dir = 'images', launch_build=False).run()
        finally:
            kill_tdw(docker_id)
    else:
        PutObjectInContainer(port=port, output_dir = 'images', launch_build=False).run()

import numpy as np
import cv2
import math
​
​
id = '0127'
depth = np.load('./new/depth_127.npy')
​
​
​
HFOV = VFOV = FOV = 54.43222897365458
​
​
fx = (depth.shape[0] / 2) / np.tan(math.radians(FOV / 2.))
fy = (depth.shape[1] / 2) / np.tan(math.radians(FOV / 2.))
print('fx, fy:', fx, fy)
#fx = 248.8900803871498
#fy = 248.8900803871498
W, H = depth.shape
cx = W / 2.
cy = H / 2.
p = np.zeros((H, W, 3))
​
x_index = np.linspace(0, W - 1, W)
y_index = np.linspace(0, H - 1, H)
xx, yy = np.meshgrid(x_index, y_index)
xx = (xx - cx) / fx * depth
yy = (yy - cy) / fy * depth
​
mH = 600
mW = 800
occ = np.zeros((mH, mW))
map = np.zeros((mH, mW, 3))
map.fill(255)
with open(f'{id}.txt', 'w') as f:
    for i in range(xx.shape[0]):
        for j in range(xx.shape[1]):
            if depth[i, j] < 99:
                f.write(f'{xx[i, j]} {yy[i, j]} {depth[i, j]}\n')

from typing import List
from sticky_mitten_avatar.avatars import Arm
from sticky_mitten_avatar import StickyMittenAvatarController
import time

class PutObjectInContainer(StickyMittenAvatarController):
    """
    1. Create a sticky mitten avatar, a jug, and a container. Add an overhead camera for image capture.
    2. The avatar picks up the jug.
    3. The avatar goes to the container.
    4. The avatar puts the jug in the container.

    Save an image per frame.
    """

    def __init__(self, port: int = 1071, launch_build: bool = True):
        """
        :param port: The port number.
        :param launch_build: If True, automatically launch the build.
        """

        super().__init__(port=port, launch_build=launch_build, \
                audio=False, id_pass=False, demo=True)

        # Save images every frame, if possible.
        self.o_id = 0
        self.bowl_id = 1

    def _get_scene_init_commands(self, scene: str = None, layout: int = None) -> List[dict]:
        commands = super()._get_scene_init_commands()
        # Add a jug.
        self.o_id, jug_commands = self._add_object("jug05",
                                                   position={"x": -0.2, "y": 0, "z": 0.285},
                                                   scale={"x": 0.8, "y": 0.8, "z": 0.8})
        commands.extend(jug_commands)
        # Add a container.
        bowl_position = {"x": 1.2, "y": 0, "z": 0.25}
        self.bowl_id, bowl_commands = self._add_object("basket_18inx18inx12iin_wood_mesh",
                                                       position=bowl_position,
                                                       rotation={"x": 0, "y": 30, "z": 0},
                                                       scale={"x": 1.3, "y": 1, "z": 1.3})
        print('bowl:', self.bowl_id)
        commands.extend(bowl_commands)
        return commands
        
    def save_infomation(self, output_directory='.'):
        frame_count = self.frame._frame_count
        print('frame:', frame_count)
        #print('position:', self.frame.avatar_transform.position)        
        self.frame.save_images(output_directory=output_directory)
        #for id in self.frame.avatar_body_part_transforms:
        #    print(id)
        '''np.save(f'{output_directory}/{frame_count}.npy', self.frame.camera_matrix)
        
        import pickle
        info = {}
        info['pos'] = self.frame.avatar_transform.position
        info['rotation'] = self.frame.avatar_transform.rotation'''
        
    def run(self) -> None:
        """
        Run a single trial. Save images per frame.
        """
        start = time.time()
        self.init_scene()
        print('init_scene time: ', time.time() - start)
        output_directory = 'phy3'
        # Add a third-person camera.
        #self.add_overhead_camera({"x": 3.08, "y": 1.25, "z": 3.41}, target_object="a", images="cam")
        self.save_infomation(output_directory)
        
        # Pick up the object.
        start = time.time()
        '''print('grasp:', self.grasp_object(object_id=self.bowl_id, arm=Arm.left))
        print('grasp time:', time.time() - start)
        self.save_infomation(output_directory)
        # Lift the object up a bit.
        #self.move_forward_by(distance=0.3, move_stopping_threshold=0.1)
        print('move arm:', self.reach_for_target(
                target={"x": -0.2, "y": 0.8, "z": 0.42}, \
                arm=Arm.left, \
                stop_on_mitten_collision=True
                ))
        self.save_infomation(output_directory)'''
        
        
        # Go to the bowl.
        print('turn to:', self.turn_to(target=self.bowl_id, force=100))
        self.save_infomation(output_directory)
        print('go to: ', self.go_to(target=self.bowl_id, move_stopping_threshold=0.2))
        self.save_infomation(output_directory)
        #print(self.turn_to(target=self.bowl_id, force=100))
        #self.save_infomation(output_directory)
        # Lift the object up a bit.
        #print('move arm:', self.reach_for_target(target={"x": -0.1, "y": 0.8, "z": 0.2}, arm=Arm.left))
        #self.save_infomation(output_directory)
        # Drop the object in the container.
        start = time.time()
        #print('drop:', self.drop(arm=Arm.left))
        #print('drop time:', time.time() - start)
        #self.save_infomation(output_directory)
        
        start = time.time()
        print('grasp:', self.grasp_object(object_id=self.bowl_id, arm=Arm.left))
        print('hold:', self.frame.held_objects[Arm.left])
        print('grasp time:', time.time() - start)
        self.save_infomation(output_directory)
        #print('move arm:', self.reach_for_target(target={"x": -0.2, "y": 0.3, "z": 0.2}, arm=Arm.left))
        #self.save_infomation(output_directory)        
        
        status = self.turn_by(angle=-15, force=500)
        self.save_infomation(output_directory)
        print('hold:', self.frame.held_objects[Arm.left])
        print(self.move_forward_by(distance=0.5, move_stopping_threshold=0.1))
        self.save_infomation(output_directory)
        print('hold:', self.frame.held_objects[Arm.left])
        status = self.turn_by(angle=-15, force=500)
        self.save_infomation(output_directory)
        print('hold:', self.frame.held_objects[Arm.left])
        print(self.move_forward_by(distance=0.5, move_stopping_threshold=0.1))
        self.save_infomation(output_directory)
        print('hold:', self.frame.held_objects[Arm.left])
        # Stop the build.
        self.end()


if __name__ == "__main__":
    port = 10075
    PutObjectInContainer(port=port, launch_build=False).run()

from pathlib import Path
from typing import Dict, List, Union, Optional, Tuple
from tdw.tdw_utils import TDWUtils
from tdw.output_data import Images, CameraMatrices
from sticky_mitten_avatar.avatars import Arm
from sticky_mitten_avatar import StickyMittenAvatarController
from sticky_mitten_avatar.util import get_data, get_angle
from queue import Queue
import numpy as np
import pyastar
import time
import cv2
import math

DISTANCE = 0.5
ANGLE = 15.
SCALE = 4
#ANGLE2 = 45.

class Nav(StickyMittenAvatarController):
    """
    1. map
    2. global goal
    3. 
    """

    def __init__(self, output_dir, port: int = 1071, launch_build: bool = True):
        """
        :param port: The port number.
        :param launch_build: If True, automatically launch the build.
        """

        super().__init__(port=port, launch_build=launch_build, audio=False, \
                id_pass=True, demo=False)
        self.output_dir = Path(output_dir)
        if not self.output_dir.exists():
            self.output_dir.mkdir(parents=True)
        self.output_dir = str(self.output_dir.resolve())    
        
        
    def save_infomation(self, output_directory='.'):
        frame_count = self.frame._frame_count
        print('frame:', frame_count)
        self.frame.save_images(output_directory=output_directory)
        np.save(f'{output_directory}/{frame_count}.npy', self.frame.camera_matrix)
        
        
        '''tmp_map = np.zeros((self.map_shape[0], self.map_shape[1], 3))
        tmp_map.fill(255)
        tmp_map[self.map > 0.7] = [100, 100, 100]
        cv2.imwrite(f'./map/map_{frame_count}.png', tmp_map)'''
        
        import pickle
        info = {}
        info['pos'] = self.frame.avatar_transform.position
        info['rotation'] = self.frame.avatar_transform.rotation
    
    def check_occupied(self, x: float, z: float):
        """
        Check (x, z) if occupied or free or out-side
        0: occupied; 1: free; 2: out_side
        """
        
        if self.occupancy_map is None or self._scene_bounds is None:
            assert False
            return False, 0, 0, 0
        i = int(round((x - self._scene_bounds["x_min"]) * SCALE))
        j = int(round((z - self._scene_bounds["z_min"]) * SCALE))
        i = min(i, self.occupancy_map.shape[0] - 1)
        i = max(i, 0)
        j = min(j, self.occupancy_map.shape[1] - 1)
        j = max(j, 0)
        try:
            t = self.occupancy_map[i][j]
        except:
            print('error:', i, j)
        return True, i, j, self.occupancy_map[i][j]
    
    def get_occupancy_position(self, i, j):
        """
        Converts the position (i, j) in the occupancy map to (x, z) coordinates.

        :param i: The i coordinate in the occupancy map.
        :param j: The j coordinate in the occupancy map.
        :return: Tuple: True if the position is in the occupancy map; x coordinate; z coordinate.
        """

        if self.occupancy_map is None or self._scene_bounds is None:
            return False, 0, 0,
        x = self._scene_bounds["x_min"] + (i / SCALE)
        z = self._scene_bounds["z_min"] + (j / SCALE)
        return True, x, z
    
    def generate_goal(self):
        while True:
            x, z = np.random.random_sample(), np.random.random_sample()
            x = (self._scene_bounds["x_max"] - self._scene_bounds["x_min"]) * \
                                                        x + self._scene_bounds["x_min"]
            z = (self._scene_bounds["z_max"] - self._scene_bounds["z_min"]) * \
                                                        z + self._scene_bounds["z_min"]            
            rep = self.check_occupied(x, z)
            sx, _, sz = self.frame.avatar_transform.position
            if rep[3] == 1 and self.l2_distance((sx, sz), (x, z)) > 7:
                return x, z
        return None
        
    def find_shortest_path(self, st, goal, map = None):
        if map is None:
            map = self.map
        #map: 0-free, 1-occupied
        st_x, _, st_z = st
        g_x, g_z = goal
        _, st_i, st_j, t = self.check_occupied(st_x, st_z)
        #assert t == 1
        _, g_i, g_j, t = self.check_occupied(g_x, g_z)
        #assert t == 1
        dist_map = np.ones_like(map, dtype=np.float32)
        dist_map[map == 1] = 50000
        
        #print('min dist:', dist_map.min())
        #print('max dist:', dist_map.max())
        #dist_map
        path = pyastar.astar_path(dist_map, (st_i, st_j),
            (g_i, g_j), allow_diagonal=False)
        return path
    
    def l2_distance(self, st, g):
        return ((st[0] - g[0]) ** 2 + (st[1] - g[1]) ** 2) ** 0.5
    
    def check_goal(self, thresold = 0.5):
        x, _, z = self.frame.avatar_transform.position
        gx, gz = self.goal
        d = self.l2_distance((x, z), (gx, gz))
        return d < thresold
        
    def draw(self, traj, value):
        if not isinstance(value, np.ndarray):
            value = np.array(value)
        for dx in range(1):
            for dz in range(1):
                x = traj[0] + dx
                z = traj[1] + dz
                self.paint[x, z] = value
    
    def dep2map(self):  
        #start = time.time()
        pre_map = np.zeros_like(self.map)   
        try:
            depth = self.frame.get_depth_values()
        except:
            return pre_map
        #camera info
        FOV = 54.43222897365458
        W, H = depth.shape
        cx = W / 2.
        cy = H / 2.
        fx = cx / np.tan(math.radians(FOV / 2.))
        fy = cy / np.tan(math.radians(FOV / 2.))
        
        #Ego
        x_index = np.linspace(0, W - 1, W)
        y_index = np.linspace(0, H - 1, H)
        xx, yy = np.meshgrid(x_index, y_index)
        xx = (xx - cx) / fx * depth
        yy = (yy - cy) / fy * depth
        
        pc = np.stack((xx, yy, depth, np.ones((xx.shape[0], xx.shape[1]))))  #3, 256, 256
        pc = pc.reshape(4, -1)
        
        E = self.frame.camera_matrix
        inv_E = np.linalg.inv(np.array(E).reshape((4, 4)))
        rot = np.array([[1, 0, 0, 0],
                        [0, -1, 0, 0],
                        [0, 0, -1, 0],
                        [0, 0, 0, 1]])
        inv_E = np.dot(inv_E, rot)
        
        rpc = np.dot(inv_E, pc).reshape(4, W, H)
        #pre_map = np.zeros_like(self.map)
        #print('dep2pc time: ', time.time() - start)
        rpc = rpc.reshape(4, -1)
        X = np.rint((rpc[0, :] - self._scene_bounds["x_min"]) * SCALE)
        X = np.maximum(X, 0)
        X = np.minimum(X, self.map_shape[0] - 1)
        Z = np.rint((rpc[2, :] - self._scene_bounds["z_min"]) * SCALE)
        Z = np.maximum(Z, 0)
        Z = np.minimum(Z, self.map_shape[1] - 1)
        depth = depth.reshape(-1)
        index = np.where((depth < 99) & (rpc[1, :] > 0.2) & (rpc[1, :] < 1.4))
        X = X[index]
        Z = Z[index]
        #print(X.dtype, X.shape)
        X = X.astype(np.int32)
        Z = Z.astype(np.int32)
        pre_map[X, Z] = 1
        '''        
        for i in range(W):
            for j in range(H):
                x, y, z = rpc[0:3, i, j]
                if depth[i, j] < 99 and y > 0.2 and y < 2.8:                    
                    _, id_i, id_j, _ = self.check_occupied(x, z)
                    pre_map[id_i, id_j] = 1
        '''
        #print('dep2pc time: ', time.time() - start)
        return pre_map
    
    
    def _get_scene_init_commands(self, scene = None, layout=None): 
        commands = super()._get_scene_init_commands(scene=scene, layout=layout)
        self.object_id = None
        if self.add_object:
            name = "round_bowl_small_beech"        
            position = {"x": -7.830219268798828, "y": 0, "z": -2.2613542079925537}
            #-7.830219268798828 0.0023899078369140625 -2.7613542079925537
            rotation={"x": 0, "y": 30, "z": 0}
            scale={"x": 1, "y": 1, "z": 1}
            self.object_id, add_commands = self._add_object(name,
                                                        position=position,
                                                        rotation=rotation,
                                                        scale=scale)
            print('object id:', self.object_id)
            commands.extend(add_commands)
        return commands
        
    
    
    def run(self, scene='5a', layout=2) -> None:
        """
        Run a single trial. Save images per frame.
        """
        
        output_directory = 'nav_test1'
        
        start = time.time()
        self.add_object = True
        self.init_scene(scene=scene, layout=layout)
        '''self.init_scene()
        self.occupancy_map = np.ones((int(107 * SCALE / 4), int(107 * SCALE / 4)))
        self._scene_bounds = {        
            "x_max": 13.107263565063477,
            "x_min": -13.584211349487305,
            "z_max": 13.076096534729004,
            "z_min": -3.587161064147949
        }'''
        print('init scene time:', time.time() - start)
                
        print('go to: ', self.go_to(target=self.object_id, move_stopping_threshold=0.2))
        self.save_infomation(output_directory)
        print('turn to: ', self.turn_to(target=self.object_id, force=100))
        self.save_infomation(output_directory)
        start = time.time()
        print('grasp:', self.grasp_object(object_id=self.object_id, arm=Arm.left))
        print('grasp time:', time.time() - start)
        self.save_infomation(output_directory)
        print(self.frame.held_objects[Arm.left])
        #return None
        #0: free, 1: occupied
        self.gt_map = np.zeros_like(self.occupancy_map)
        self.gt_map[self.occupancy_map == 0] = 1
        
        self.map = np.zeros_like(self.occupancy_map)
        
        
        #0: occupied; 1: free; 2: out-side
        
        #paint map
        print('occupancy_map dtype:', self.occupancy_map.dtype)
        print('occupancy_map shape:', self.occupancy_map.shape)
        W, H = self.occupancy_map.shape
        self.map_shape = self.occupancy_map.shape        
        self.paint = np.zeros((W, H, 3))
        self.paint.fill(100)
        self.paint[self.occupancy_map == 1, 0:3] = 255
        
        x, y, z = self.frame.avatar_transform.position
        self.position = self.frame.avatar_transform.position
        print('init position:', x, y, z)
        print('init rotation:', self.frame.avatar_transform.rotation)
        print('type:', type(self.occupancy_map))
        assert self.check_occupied(x, z)[3] == 1
        W, H = self.occupancy_map.shape
        
        self.step = 0
        action = 0
        self.goal = self.generate_goal()
        #self.goal = (2.250844070014411, 3.9478932125618513)
        #self.goal = (-3.3103403049266413, -2.2811090762226776)
        
        print('goal position:', self.goal)
        print('goal:', self.check_occupied(self.goal[0], self.goal[1]))
        x, y, z = self.frame.avatar_transform.position
        print('distance: ', self.l2_distance((x, z), self.goal))
        f = open('nav.log', 'w')
        status = ''
        traj = []
        path = self.find_shortest_path(self.position, self.goal, self.gt_map)
        Astar_map = self.paint.copy()
        k_f = max(255 / len(path), 2.5)
        for i in range(len(path)):
            Astar_map[path[i][0], path[i][1]] = [i * k_f, 255, 255 - i * k_f]
        _, i, j, _ = self.check_occupied(self.goal[0], self.goal[1])   
        Astar_map[i, j] = [255, 0, 255]
        Astar_map = cv2.resize(Astar_map, dsize=(0, 0), fx = 4, fy = 4)
        cv2.imwrite('Astar_map.jpg', Astar_map)
        
        while not self.check_goal() and self.step < 100:
            step_time = time.time()
            self.step += 1
            self.position = self.frame.avatar_transform.position
            
            pre_map = self.dep2map()
            self.map = np.maximum(self.map, pre_map)
            
            
            path = self.find_shortest_path(self.position, self.goal, self.map)
            i, j = path[min(3, len(path) - 1)]
            T, x, z = self.get_occupancy_position(i, j)
            assert T == True
            local_goal = [x, z]
            angle = get_angle(forward=np.array(self.frame.avatar_transform.forward),
                                origin=np.array(self.frame.avatar_transform.position),
                                position=np.array([local_goal[0], 0, local_goal[1]]))
            #print('angle:', angle)
            action_time = time.time()
            if np.abs(angle) < ANGLE:
                status = self.move_forward_by(distance=DISTANCE, move_stopping_threshold=0.05)                
                x, y, z = self.frame.avatar_transform.position
                _, i, j, _ = self.check_occupied(x, z)
                traj.append((i, j))
                action = 0
            elif angle > 0:
                status = self.turn_by(angle=-ANGLE, force=500)
                action = 1
            else:
                status = self.turn_by(angle=ANGLE, force=500)
                action = 2
                
            if self.step % 10 == 0:
                print(self.frame.held_objects[Arm.left])
                #self.save_infomation(output_directory)
            action_time = time.time() - action_time
            step_time = time.time() - step_time
            x, y, z = self.frame.avatar_transform.position            
            f.write('step: {}, position: {}, action: {}, goal: {}\n'.format(
                    self.step, \
                    self.frame.avatar_transform.position, \
                    action, \
                    self.goal                    
                ))
            f.write('local_goal: {}, distance: {}, angle: {}, forward: {}\n'.format(
                    local_goal, \
                    self.l2_distance((x, z), self.goal), \
                    angle, \
                    self.frame.avatar_transform.forward
                ))
            f.write('status: {}, action time: {}, step time: {}\n'.format(
                status, action_time, step_time))
        x, y, z = self.frame.avatar_transform.position
        print(f'final step: {self.step}')
        print(f'final distance: {self.l2_distance((x, z), self.goal)}')
        print(f'success: {self.check_goal()}')
        print(f'final time: {time.time() - start}')
        W, H = self.map.shape
        self.paint = np.zeros((W, H, 3))
        self.paint.fill(100)
        self.paint[self.map == 0, 0:3] = 255
        k_f = max(255 / len(traj), 2.5)
        for i in range(len(traj)):
            self.draw(traj[i], [i * k_f, 255, 255 - i * k_f])
        _, i, j, _ = self.check_occupied(self.goal[0], self.goal[1])   
        self.draw((i, j), [255, 0, 255])
        #res = cv2.resize(img, dsize=(54, 140))
        H, W, _ = self.paint.shape
        print(self.paint.shape)
        self.paint = cv2.resize(self.paint, dsize=(0, 0), fx = 4, fy = 4)
        cv2.imwrite('map.jpg', self.paint)   
        
        self.end()


if __name__ == "__main__":
    port = 10075
    Nav(port=port, output_dir = 'images', launch_build=False).run()