def load_image(image_path, new_size=None): """ Load image. Args: image_path: Path to image to load. new_size: new image dimensions(tuple). Returns: numpy array(image), image_path """ image_path = get_abs_path(image_path, verify=True) image = cv2.imread(image_path) if image is None: LOGGER.warning(f'Failed to read image: {image_path}') return if new_size: return cv2.resize(image, new_size) return image, image_path
def update_data(self, bbs_aug, frame_before, image_aug, new_name, new_path): """ Update new bounding boxes data and save augmented image. Args: bbs_aug: Augmented bounding boxes frame_before: pandas DataFrame containing pre-augmentation data. image_aug: Augmented image as numpy nd array. new_name: new image name to save image. new_path: path to save the image. Returns: None """ frame_after = pd.DataFrame( bbs_aug.remove_out_of_image().clip_out_of_image().bounding_boxes, columns=['x1y1', 'x2y2'], ) if (frame_after.empty ): # some post-augmentation photos do not contain bounding boxes LOGGER.warning( f'\nskipping image: {new_name}: no bounding boxes after ' f'augmentation') return frame_after = pd.DataFrame( np.hstack( (frame_after['x1y1'].tolist(), frame_after['x2y2'].tolist())), columns=['x1', 'y1', 'x2', 'y2'], ).astype('int64') frame_after['object_type'] = frame_before['object_type'].values frame_after['object_id'] = frame_before['object_id'].values frame_after['image'] = new_name for index, row in frame_after.iterrows(): x1, y1, x2, y2, object_type, object_id, image_name = row bx, by, bw, bh = self.calculate_ratios(x1, y1, x2, y2) self.augmentation_data.append( [image_name, object_type, object_id, bx, by, bw, bh]) cv2.imwrite(new_path, image_aug)