Upload folder using huggingface_hub

lmdeploy_infer.py

@@ -0,0 +1,243 @@
+import base64
+import os
+import ast
+from io import BytesIO
+from typing import List, Union
+
+from PIL import Image, ImageFile
+import numpy as np
+from scipy.spatial.transform import Rotation
+
+from lmdeploy import pipeline, TurbomindEngineConfig, ChatTemplateConfig
+IMAGE_TOKEN = '<IMAGE_TOKEN>'
+
+def normalize_quaternion(quat):
+ return np.array(quat) / np.linalg.norm(quat, axis=-1, keepdims=True)
+
+def quaternion_to_discrete_euler(quaternion, bins_num=256):
+ euler = Rotation.from_quat(quaternion).as_euler('xyz', degrees=True) + 180
+ resolution = 360 / bins_num
+ disc = np.around((euler / resolution)).astype(int)
+ disc[disc == bins_num] = 0
+ return disc
+
+def discrete_euler_to_quaternion(discrete_euler, bins_num=256):
+ resolution = 360 / bins_num
+ euler = (discrete_euler * resolution) - 180
+ return Rotation.from_euler('xyz', euler, degrees=True).as_quat()
+
+
+class RotationActionDiscretizer:
+ def __init__(self, bins_num=256, min_action=-1, max_action=1):
+ """
+ Note: the input action is quaternion
+ Args: bins_num: Number of bins to discretize the rotation space into.
+ """
+ self.bins_num = bins_num
+
+ def discretize(self, action: Union[np.ndarray, List[float]], degrees=False):
+ # Check if the input action is quaternion or euler
+ if len(action) == 4:
+ return quaternion_to_discrete_euler(normalize_quaternion(action), bins_num=self.bins_num)
+ else:
+ return quaternion_to_discrete_euler(
+ normalize_quaternion(Rotation.from_euler('xyz', action, degrees=degrees).as_quat()),
+ bins_num=self.bins_num
+ )
+
+ def undiscretize(self, discrete_action):
+ return normalize_quaternion(discrete_euler_to_quaternion(discrete_action, bins_num=self.bins_num))
+ 
+ def get_action_space(self):
+ return self.bins_num
+ 
+ def generate_discrete_special_tokens(self)-> List[str]:
+ return [f"<rot{i}>" for i in range(self.bins_num)]
+ 
+ def map_4d_quaternion_to_special_tokens(self, action) -> List[str]:
+ discretiezd_action = self.discretize(action)
+ return [f"<rot{action}>" for action in discretiezd_action]
+ 
+ def map_roll_pitch_yaw_to_special_tokens(self, roll_pitch_yaw: Union[np.ndarray, List[float]], degrees=False) -> List[str]:
+ discretized_action = self.discretize(roll_pitch_yaw, degrees)
+ return [f"<rot{a}>" for a in discretized_action]
+
+
+class TranslationActionDiscretizer:
+ def __init__(self, bins_num=256, min_action=-1, max_action=1):
+ self.bins_num = bins_num
+ self.min_action = min_action
+ self.max_action = max_action
+ 
+ # Create Uniform Bins + Compute Bin Centers
+ self.bins = np.linspace(min_action, max_action, bins_num)
+ self.bin_centers = (self.bins[:-1] + self.bins[1:]) / 2.0
+
+ def discretize(self, action: np.ndarray):
+ action = np.clip(action, a_min=float(self.min_action), a_max=float(self.max_action))
+ discretized_action = np.digitize(action, self.bins)
+ return discretized_action
+
+ def undiscretize(self, discrete_action):
+ """
+ NOTE =>> Because of the way the actions are discretized w.r.t. the bins (and not the bin centers), the
+ digitization returns bin indices between [1, # bins], inclusive, when there are actually only
+ (# bins - 1) bin intervals.
+
+ Therefore, if the digitization returns the last possible index, we map this to the last bin interval.
+
+ EXAMPLE =>> Let's say self._bins has 256 values. Then self._bin_centers has 255 values. Digitization returns
+ indices between [1, 256]. We subtract 1 from all indices so that they are between [0, 255]. There
+ is still one index (i==255) that would cause an out-of-bounds error if used to index into
+ self._bin_centers. Therefore, if i==255, we subtract 1 from it so that it just becomes the index of
+ the last bin center. We implement this simply via clipping between [0, 255 - 1].
+ """
+ 
+ discrete_action = np.clip(discrete_action - 1, a_min=0, a_max=self.bin_centers.shape[0] - 1)
+ undiscretized_action = self.bin_centers[discrete_action]
+ 
+ # Clamp the result to the action bounds
+ return np.clip(undiscretized_action, self.min_action, self.max_action)
+ 
+ def get_action_space(self):
+ return self.bins_num
+ 
+ def generate_discrete_special_tokens(self)-> List[str]:
+ return [f"<loc{i}>" for i in range(self.bins_num)]
+ 
+ def map_3d_action_to_special_tokens(self, action) -> List[str]:
+ discretiezd_action = self.discretize(action)
+ return [f"<loc{action}>" for action in discretiezd_action]
+ 
+ 
+class OpennessActionDiscretizer:
+ def __init__(self, bins_num=256, min_openness=0, max_openness=1):
+ """
+ Args:
+ bins_num: Number of bins to discretize the openness space into.
+ min_openness: Minimum openness of the gripper.
+ max_openness: Maximum openness of the gripper.
+ """
+ self.bins_num = bins_num
+ self.min_openness = min_openness
+ self.max_openness = max_openness
+ 
+ # Create Uniform Bins + Compute Bin Centers
+ self.bins = np.linspace(min_openness, max_openness, bins_num)
+ self.bin_centers = (self.bins[:-1] + self.bins[1:]) / 2.0
+
+ def discretize(self, openness: float):
+ openness = np.clip(openness, a_min=self.min_openness, a_max=self.max_openness)
+ discretized_openness = np.digitize(openness, self.bins)
+ return discretized_openness
+
+ def undiscretize(self, discrete_openness):
+ discrete_openness = np.clip(discrete_openness - 1, a_min=0, a_max=self.bin_centers.shape[0] - 1)
+ return self.bin_centers[discrete_openness]
+ 
+ def get_action_space(self):
+ return self.bins_num
+ 
+ def generate_discrete_special_tokens(self) -> List[str]:
+ return [f"<open{i}>" for i in range(self.bins_num)]
+ 
+ def map_openness_to_special_tokens(self, openness) -> List[str]:
+ discretized_openness = self.discretize(openness)
+ return [f"<open{discretized_openness}>"]
+
+# def construct_lmdeploy_tasks(jsonl_path):
+# data = load_jsonl(jsonl_path)
+ 
+# lmdeploy_tasks = []
+# for sample_idx, item in enumerate(data):
+ 
+# langs = item["conversations"][0]["value"]
+# langs = langs.replace("<image>", IMAGE_TOKEN)
+# image_urls = [
+# os.path.join(sample_save_folder, f"{sample_idx}_{im_idx}.png") for im_idx in range(len(item["image"]))
+# ]
+# gt_lang = item["conversations"][1]["value"]
+# lmdeploy_tasks.append((langs, image_urls, gt_lang))
+ 
+# return lmdeploy_tasks
+
+def load_image_from_base64(image: Union[bytes, str]) -> Image.Image:
+ """load image from base64 format."""
+ return Image.open(BytesIO(base64.b64decode(image)))
+
+def load_image(image_url: Union[str, Image.Image]) -> Image.Image:
+ """load image from url, local path or openai GPT4V."""
+ FETCH_TIMEOUT = int(os.environ.get('LMDEPLOY_FETCH_TIMEOUT', 10))
+ headers = {
+ 'User-Agent':
+ 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 '
+ '(KHTML, like Gecko) Chrome/58.0.3029.110 Safari/537.3'
+ }
+ try:
+ ImageFile.LOAD_TRUNCATED_IMAGES = True
+ if isinstance(image_url, Image.Image):
+ img = image_url
+ else:
+ # Load image from local path
+ img = Image.open(image_url)
+
+ # check image valid
+ img = img.convert('RGB')
+ except Exception as error:
+ if isinstance(image_url, str) and len(image_url) > 100:
+ image_url = image_url[:100] + ' ...'
+ print(f'{error}, image_url={image_url}')
+ # use dummy image
+ img = Image.new('RGB', (32, 32))
+
+ return img
+
+model = '/mnt/petrelfs/huangsiyuan/VLA/InternVL/internvl_chat/output/internvla_8b_1node_with_visual_traces_wo_sp_token_w_cam/VLA8B_V1'
+pipe = pipeline(model, backend_config=TurbomindEngineConfig(session_len=2048), chat_template_config=ChatTemplateConfig(model_name='internvl2-internlm2'))
+
+TRANS_MAX = 0.275
+TRANS_MIN = -0.275
+
+ROT_MIN = -0.350
+ROT_MAX = 0.395
+
+OPEN_MIN = -0.388
+OPEN_MAX = 0.300
+
+translation_bins_num = 256
+rotation_bins_num = 256
+openness_bins_num = 256
+translation_action_discretizer = TranslationActionDiscretizer(bins_num=translation_bins_num, max_action=TRANS_MAX, min_action=TRANS_MIN)
+rotation_action_discretizer = RotationActionDiscretizer(bins_num=rotation_bins_num, min_action=ROT_MIN, max_action=ROT_MAX)
+openness_action_discretizer = OpennessActionDiscretizer(bins_num=openness_bins_num, min_openness=OPEN_MIN, max_openness=OPEN_MAX)
+
+VQA_FORMAT = f"{IMAGE_TOKEN}\n {IMAGE_TOKEN}\n Given the observation images from the wrist camera mounted at CAM_PARAM and the overhead camera mounted at CAM_PARAM, please provide the action that the robot should take to finish the task: TASK"
+# question_template = "<image>\n <image>\n Given the observation images from the wrist camera mounted at <cam>[256,89,256,236,129,181]</cam> and the overhead camera mounted at <cam>[82,1,256,54,128,98]</cam>, please provide the action that the robot should take to finish the task: place a chess piece on the chessboar"
+
+# cam_params
+wrist_cam_pose = [0.3618544138321802, -0.08323374464523976, 0.41759402329169787, 2.6584232953914344, 0.035482430406705845, 1.2906347836099603]
+overhead_cam_pose = [-0.09877916942983442, -0.3919519409041736, 0.4780865865815033, -1.8237694898473762, -0.012183613523460979, -0.746683044221379]
+cam_pose_list = [wrist_cam_pose, overhead_cam_pose]
+for cam_pose in cam_pose_list:
+ cam_xyz_token = translation_action_discretizer.discretize(np.array(cam_pose[:3]))
+ cam_rpy_token = rotation_action_discretizer.discretize(np.array(cam_pose[3:6]))
+ cam_action_tokens = [cam_xyz_token[0], cam_xyz_token[1], cam_xyz_token[2], cam_rpy_token[0], cam_rpy_token[1], cam_rpy_token[2]]
+ cam_action_tokens_str = "<cam>[" + ",".join(map(str, cam_action_tokens)) + "]</cam>"
+ VQA_FORMAT = VQA_FORMAT.replace("CAM_PARAM", cam_action_tokens_str, 1)
+ 
+# task lang
+task = "Pick up the green object from the table and put it in the bowl"
+VQA_FORMAT = VQA_FORMAT.replace("TASK", task)
+
+img1 = "/mnt/petrelfs/huangsiyuan/VLA/droid_action_tasks_internvl/sample_images/2_0.png"
+img2 = "/mnt/petrelfs/huangsiyuan/VLA/droid_action_tasks_internvl/sample_images/2_1.png"
+images = [load_image(img1), load_image(img2)] # only need to return the PIL.Image object
+response = pipe((VQA_FORMAT, images))
+print(response.text)
+print("gt: [124,137,104,126,130,129,233]")
+action_list = np.array(ast.literal_eval(response.text))
+xyz = translation_action_discretizer.undiscretize(action_list[:3])
+rpy = rotation_action_discretizer.undiscretize(action_list[3:6])
+openness = openness_action_discretizer.undiscretize(action_list[6])
+
+print(f"xyz: {xyz}, rpy: {rpy}, openness: {openness}")