def process_all(images_path, depth_path, output_path): img_names = my_utils.os_listdir(images_path) depth_names = my_utils.os_listdir(depth_path) beta = 0 pbar = tqdm.tqdm(total=len(img_names)) for name_file, depth_file in zip(img_names, depth_names): pbar.update(1) img = cv2.imread(os.path.join(images_path, name_file)) gray_img = cv2.cvtColor(img, cv2.COLOR_BGRA2GRAY) # divided by 256 to convert it into metres original_depth = cv2.imread(os.path.join(depth_path, depth_file), cv2.IMREAD_UNCHANGED) / 256 smooth_depth = improve_depth(gray_img, original_depth, threshold=beta) np.save(os.path.join(output_path, name_file), smooth_depth)
def load_streak_database(self): ''' Function to load and store the texture maps. Streaks are stored in a list. ''' if not os.path.exists(self.streaks_path): print("No existing path for streak database (", self.streaks_path, ")") exit(-1) tmp = [] norm_coeff_path = self.norm_coeff_path norm_coeffs = {} with open(norm_coeff_path, 'r') as file: lines = file.readlines() for line in lines: if line[:2] == 'cv': coeff = int(line[2:]) continue norm_coeffs.update({ coeff: [float(v) for v in line.split('\n')[0].split(' ')[:-1]] }) for file_name in my_utils.os_listdir(self.streaks_path): name = os.path.splitext(file_name)[0] coeff, osc = name.split('_') if len(coeff) == 3: coeff = int(coeff[-1:]) else: coeff = int(coeff[-2:]) osc = int(osc[-1:]) drop_image = cv2.imread(os.path.join(self.streaks_path, file_name), cv2.IMREAD_ANYDEPTH) drop_image = cv2.cvtColor(drop_image, cv2.COLOR_GRAY2BGR) drop_image_norm = ((255.0 * norm_coeffs[coeff][osc] * drop_image) / 65535.0).astype(np.uint8) tmp.append(drop_image_norm) self.ratio = np.append(self.ratio, tmp[-1].shape[1] / tmp[-1].shape[0]) self.ratio = np.unique(self.ratio) self.streaks_light = np.array(tmp)
def run(self): process_t0 = time.time() folders_num = len(self.images) # case for any number of sequences and supported rain intensities for folder_idx, sequence in enumerate(self.sequences): folder_t0 = time.time() print('\nSequence: ' + sequence) sim_num = len(self.particles[sequence]) depth_folder = self.depth[sequence] for sim_idx, sim_weather in enumerate(self.weather): weather, fallrate = sim_weather["weather"], sim_weather[ "fallrate"] out_seq_dir = os.path.join(self.output_root, sequence) out_dir = os.path.join(out_seq_dir, weather, '{}mm'.format(fallrate)) sim_file = self.particles[sequence][sim_idx] # Resolve output path path_exists = os.path.exists(out_dir) if path_exists: if self.conflict_strategy == "skip": pass elif self.conflict_strategy == "overwrite": pass elif self.conflict_strategy == "rename_folder": out_dir_, out_shift = out_dir, 0 while os.path.exists(out_dir_ + '_copy%05d' % out_shift): out_shift += 1 out_dir = out_dir_ + '_copy%05d' % out_shift else: raise NotImplementedError # Create directory os.makedirs(out_dir, exist_ok=True) # Default fog-like rain parameters fog_params = { "rain_intensity": fallrate, "focal": self.focal, "f_number": self.f_number, "angle": 90, "exposure": self.exposure, "camera_gain": self.camera_gain } if "nuscenes" in self.dataset: files = self.images[sequence] depth_files = self.depth[sequence] depth_file = depth_files[0] assert depth_file.endswith( ".npy" ), "nuscenes processing only works with .npy for depth" # depth = np.load(depth_file) if "gan" in self.dataset: # HARDCODED since these values are not known in nuscenes_gan imW, imH = (1600, 900) else: img = cv2.imread(files[0]) imH, imW = img.shape[0:2] else: files = natsorted( np.array([ os.path.join(self.images[sequence], picture) for picture in my_utils.os_listdir( self.images[sequence]) ])) depth_files = natsorted( np.array([ os.path.join(depth_folder, depth) for depth in my_utils.os_listdir(depth_folder) ])) im = files[0] if im.endswith(".png"): imH, imW = cv2.imread(im).shape[0:2] elif im.endswith(".npy"): imH, imW = np.load(im).shape[0:2] else: raise Exception("Invalid extension", im) imH = imH // self.settings["render_scale"] imW = imW // self.settings["render_scale"] if self.camera_gain: fog_params["camera_gain"] = self.camera_gain print('Simulation: rain {}mm/hr'.format(fallrate)) # loading StreaksDBManager, RainRenderer, FOVComputation and EnvMapGenerator self.db = DBManager(streaks_path_xml=sim_file, streaks_path=self.texture, norm_coeff_path=self.norm_coeff) self.renderer = RainRenderer(focal=self.focal, f_number=self.f_number, focus_plane=6, radius=10, fov=165) self.fov_comp = FovComputation(camera=np.array([0, 0, 0])) map_generator = EnvironmentMapGenerator(self.focal, imW, imH) # loading fog class FOG = add_attenuation.FogRain(**fog_params) # loading calib files for frames calib_files = self.calib[sequence] # loading streaks from Streaks Database self.db.load_streak_database() # creating drops based on the simulator file self.db.load_streaks_from_xml(self.dataset, self.settings, [imW, imH], use_pickle=False, verbose=self.verbose) env_map_input = self.env_type if self.irrad_type == 'garg' else self.env_type + '_' + self.irrad_type frame_render_dict = list(self.db.streaks_simulator.values()) f_start, f_end, f_step = self.frame_start, self.frame_end, self.frame_step f_end = len(files) if f_end is None else min(f_end, len(files)) if self.frames: # prone to go "boom", so we clip and remove 'wrong' ids idx = np.unique(np.clip(self.frames, 0, f_end - 1)).tolist() else: idx = list(range(f_start, f_end, f_step)) # to make it f_num = len(idx) sim_t0 = time.time() print("{} images".format(len(idx))) frames_exist_nb = 0 for f_idx, i in enumerate(idx): image_file = files[i] depth_file = depth_files[i] # Compute frame index (independent of starting frame) to allow deterministic / reproducible rendering if self.dataset == 'nuscenes': # It could be useful for other dataset, but, for the moment, lets consign this little gem of a # code to nuscenes # If f_end was not supplied, we could see if the number of file is not equal to the number of # simulated drop... if so, lets remap them render_ix = np.linspace(0, len(frame_render_dict), len(files), endpoint=False, dtype=int) f_name_idx = render_ix[i] else: f_name_idx = i assert os.path.exists( image_file), "Image file {} does not exist".format( image_file) assert os.path.exists( depth_file), "Depth file {} does not exist".format( depth_file) # Ensure deterministic behavior np.random.seed(f_name_idx) frame_t0 = time.time() frame = frame_render_dict[f_name_idx % len(frame_render_dict)] file_name = os.path.split(image_file)[-1] out_rainy_path = os.path.join( out_dir, 'rainy_image', '{}.png'.format(file_name[:-4])) out_rainy_mask_path = os.path.join( out_dir, 'rain_mask', '{}.png'.format(file_name[:-4])) out_env_path = os.path.join( out_seq_dir, 'envmap', '{}.png'.format(file_name[:-4])) frame_exists = os.path.exists( out_rainy_path) or os.path.exists(out_rainy_mask_path) if frame_exists: if self.conflict_strategy == "skip": frames_exist_nb += 1 continue elif self.conflict_strategy == "overwrite": pass else: raise NotImplementedError # TODO adding functions of depth weighting for other dataset if USE_DEPTH_WEIGHTING == 1 and calib_files is not None: # Compute the drop depth map to allow weighting envmap from drop FOV depth_drop_evaluator = DropDepthMap( filename=calib_files[0]) # flush should happens after a while if self.verbose: sys.stdout.write('\r' + my_utils.process_eta_str( process_t0, folder_idx, folders_num, folder_t0, sim_idx, sim_num, sim_t0, f_idx, f_num, frame_t0) + ' ') # two copies of bg because one is used for rain drop calculation # and the other is changed after adding each drop bg = cv2.imread(image_file) / 255.0 if self.settings["render_scale"] != 1: bg = cv2.resize( bg, (int(bg.shape[1] // self.settings["render_scale"]), int(bg.shape[0] // self.settings["render_scale"]))) if FOG_ATT == 1: # Depth map is used in weighting the luminance effect of the environment map on a single drop if depth_file.endswith(".png"): depth = cv2.imread(depth_file, cv2.IMREAD_UNCHANGED) if depth is None: print('Missing/Corrupted depth data (%s)' % depth_file) continue depth = depth.astype(np.float32) / 256. elif depth_file.endswith(".npy"): depth = np.load(depth_file) else: raise Exception("Invalid extension") # Apply depth and render scale depthHW = np.array([ int((depth.shape[0] * self.settings["depth_scale"]) // self.settings["render_scale"]), int((depth.shape[1] * self.settings["depth_scale"]) // self.settings["render_scale"]) ]) if not np.all(depth.shape[:2] == depthHW): depth = cv2.resize(depth, (depthHW[1], depthHW[0])) assert (np.all(depth.shape[:2] <= bg.shape[:2]) ), "Depth cannot be larger than the image" # Strategy to apply if RGB and Depth size mismatch if not np.all(depth.shape[:2] == bg.shape[:2]): # print("\nDepth {} size ({},{}) differs from image ({},{}). Will assume depth is crop centered.".format(image_file, depth.shape[0], depth.shape[1], bg.shape[0], bg.shape[1])) bg = my_utils.crop_center(bg, depth.shape[0], depth.shape[1]) else: # In that case, no need for depth, but it's still used down in the code, for more less nothing depth = np.zeros((bg.shape[1], bg.shape[0]), np.float) rainy_bg = FOG.fog_rain_layer(bg, depth) # rain layer is the image of the rendered rain drops blended with the background rain_layer = np.zeros((bg.shape[0], bg.shape[1], 4), np.float64) # rainy_mask keeps track of the pixels of the background that have already been occupied by rain. # This is used in cases of occluding or partially occluded drops rainy_mask = np.zeros((bg.shape[0], bg.shape[1]), np.float64) rainy_saturation_mask = np.zeros( (bg.shape[0], bg.shape[1], 3), np.float64) # Environment map of the frame using (Christopher Cameron, 2005): # http://www.cs.cmu.edu/afs/andrew/scs/cs/15-463/f05/pub/www/projects/fproj/cmcamero/report.pdf if 'ours' in env_map_input: # print('\nGenerating environment map') self.BGR_env_map = map_generator.generate_map(rainy_bg) elif 'pano' in env_map_input: # print('\nLoading Environment Pano') self.BGR_env_map = cv2.imread( os.path.join('../data', 'panos', file_name)) / 255.0 else: raise NotImplementedError self.env_map_xyY = my_utils.convert_rgb_to_xyY( self.BGR_env_map[..., ::-1]) self.env_map_xyY[np.isnan(self.env_map_xyY)] = 0 self.solid_angle_map = solid_angle.get_solid_angles( self.BGR_env_map) # Render only streaks inside the frame streak_dict = frame.streaks streak_dict = { keys: values for keys, values in streak_dict.items() if 1 <= values.max_width < max(imH, imW) and 1 <= values.length < max(imH, imW) and ( (0 <= values.image_position_start[0] < imW and 0 <= values.image_position_start[1] < imH) or (0 <= values.image_position_end[0] < imW and 0 <= values.image_position_end[1] < imH)) } if USE_DEPTH_WEIGHTING == 1: xyz_coord = depth_drop_evaluator.get_world_points( depth) assert streak_dict.__len__() <= 2 ** 16, \ "Assert that the number of drops doesn't overpass the uint16 rain_mask capacity" streak_list = list(streak_dict.values()) drop_num = len(streak_list) drop_process_t0 = time.time() for drop_idx, drop_dict in enumerate(streak_list): # Returns the rainy image, rainy_mask, drop, blended drop and the starting coord of # the drop in image rainy_bg, rainy_mask, rainy_saturation_mask, \ drop, blended_drop, minC = self.compute_drop(bg, drop_dict, rainy_bg, rainy_mask, rainy_saturation_mask) if blended_drop is not None: rain_layer = self.renderer.make_rain_layer( drop, blended_drop, rain_layer, rainy_mask, minC) else: print( "Trace: rain drop {} in sequence {} in image {} ({})" .format(drop_idx, sequence, f_idx, f_name_idx)) # Compute progress avg_drop_time = (time.time() - drop_process_t0) / (drop_idx + 1) if self.verbose or drop_idx == 0: sys.stdout.write('\r' + my_utils.process_eta_str( process_t0, folder_idx, folders_num, folder_t0, sim_idx, sim_num, sim_t0, f_idx, f_num, frame_t0, drop_idx, drop_num) + '\t\t' + '%.1fms /drop' % (1000. * avg_drop_time) + ' ') # Create all output directories os.makedirs(os.path.dirname(out_rainy_path), exist_ok=True) os.makedirs(os.path.dirname(out_rainy_mask_path), exist_ok=True) if self.save_envmap: os.makedirs(os.path.dirname(out_env_path), exist_ok=True) # mean contrast adjusted image rainy_bg_mean = np.mean(rainy_bg) bg_mean = np.mean(bg) difference_mean = rainy_bg_mean - bg_mean rainy_bg_copy = rainy_bg - difference_mean plt.imsave(out_rainy_path, np.clip(rainy_bg_copy[..., ::-1], 0, 1)) plt.imsave(out_rainy_mask_path, rainy_mask) if self.save_envmap: plt.imsave(out_env_path, self.BGR_env_map[..., ::-1]) if frames_exist_nb > 0: print("Skipped {}/{} already existing renderings".format( frames_exist_nb, len(idx))) print("\n\nEnd of the simulation")