diff --git a/omnigibson/scene_graphs/graph_builder.py b/omnigibson/scene_graphs/graph_builder.py index 81d668309..1f987748c 100644 --- a/omnigibson/scene_graphs/graph_builder.py +++ b/omnigibson/scene_graphs/graph_builder.py @@ -4,7 +4,6 @@ import torch as th from matplotlib import pyplot as plt from PIL import Image -from torchvision.transforms import ToPILImage, ToTensor from omnigibson import object_states from omnigibson.object_states.factory import get_state_name @@ -12,6 +11,7 @@ from omnigibson.robots import BaseRobot from omnigibson.sensors import VisionSensor from omnigibson.utils import transform_utils as T +from omnigibson.utils.numpy_utils import pil_to_tensor def _formatted_aabb(obj): @@ -289,14 +289,10 @@ def _draw_graph(): robot_view = (robot_camera_sensor.get_obs()[0]["rgb"][..., :3]).to(th.uint8) imgheight, imgwidth, _ = robot_view.shape - pil_transform = ToPILImage() - torch_transform = ToTensor() - - # check imgheight and imgwidth; if they are too small, we need to upsample the image to 1280x1280 + # check imgheight and imgwidth; if they are too small, we need to upsample the image to 640x640 if imgheight < 640 or imgwidth < 640: - robot_view = torch_transform( - pil_transform((robot_view.permute(2, 0, 1).cpu())).resize((640, 640), Image.BILINEAR) - ).permute(1, 2, 0) + # Convert to PIL Image to upsample, then write back to tensor + robot_view = pil_to_tensor(Image.fromarray(robot_view.cpu().numpy()).resize((640, 640), Image.BILINEAR)) imgheight, imgwidth, _ = robot_view.shape figheight = 4.8 diff --git a/omnigibson/utils/numpy_utils.py b/omnigibson/utils/numpy_utils.py index 1d540c27b..fda0a3164 100644 --- a/omnigibson/utils/numpy_utils.py +++ b/omnigibson/utils/numpy_utils.py @@ -15,3 +15,7 @@ def vtarray_to_torch(vtarray, dtype=th.float32, device="cpu"): else: assert device.startswith("cuda") return th.tensor(np.array(vtarray), dtype=dtype, device=device) + + +def pil_to_tensor(pil_image): + return th.tensor(np.array(pil_image), dtype=th.uint8) diff --git a/setup.py b/setup.py index f69a32992..82eae318b 100644 --- a/setup.py +++ b/setup.py @@ -22,8 +22,8 @@ packages=find_packages(), install_requires=[ "gymnasium>=0.28.1", - "numpy~=1.23.5", - "scipy~=1.10.1", + "numpy>=1.23.5", + "scipy>=1.10.1", "GitPython~=3.1.40", "transforms3d~=0.4.1", "networkx~=3.2.1", @@ -35,7 +35,7 @@ "h5py~=3.10.0", "cryptography~=41.0.7", "bddl~=3.5.0", - "opencv-python~=4.8.1", + "opencv-python>=4.8.1", "nest_asyncio~=1.5.8", "imageio~=2.33.1", "imageio-ffmpeg~=0.4.9", @@ -46,8 +46,7 @@ "aenum~=3.1.15", "rtree~=1.2.0", "graphviz~=0.20", - "numba~=0.60.0", - "torchvision~=0.18.1", + "numba>=0.60.0", ], extras_require={ "isaac": ["isaacsim-for-omnigibson>=4.1.0"],