DynamicScene/DynamicGeneration
8
1import json2import os3import cv24import numpy as np5import math6import heapq7import datetime8from moviepy.editor import ImageSequenceClip9import tqdm10from utils import get_scene_dir_path, OBJECT_PICTURE_DIR11 12 13class Video_Generator():14 def __init__(self):15 self.frame_radius = 2016 self.frame_thickness = 117 self.frame_outline_color = (0,0,0) #black18 self.frame_filling_color = (255, 255, 255) #white19 self.traj_count = {}20 21 def get_img_coord(self, points, camera_pose, image_width, image_height):22 transformation_matrix = compute_transformation_matrix(camera_pose)23 camera_positions = world_to_camera(points, transformation_matrix)24 projects_points = project_to_2d(camera_positions, camera_pose, image_width, image_height)25 return projects_points26 27 def draw_objframe(self, obj_type, point, background):28 image = background29 obj_path = os.path.join(OBJECT_PICTURE_DIR, '{}.png'.format(obj_type))30 if not os.path.exists(obj_path):31 obj_path = os.path.join(OBJECT_PICTURE_DIR, 'Phone.png')32 obj_img = cv2.imread(obj_path)33 34 #draw frame35 center = (int(point[0]-1.2*self.frame_radius), int(point[1]-1.2*self.frame_radius))36 cv2.circle(image, center, self.frame_radius, self.frame_filling_color, -1)37 cv2.circle(image, center, self.frame_radius, self.frame_outline_color, self.frame_thickness)38 theta = np.pi/839 line_start1 = (int(center[0]+self.frame_radius*np.sin(theta)), int(center[1]+self.frame_radius*np.cos(theta)))40 line_start2 = (int(center[0]+self.frame_radius*np.cos(theta)), int(center[1]+self.frame_radius*np.sin(theta)))41 line_end = (int(center[0] + 1.2*self.frame_radius), int(center[1] + 1.2*self.frame_radius))42 cv2.line(image, line_start1, line_end, self.frame_outline_color, self.frame_thickness)43 cv2.line(image, line_start2, line_end, self.frame_outline_color, self.frame_thickness)44 cv2.circle(image, line_end, 3, (0,0,255), -1)45 46 #put object47 obj_resized = cv2.resize(obj_img, (self.frame_radius, self.frame_radius))48 49 x_start = max(0, obj_resized.shape[0]//2-center[1])50 x_end = min(obj_resized.shape[0], obj_resized.shape[0]//2 + image.shape[0] - center[1])51 y_start = max(0, obj_resized.shape[1]//2-center[0])52 y_end = min(obj_resized.shape[1], obj_resized.shape[1]//2 + image.shape[1] - center[0])53 54 img_x_start = max(0, center[1]-obj_resized.shape[0]//2)55 img_x_end = min(image.shape[0], center[1]+obj_resized.shape[0]//2)56 img_y_start = max(0, center[0]-obj_resized.shape[1]//2) 57 img_y_end = min(image.shape[1], center[0]+obj_resized.shape[1]//2)58 59 image[img_x_start:img_x_end, img_y_start: img_y_end] = obj_resized[x_start:x_end, y_start:y_end]60 61 def add_description(self, activity_name, time, obj_list, receptacle_list, background):62 image = background.copy()63 descrpition_width = 30064 description_bg = np.zeros((background.shape[0], descrpition_width,3), np.uint8)65 res = np.hstack((image, description_bg))66 font = cv2.FONT_HERSHEY_COMPLEX67 font_scale = 0.568 font_color = (0,0,0)69 thickness = 170 line_type = 871 # text_size = cv2.getTextSize(text,font, font_scale, line_type)[0]72 text_x = background.shape[0] + 1073 text_y = descrpition_width//2 -5074 text_y = 5075 line_interval = 3076 cv2.rectangle(res, (background.shape[1],0), (background.shape[1]+descrpition_width, background.shape[0]), (255,255,255),-1)77 texts = ['activity:', activity_name, 'time: ', str(time), 'object movement: ']78 for i, text in enumerate(texts):79 if i%2==0:80 text_x = background.shape[0] + 1081 font_color = (0,0,0)82 else:83 text_x = background.shape[0] + 3084 font_color = (0,0,255)85 cv2.putText(res, text, (text_x, text_y + i*line_interval), font, font_scale, font_color,thickness, line_type)86 87 start_line = 588 for i in range(len(obj_list)):89 obj = obj_list[i].split('|')[0]90 recep = receptacle_list[i].split('|')[0]91 # obj_move_text = '{} -> {}'.format(obj, recep)92 text_x = background.shape[0] + 12093 font_color = (0,0,255)94 95 obj_text_size = cv2.getTextSize(obj, font, font_scale, thickness)[0][0]96 cv2.putText(res, obj, (text_x - 20 - obj_text_size, text_y + (start_line+i)*line_interval), font, font_scale, font_color,thickness, line_type)97 cv2.putText(res, '->', (text_x, text_y + (start_line+i)*line_interval), font, font_scale, font_color,thickness, line_type)98 cv2.putText(res, recep, (text_x + 40, text_y + (start_line+i)*line_interval), font, font_scale, font_color,thickness, line_type)99 100 return res101 102 def draw_traj(self, info, image):103 last_point = info['last_point']104 point = info['point']105 is_end = info['end']106 is_arrow = info['arrow']107 radius = 3108 next_point = (int(point[0]), int(point[1]))109 if last_point is None:110 111 start_color = (0,0,255)112 end_color = (255, 255, 0)113 cv2.circle(image, next_point, radius, start_color, -1)114 return115 pre_point = (int(last_point[0]), int(last_point[1]))116 line_color = (0,0,0)117 line_thickness = 1118 arrowcolor = (0,255,0)119 arrow_thickness = 1120 121 #count122 count = self.traj_count.get((pre_point, next_point),0)123 self.traj_count[(pre_point, next_point)] = count + 1124 count = self.traj_count.get((next_point, pre_point),0)125 self.traj_count[(next_point, pre_point)] = count + 1126 step = 0.2127 line_thickness = min(int(1 + count * step), 5)128 129 #draw130 cv2.line(image, pre_point, next_point, line_color, line_thickness)131 if is_arrow:132 cv2.arrowedLine(image, pre_point, next_point,arrowcolor,arrow_thickness,tipLength=1.5)133 if is_end:134 end_color = (255, 255, 0)135 cv2.circle(image, next_point, radius, end_color, -1)136 137 def get_multiobj_image(self, draw_infos, background):138 image_list = []139 if len(draw_infos)<=0:140 return image_list, background141 activity_name = draw_infos[0]['activity']142 time = draw_infos[0]['time']143 object_list = [info['object'] for info in draw_infos]144 receptacle_list = [info['receptacle'] for info in draw_infos]145 146 image_infos = []147 for draw_info in draw_infos:148 obj = draw_info['object'].split('|')[0]149 points = draw_info['points']150 last_point = None151 for point_num, point in enumerate(points):152 if point_num >= len(image_infos):153 image_infos.append([])154 image_infos[point_num].append({155 'object':obj,156 'point':point,157 'last_point':last_point,158 'end':point_num == len(points)-1,159 'arrow':point_num == len(points)//3160 })161 last_point = (point[0], point[1])162 image_with_traj = background.copy() 163 for image_info in image_infos:164 #draw traj165 for info in image_info:166 self.draw_traj(info, image_with_traj)167 168 #draw obj with frame169 image = image_with_traj.copy()170 for info in image_info:171 self.draw_objframe(info['object'], info['point'], image)172 173 image = self.add_description(activity_name, time, object_list, receptacle_list, image)174 image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)175 image_list.append(image)176 177 return image_list, image_with_traj178 179 def get_receptacle_position_from_meta(self, receptacle_id, metadata):180 position = None181 objects = metadata['objects']182 183 for obj in objects:184 if obj['objectId'] == receptacle_id:185 position = obj['position']186 break187 return position188 189 def get_distance(self, point1, point2):190 return math.sqrt((point1['x'] - point2['x'])**2 + (point1['z'] - point2['z'])**2) 191 192 def get_nearest_point(self, point, reachable_points):193 min_distance = 100000194 nearest_point = None195 for rp in reachable_points:196 distance = self.get_distance(point, rp)197 if distance < min_distance:198 min_distance = distance199 nearest_point = rp200 return nearest_point201 202 def get_path(self, start_position, end_position, reachable_positions):203 res = []204 res.append(start_position)205 start_point = self.get_nearest_point(start_position, reachable_positions)206 target_point = self.get_nearest_point(end_position, reachable_positions)207 208 209 point_id = 0210 open_list = [(0, point_id, start_point)]211 came_from = {tuple((start_point['x'], start_point['z'])):0}212 cost_so_far = {tuple((start_point['x'], start_point['z'])):0}213 214 while open_list:215 current = heapq.heappop(open_list)[-1]216 217 if current == target_point:218 break219 220 for next_point in reachable_positions:221 dis = self.get_distance(current, next_point)222 if dis - 0.25 > 0.001:223 continue224 new_cost = cost_so_far[tuple((current['x'], current['z']))] + 1225 if tuple((next_point['x'], next_point['z'])) not in cost_so_far or new_cost < cost_so_far[tuple((next_point['x'], next_point['z']))]:226 cost_so_far[tuple((next_point['x'], next_point['z']))] = new_cost227 priority = new_cost + abs(next_point['x'] - current['x']) + abs(next_point['z'] - current['z'])228 point_id += 1229 heapq.heappush(open_list, (priority, point_id, next_point))230 came_from[tuple((next_point['x'], next_point['z']))] = current231 232 path = []233 current = target_point234 while current != start_point:235 path.append(current)236 current = came_from[tuple((current['x'], current['z']))]237 238 path.append(start_point)239 path.reverse()240 241 res.extend(path)242 res.append(end_position)243 return res244 245 def init_obj_traj(self, metadata):246 res = {}247 for obj in metadata['objects']:248 if obj['pickupable']:249 parentReceptacle = obj['parentReceptacles']250 if parentReceptacle is not None and 'Floor' in parentReceptacle:251 parentReceptacle.remove('Floor')252 res[obj['objectId']] = parentReceptacle[0] if (parentReceptacle is not None and len(parentReceptacle)) >0 else None253 return res254 255 def save_vedio(self, dynamic_info, background, camera_pose, metadata, reachable_positions, output_path):256 object_traj = self.init_obj_traj(metadata)257 self.traj_count = {}258 image_list = []259 paths_list = []260 for time,timeinfo in tqdm.tqdm(dynamic_info.items()):261 draw_infos = []262 for info in timeinfo:263 info['time'] = time264 target_object_id = info['object']265 target_receptacle_id = info['receptacle']266 target_object_receptacle = object_traj[target_object_id]267 268 if target_object_receptacle == target_receptacle_id:269 continue270 if target_object_receptacle is None:271 for obj in metadata['objects']:272 if obj['objectId'] == target_object_id:273 start_position = obj['position']274 break275 else:276 start_position = self.get_receptacle_position_from_meta(target_object_receptacle, metadata)277 end_position = self.get_receptacle_position_from_meta(target_receptacle_id, metadata)278 path = self.get_path(start_position, end_position, reachable_positions) #path 包括start, end279 image_width = 300280 image_height = 300281 img_path = self.get_img_coord(path, camera_pose, image_width, image_height)282 paths_list.append(img_path)283 draw_info = {284 'time':time,285 'activity':info['activity'],286 'object':target_object_id,287 'receptacle':target_receptacle_id,288 'points':img_path,289 }290 draw_infos.append(draw_info)291 292 object_traj[target_object_id] = target_receptacle_id293 294 time_images,image_with_traj = self.get_multiobj_image(draw_infos, background)295 background = image_with_traj296 image_list.extend(time_images) 297 clip = ImageSequenceClip(image_list, fps=30)298 clip.write_videofile(str(output_path), fps=30, codec="libx264")299 300 def get_dynamic_info(self, schedules):301 res = {}302 for day, day_schedules in schedules.items():303 for activity in day_schedules:304 activity_name = activity['activity']305 start_time = activity['start_time']306 content = activity['content']307 for c in content:308 c['activity'] = activity_name309 time = datetime.datetime.combine(day, start_time)310 if time not in res:311 res[time] = []312 res[time].extend(content)313 return res314 315 def generate(self, schedules, scene_file_name, vedio_path):316 metadata, camera_pose, background, reachable_points = read_scene(scene_file_name)317 318 schekeys = list(schedules.keys())319 schedules_filter = {}320 schedules_filter[schekeys[0]] = schedules[schekeys[0]]321 dynamic_info = self.get_dynamic_info(schedules_filter)322 self.save_vedio(dynamic_info, background, camera_pose, metadata, reachable_points, vedio_path)323 324def read_scene(scene_file_name):325 data_dir = get_scene_dir_path(scene_file_name)326 metadata = json.load(open(os.path.join(data_dir, 'metadata.json'),'r',encoding='utf-8'))327 camera_pose = json.load(open(os.path.join(data_dir, 'camera_pose.json'),'r',encoding='utf-8'))328 background = cv2.imread(os.path.join(data_dir, 'background.png'))329 reachable_points = json.load(open(os.path.join(data_dir, 'reachablePositions.json'),'r',encoding='utf-8'))330 331 return metadata, camera_pose, background, reachable_points332 333def compute_transformation_matrix(camera_pose):334 position = camera_pose['position']335 rotation = camera_pose['rotation']336 337 translation_matrix = np.array([338 [1, 0, 0, -position['x']],339 [0, 1, 0, -position['y']],340 [0, 0, 1, -position['z']],341 [0, 0, 0, 1]342 ])343 344 theta_x = np.radians(rotation['x'])345 rotation_matrix_x = np.array([346 [1, 0, 0, 0],347 [0, np.cos(theta_x), -np.sin(theta_x), 0],348 [0, np.sin(theta_x), np.cos(theta_x), 0],349 [0, 0, 0, 1]350 ])351 352 transformation_matrix = np.dot(rotation_matrix_x, translation_matrix)353 return transformation_matrix 354 355def world_to_camera(positions, transformation_matrix):356 camera_positions = []357 for pos in positions:358 world_pos = np.array([pos['x'], pos['y'], pos['z'], 1])359 camera_pos = np.dot(transformation_matrix, world_pos)360 camera_positions.append(camera_pos)361 return camera_positions 362 363def project_to_2d(camera_positions, camera_pose, image_width, image_height):364 fov = camera_pose['fieldOfView']365 aspect_ratio = image_width / image_height366 f = 1 / np.tan(np.radians(fov) / 2)367 projection_matrix = np.array([368 [f / aspect_ratio, 0, 0, 0],369 [0, f, 0, 0],370 [0, 0, 1, 0]371 ])372 373 projected_points = []374 for pos in camera_positions:375 projected = np.dot(projection_matrix, pos)376 projected /= projected[2]377 x = (projected[0] + 1) * image_width / 2378 y = (1 - projected[1]) * image_height / 2379 x = image_width - x380 projected_points.append((x, y))381 return projected_points 382