def generate_new_anchors(self, new_anchors_conf): """ Create new anchors according to given configuration. Args: new_anchors_conf: A dictionary containing the following keys: - anchor_no and one of the following: - relative_labels - from_xml - coordinate_labels Returns: None """ anchor_no = new_anchors_conf.get('anchor_no') if not anchor_no: raise ValueError(f'No "anchor_no" found in new_anchors_conf') labels_frame = self.get_adjusted_labels(new_anchors_conf) relative_dims = np.array( list( zip( labels_frame['relative_width'], labels_frame['relative_height'], ))) centroids, _ = k_means(relative_dims, anchor_no, frame=labels_frame) self.anchors = ( generate_anchors(self.image_width, self.image_height, centroids) / self.input_shape[0]) LOGGER.info('Changed default anchors to generated ones')
def create_new_dataset(self, new_dataset_conf): """ Create a new TFRecord dataset. Args: new_dataset_conf: A dictionary containing the following keys: - dataset_name(required) str representing a name for the dataset - test_size(required) ex: 0.1 - augmentation(optional) True or False - sequences(required if augmentation is True) - aug_workers(optional if augmentation is True) defaults to 32. - aug_batch_size(optional if augmentation is True) defaults to 64. And one of the following is required: - relative_labels: Path to csv file with the following columns: ['Image', 'Object Name', 'Object Index', 'bx', 'by', 'bw', 'bh'] - coordinate_labels: Path to csv file with the following columns: ['Image Path', 'Object Name', 'Image Width', 'Image Height', 'X_min', 'Y_min', 'X_max', 'Y_max', 'Relative Width', 'Relative Height', 'Object ID'] - from_xml: True or False to parse from xml_labels folder. """ LOGGER.info(f'Generating new dataset ...') test_size = new_dataset_conf.get('test_size') labels_frame = self.generate_new_frame(new_dataset_conf) save_tfr( labels_frame, os.path.join('data', 'tfrecords'), new_dataset_conf['dataset_name'], test_size, self, )
def k_means(relative_sizes, k, distance_func=np.median, frame=None): """ Calculate optimal anchor relative sizes. Args: relative_sizes: 2D array of relative box sizes. k: int, number of clusters. distance_func: function to calculate distance. frame: pandas DataFrame with the annotation data(for visualization purposes). Returns: Optimal relative sizes. """ box_number = relative_sizes.shape[0] last_nearest = np.zeros((box_number, )) centroids = relative_sizes[np.random.randint(0, box_number, k)] old_distances = np.zeros((relative_sizes.shape[0], k)) iteration = 0 while True: distances = 1 - iou(relative_sizes, centroids, k) print(f'Iteration: {iteration} Loss: ' f'{np.sum(np.abs(distances - old_distances))}') old_distances = distances.copy() iteration += 1 current_nearest = np.argmin(distances, axis=1) if (last_nearest == current_nearest).all(): LOGGER.info(f'Generated {len(centroids)} anchors in ' f'{iteration} iterations') return centroids, frame for anchor in range(k): centroids[anchor] = distance_func( relative_sizes[current_nearest == anchor], axis=0) last_nearest = current_nearest
def create_new_dataset(self, new_dataset_conf): """ Create a new TFRecord dataset. Args: new_dataset_conf: A dictionary containing the following keys: - dataset_name(required) str representing a name for the dataset - test_size(required) ex: 0.1 - augmentation(optional) True or False - sequences(required if augmentation is True) - aug_workers(optional if augmentation is True) defaults to 32. - aug_batch_size(optional if augmentation is True) defaults to 64. And one of the following is required: - relative_labels: Path to csv file with the following columns: ['image', 'object_name', 'object_index', 'bx', 'by', 'bw', 'bh'] - coordinate_labels: Path to csv file with the following columns: ['image_path', 'object_name', 'img_width', 'img_height', 'x_min', 'y_min', 'x_max', 'y_max', 'relative_width', 'relative_height', 'object_id'] - xml_labels_folder: Path to folder containing xml labels. """ LOGGER.info(f'Generating new dataset ...') test_size = new_dataset_conf.get('test_size') labels_frame = self.generate_new_frame(new_dataset_conf) save_tfr( labels_frame, get_abs_path('data', 'tfrecords', create_parents=True), new_dataset_conf['dataset_name'], test_size, self, )
def adjust_frame(frame, cache_file=None): """ Add relative width, relative height and object ids to annotation pandas DataFrame. Args: frame: pandas DataFrame containing coordinates instead of relative labels. cache_file: cache_file: csv file name containing current session labels. Returns: Frame with the new columns """ object_id = 1 for item in frame.columns[2:]: frame[item] = frame[item].astype(float).astype(int) frame['relative_width'] = (frame['x_max'] - frame['x_min']) / frame['img_width'] frame['relative_height'] = (frame['y_max'] - frame['y_min']) / frame['img_height'] for object_name in list(frame['object_name'].drop_duplicates()): frame.loc[frame['object_name'] == object_name, 'object_id'] = object_id object_id += 1 if cache_file: frame.to_csv(get_abs_path('output', 'data', cache_file, create_parents=True), index=False) LOGGER.info(f'Parsed labels:\n{frame["object_name"].value_counts()}') return frame
def adjust_non_voc_csv(csv_file, image_path, image_width, image_height): """ Read relative data and return adjusted frame accordingly. Args: csv_file: .csv file containing the following columns: [image, object_name, object_index, bx, by, bw, bh] image_path: Path prefix to be added. image_width: image width. image_height: image height Returns: pandas DataFrame with the following columns: ['image_path', 'object_name', 'img_width', 'img_height', 'x_min', 'y_min', 'x_max', 'y_max', 'relative_width', 'relative_height', 'object_id'] """ image_path = get_abs_path(image_path, verify=True) coordinates = [] old_frame = pd.read_csv(get_abs_path(csv_file, verify=True)) new_frame = pd.DataFrame() new_frame['image_path'] = old_frame['image'].apply( lambda item: get_abs_path(image_path, item)) new_frame['object_name'] = old_frame['object_name'] new_frame['img_width'] = image_width new_frame['img_height'] = image_height new_frame['relative_width'] = old_frame['bw'] new_frame['relative_height'] = old_frame['bh'] new_frame['object_id'] = old_frame['object_index'] + 1 for index, row in old_frame.iterrows(): image, object_name, object_index, bx, by, bw, bh = row co = ratios_to_coordinates(bx, by, bw, bh, image_width, image_height) coordinates.append(co) ( new_frame['x_min'], new_frame['y_min'], new_frame['x_max'], new_frame['y_max'], ) = np.array(coordinates).T new_frame[['x_min', 'y_min', 'x_max', 'y_max']] = new_frame[['x_min', 'y_min', 'x_max', 'y_max']].astype('int64') print(f'Parsed labels:\n{new_frame["object_name"].value_counts()}') classes = new_frame['object_name'].drop_duplicates() LOGGER.info( f'Adjustment from existing received {len(new_frame)} labels containing ' f'{len(classes)} classes') LOGGER.info(f'Added prefix to images: {image_path}') return new_frame[[ 'image_path', 'object_name', 'img_width', 'img_height', 'x_min', 'y_min', 'x_max', 'y_max', 'relative_width', 'relative_height', 'object_id', ]]
def predict_photos(self, photos, trained_weights, batch_size=32, workers=16, output_dir=None): """ Predict a list of image paths and save results to output folder. Args: photos: A list of image paths. trained_weights: .weights or .tf file batch_size: Prediction batch size. workers: Parallel predictions. output_dir: Path to output dir, defaults to output/detections Returns: None """ self.create_models( reverse_v4=True if trained_weights.endswith('tf') else False) self.load_weights(get_abs_path(trained_weights, verify=True)) to_predict = photos.copy() saved_paths = set() with ThreadPoolExecutor(max_workers=workers) as executor: predicted = 1 total_photos = len(photos) while to_predict: current_batch = [ to_predict.pop() for _ in range(min(batch_size, len(to_predict))) if to_predict ] future_predictions = { executor.submit( self.predict_on_image, image, output_dir, ): image for image in current_batch } for future_prediction in as_completed(future_predictions): saved_path = future_prediction.result() saved_paths.add(saved_path) completed = f'{predicted}/{total_photos}' current_image = future_predictions[future_prediction] percent = (predicted / total_photos) * 100 print( f'\rpredicting {os.path.basename(current_image)} ' f'{completed}\t{percent}% completed', end='', ) predicted += 1 print() for saved_path in saved_paths: LOGGER.info(f'Saved prediction: {saved_path}')
def detect_video(self, video, trained_weights, codec='mp4v', display=False, output_dir=None): """ Perform detection on a video, stream(optional) and save results. Args: video: Path to video file. trained_weights: .tf or .weights file codec: str ex: mp4v display: If True, detections will be displayed during the detection operation. output_dir: Path to output dir, defaults to output/detections Returns: None """ self.create_models( reverse_v4=True if trained_weights.endswith('tf') else False) self.load_weights(trained_weights) vid = cv2.VideoCapture(video) length = int(vid.get(cv2.CAP_PROP_FRAME_COUNT)) width = int(vid.get(cv2.CAP_PROP_FRAME_WIDTH)) height = int(vid.get(cv2.CAP_PROP_FRAME_HEIGHT)) fps = int(vid.get(cv2.CAP_PROP_FPS)) current = 1 codec = cv2.VideoWriter_fourcc(*codec) out = (get_abs_path( output_dir, 'predicted_vid.mp4', create_parents=True) if output_dir else get_abs_path('output', 'detections', 'predicted_vid.mp4', create_parents=True)) writer = cv2.VideoWriter(out, codec, fps, (width, height)) while vid.isOpened(): _, frame = vid.read() detections, adjusted = self.detect_image(frame, f'frame_{current}') self.draw_on_image(adjusted, detections) writer.write(adjusted) completed = f'{(current / length) * 100}% completed' print( f'\rframe {current}/{length}\tdetections: ' f'{len(detections)}\tcompleted: {completed}', end='', ) if display: cv2.destroyAllWindows() cv2.imshow(f'frame {current}', adjusted) current += 1 if cv2.waitKey(1) == ord('q'): LOGGER.info(f'Video detection aborted {current}/{length} ' f'frames completed') break
def create_models(self, reverse_v4=False): """ Create training and inference yolo models. Returns: training, inference models """ input_initial = self.apply_func(Input, shape=self.input_shape) cfg_out = self.read_dark_net_cfg() cfg_parser = configparser.ConfigParser() cfg_parser.read_file(cfg_out) self.output_indices = [] self.previous_layer = input_initial for section in cfg_parser.sections(): self.create_section(section, cfg_parser) if len(self.output_indices) == 0: self.output_indices.append(len(self.model_layers) - 1) self.output_layers.extend( [self.model_layers[i] for i in self.output_indices]) if '4' in self.model_configuration and reverse_v4: self.output_layers.reverse() self.training_model = Model(inputs=input_initial, outputs=self.output_layers) output_0, output_1, output_2 = self.output_layers boxes_0 = self.apply_func( Lambda, output_0, lambda item: get_boxes(item, self.anchors[self.masks[0]], self. classes), ) boxes_1 = self.apply_func( Lambda, output_1, lambda item: get_boxes(item, self.anchors[self.masks[1]], self. classes), ) boxes_2 = self.apply_func( Lambda, output_2, lambda item: get_boxes(item, self.anchors[self.masks[2]], self. classes), ) outputs = self.apply_func( Lambda, (boxes_0[:3], boxes_1[:3], boxes_2[:3]), lambda item: self.get_nms(item), ) self.inference_model = Model(input_initial, outputs, name='inference_model') LOGGER.info('Training and inference models created') return self.training_model, self.inference_model
def relative_to_coordinates(self, out_file=None): """ Convert relative coordinates in self.mapping to coordinates. Args: out_file: path to new converted csv. Returns: pandas DataFrame with the new coordinates. """ items_to_save = [] for index, data in self.mapping.iterrows(): image_name, object_name, object_index, bx, by, bw, bh = data x1, y1, x2, y2 = ratios_to_coordinates(bx, by, bw, bh, self.image_width, self.image_height) items_to_save.append([ image_name, x1, y1, x2, y2, object_name, object_index, bx, by, bw, bh, ]) new_data = pd.DataFrame( items_to_save, columns=[ 'image', 'x1', 'y1', 'x2', 'y2', 'object_type', 'object_id', 'bx', 'by', 'bw', 'bh', ], ) new_data[['x1', 'y1', 'x2', 'y2']] = new_data[['x1', 'y1', 'x2', 'y2']].astype('int64') if out_file: new_data.to_csv(out_file, index=False) LOGGER.info(f'Converted labels in {self.labels_file} to coordinates') return new_data
def save_tfr(data, output_folder, dataset_name, test_size, trainer=None): """ Transform and save dataset into TFRecord format. Args: data: pandas DataFrame with adjusted labels. output_folder: Path to folder where TFRecord(s) will be saved. dataset_name: str name of the dataset. test_size: relative test subset size. trainer: core.Trainer object Returns: None """ assert (0 < test_size < 1), f'test_size must be 0 < test_size < 1 and {test_size} is given' data['object_name'] = data['object_name'].apply( lambda x: x.encode('utf-8')) data['object_id'] = data['object_id'].astype(int) data[data.dtypes[data.dtypes == 'int64'].index] = data[data.dtypes[ data.dtypes == 'int64'].index].apply(abs) data.to_csv( get_abs_path('data', 'tfrecords', 'full_data.csv', create_parents=True), index=False, ) groups = np.array(data.groupby('image_path')) np.random.shuffle(groups) separation_index = int((1 - test_size) * len(groups)) training_set = groups[:separation_index] test_set = groups[separation_index:] training_frame = pd.concat([item[1] for item in training_set]) test_frame = pd.concat([item[1] for item in test_set]) training_frame.to_csv( get_abs_path('data', 'tfrecords', 'training_data.csv', create_parents=True), index=False, ) test_frame.to_csv( get_abs_path('data', 'tfrecords', 'test_data.csv', create_parents=True), index=False, ) training_path = get_abs_path(output_folder, f'{dataset_name}_train.tfrecord') test_path = get_abs_path(output_folder, f'{dataset_name}_test.tfrecord') write_tf_record(training_path, training_set, data, trainer) LOGGER.info(f'Saved training TFRecord: {training_path}') write_tf_record(test_path, test_set, data, trainer) LOGGER.info(f'Saved validation TFRecord: {test_path}')
def predict_photos( self, photos, trained_weights, batch_size=32, workers=16, output_dir=None ): """ Predict a list of image paths and save results to output folder. Args: photos: A list of image paths. trained_weights: .weights or .tf file batch_size: Prediction batch size. workers: Parallel predictions. output_dir: Path to output dir, defaults to output/detections Returns: None """ self.create_models() self.load_weights(trained_weights) to_predict = photos.copy() with ThreadPoolExecutor(max_workers=workers) as executor: predicted = 1 done = [] total_photos = len(photos) while to_predict: current_batch = [ to_predict.pop() for _ in range(batch_size) if to_predict ] future_predictions = { executor.submit( self.predict_on_image, image, Path(output_dir, image).absolute().resolve() if output_dir else None, ): image for image in current_batch } for future_prediction in as_completed(future_predictions): future_prediction.result() completed = f'{predicted}/{total_photos}' current_image = future_predictions[future_prediction] percent = (predicted / total_photos) * 100 print( f'\rpredicting {os.path.basename(current_image)} ' f'{completed}\t{percent}% completed', end='', ) predicted += 1 done.append(current_image) for item in done: LOGGER.info(f'Saved prediction: {item}')
def write_tf_record(output_path, groups, data, trainer=None): """ Write data to TFRecord. Args: output_path: Full path to save. groups: pandas GroupBy object. data: pandas DataFrame trainer: core.Trainer object. Returns: None """ print(f'Processing {os.path.split(output_path)[-1]}') if trainer: if 'train' in output_path: trainer.train_tf_record = output_path if 'test' in output_path: trainer.valid_tf_record = output_path with tf.io.TFRecordWriter(output_path) as r_writer: for current_image, (image_path, objects) in enumerate(groups, 1): print( f'\rBuilding example: {current_image}/{len(groups)} ... ' f'{os.path.split(image_path)[-1]} ' f'{round(100 * (current_image / len(groups)))}% completed', end='', ) separate_data = pd.DataFrame(objects, columns=data.columns).T.to_numpy() ( image_width, image_height, x_min, y_min, x_max, y_max, ) = separate_data[2:8] x_min /= image_width x_max /= image_width y_min /= image_height y_max /= image_height try: image_data = open(image_path, 'rb').read() key = hashlib.sha256(image_data).hexdigest() training_example = create_example(separate_data, key, image_data) r_writer.write(training_example.SerializeToString()) except Exception as e: LOGGER.error(e) print()
def __init__( self, labels_file, augmentation_map, workers=32, converted_coordinates_file=None, image_folder=None, ): """ Initialize augmentation session. Args: labels_file: cvv file containing relative image labels augmentation_map: A structured dictionary containing categorized augmentation sequences. workers: Parallel threads. converted_coordinates_file: csv file containing converted from relative to coordinates. image_folder: Folder containing images other than data/photos/ """ assert all([ia, iaa, iap]) self.labels_file = labels_file self.mapping = pd.read_csv(labels_file) self.image_folder = Path('data', 'photos').absolute().resolve() if image_folder: self.image_folder = Path(image_folder).absolute().resolve() self.image_paths = [ (Path(self.image_folder) / image).absolute().resolve() for image in os.listdir(self.image_folder) if not image.startswith('.') ] self.image_paths_copy = self.image_paths.copy() if not self.image_paths: LOGGER.error( f'Augmentation aborted: no photos found in {self.image_folder}' ) raise ValueError(f'No photos given') self.image_width, self.image_height = imagesize.get( self.image_paths[0]) self.converted_coordinates = (pd.read_csv(converted_coordinates_file) if converted_coordinates_file else self.relative_to_coordinates()) self.converted_groups = self.converted_coordinates.groupby('image') self.augmentation_data = [] self.augmentation_sequences = [] self.augmentation_map = augmentation_map self.workers = workers self.augmented_images = 0 self.total_images = len(self.image_paths) self.session_id = np.random.randint(10**6, (10**7))
def check_tf_records(self): """ Ensure tfrecords are specified to start training. Returns: None """ if not self.train_tf_record: issue = 'No training TFRecord specified' LOGGER.error(issue) raise ValueError(issue) if not self.valid_tf_record: issue = 'No validation TFRecord specified' LOGGER.error(issue) raise ValueError(issue)
def parse_voc_folder(folder_path, voc_conf): """ Parse a folder containing voc xml annotation files. Args: folder_path: Folder containing voc xml annotation files. voc_conf: Path to voc json configuration file. Returns: pandas DataFrame with the annotations. """ assert os.path.exists(folder_path) cache_path = os.path.join('output', 'data', 'parsed_from_xml.csv') if os.path.exists(cache_path): frame = pd.read_csv(cache_path) print(f'Labels retrieved from cache:' f'\n{frame["Object Name"].value_counts()}') return frame image_data = [] frame_columns = [ 'Image Path', 'Object Name', 'Image Width', 'Image Height', 'X_min', 'Y_min', 'X_max', 'Y_max', ] xml_files = [ file_name for file_name in os.listdir(folder_path) if file_name.endswith('.xml') ] for file_name in xml_files: annotation_path = os.path.join(folder_path, file_name) image_labels = parse_voc_file(annotation_path, voc_conf) image_data.extend(image_labels) frame = pd.DataFrame(image_data, columns=frame_columns) classes = frame['Object Name'].drop_duplicates() LOGGER.info(f'Read {len(xml_files)} xml files') LOGGER.info(f'Received {len(frame)} labels containing ' f'{len(classes)} classes') if frame.empty: raise ValueError( f'No labels were found in {os.path.abspath(folder_path)}') frame = adjust_frame(frame, 'parsed_from_xml.csv') return frame
def clear_outputs(): """ Clear output folder. Returns: None """ for folder_name in os.listdir(get_abs_path('output', verify=True)): if not folder_name.startswith('.'): full_path = get_abs_path('output', folder_name) for file_name in os.listdir(full_path): full_file_path = get_abs_path(full_path, file_name) if os.path.isdir(full_file_path): shutil.rmtree(full_file_path) else: os.remove(full_file_path) LOGGER.info(f'Deleted old output: {full_file_path}')
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 clear_outputs(): """ Clear output folder. Returns: None """ for folder_name in os.listdir(os.path.join('..', 'output')): if not folder_name.startswith('.'): full_path = (Path(os.path.join( 'output', folder_name)).absolute().resolve()) for file_name in os.listdir(full_path): full_file_path = Path(os.path.join(full_path, file_name)) if os.path.isdir(full_file_path): shutil.rmtree(full_file_path) else: os.remove(full_file_path) LOGGER.info(f'Deleted old output: {full_file_path}')
def save_fig(title, save_figures=True): """ Save generated figures to output folder. Args: title: Figure title also the image to save file name. save_figures: If True, figure will be saved Returns: None """ if save_figures: saving_path = get_abs_path( 'output', 'plots', f'{title}.png', create_parents=True ) if os.path.exists(saving_path): return plt.savefig(saving_path) LOGGER.info(f'Saved figure {saving_path}') plt.close()
def save_fig(title, save_figures=True): """ Save generated figures to output folder. Args: title: Figure title also the image to save file name. save_figures: If True, figure will be saved Returns: None """ if save_figures: saving_path = str( Path(os.path.join('output', 'plots', f'{title}.png')).absolute().resolve() ) if os.path.exists(saving_path): return plt.savefig(saving_path) LOGGER.info(f'Saved figure {saving_path}') plt.close()
def create_sequences(self, sequences): """ Create sequences for imgaug.augmenters.Sequential(). Args: sequences: A list of dictionaries with each dictionary containing the following: -sequence_group: str, one of self.augmentation_map keys including: ['meta', 'arithmetic', 'artistic', 'blend', 'gaussian_blur', 'color', 'contrast', 'convolution', 'edges', 'flip', 'geometric', 'corrupt_like', 'pi_like', 'pooling', 'segmentation', 'size'] -no: augmentation number ('no' key in self.augmentation_map > chosen sequence_group) Example: sequences = ( [[{'sequence_group': 'meta', 'no': 5}, {'sequence_group': 'arithmetic', 'no': 3}], [{'sequence_group': 'arithmetic', 'no': 2}]] ) Returns: The list of augmentation sequences that will be applied over images. """ total_target = (len(sequences) * len(self.image_paths)) + len( self.image_paths) LOGGER.info(f'Total images(old + augmented): {total_target}') for group in sequences: some_ofs = [ self.augmentation_map[item['sequence_group']][item['no'] - 1] for item in group[0] ] one_ofs = [ self.augmentation_map[item['sequence_group']][item['no'] - 1] for item in group[1] ] some_of_aug = [item['augmentation'] for item in some_ofs] one_of_aug = [item['augmentation'] for item in one_ofs] some_of_seq = iaa.SomeOf((0, 5), [eval(item) for item in some_of_aug]) one_of_seq = iaa.OneOf([eval(item) for item in one_of_aug]) self.augmentation_sequences.append( iaa.Sequential([some_of_seq, one_of_seq], random_order=True)) return self.augmentation_sequences
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)
def read_tfr( tf_record_file, classes_file, feature_map, max_boxes, classes_delimiter='\n', new_size=None, get_features=False, ): """ Read and load dataset from TFRecord file. Args: tf_record_file: Path to TFRecord file. classes_file: file containing classes. feature_map: A dictionary of feature names mapped to tf.io objects. max_boxes: Maximum number of boxes per image. classes_delimiter: delimiter in classes_file. new_size: w, h new image size get_features: If True, features will be returned. Returns: MapDataset object. """ tf_record_file = str(Path(tf_record_file).absolute().resolve().as_posix()) text_init = tf.lookup.TextFileInitializer(classes_file, tf.string, 0, tf.int64, -1, delimiter=classes_delimiter) class_table = tf.lookup.StaticHashTable(text_init, -1) files = tf.data.Dataset.list_files(tf_record_file) dataset = files.flat_map(tf.data.TFRecordDataset) LOGGER.info(f'Read TFRecord: {tf_record_file}') return dataset.map(lambda x: read_example( x, feature_map, class_table, max_boxes, new_size, get_features))
def load_weights(self, weights_file): """ Load DarkNet weights or checkpoint/pre-trained weights. Args: weights_file: .weights or .tf file path. Returns: None """ assert (suffix := Path(weights_file).suffix) in [ '.tf', '.weights', ], 'Invalid weights file' assert ( self.classes == 80 if suffix == '.weights' else 1 ), f'DarkNet model should contain 80 classes, {self.classes} is given.' if suffix == '.tf': self.training_model.load_weights(get_abs_path(weights_file)) LOGGER.info(f'Loaded weights: {weights_file} ... success') return with open(get_abs_path(weights_file, verify=True), 'rb') as weights_data: LOGGER.info(f'Loading pre-trained weights ...') major, minor, revision, seen, _ = np.fromfile(weights_data, dtype=np.int32, count=5) self.model_layers = [ layer for layer in self.training_model.layers if id(layer) not in [id(item) for item in self.output_layers] ] self.model_layers.sort( key=lambda layer: int(layer.name.split('_')[1])) self.model_layers.extend(self.output_layers) for i, layer in enumerate(self.model_layers): current_read = weights_data.tell() total_size = os.fstat(weights_data.fileno()).st_size if current_read == total_size: print() break print( f'\r{round(100 * (current_read / total_size))}' f'%\t{current_read}/{total_size}', end='', ) if 'conv2d' not in layer.name: continue next_layer = self.model_layers[i + 1] b_norm_layer = (next_layer if 'batch_normalization' in next_layer.name else None) filters = layer.filters kernel_size = layer.kernel_size[0] input_dimension = layer.get_input_shape_at(-1)[-1] convolution_bias = (np.fromfile( weights_data, dtype=np.float32, count=filters) if b_norm_layer is None else None) bn_weights = (np.fromfile( weights_data, dtype=np.float32, count=4 * filters).reshape( (4, filters))[[1, 0, 2, 3]] if (b_norm_layer is not None) else None) convolution_shape = ( filters, input_dimension, kernel_size, kernel_size, ) convolution_weights = (np.fromfile( weights_data, dtype=np.float32, count=np.product(convolution_shape), ).reshape(convolution_shape).transpose([2, 3, 1, 0])) if b_norm_layer is None: try: layer.set_weights( [convolution_weights, convolution_bias]) except ValueError: pass if b_norm_layer is not None: layer.set_weights([convolution_weights]) b_norm_layer.set_weights(bn_weights) assert len(weights_data.read()) == 0, 'failed to read all data' LOGGER.info(f'\nLoaded weights: {weights_file} ... success')
def get_dataset_next(dataset): try: return next(dataset) except tf.errors.UnknownError as e: LOGGER.error(f'Error occurred during reading from dataset\n{e}')
def augment_photos_folder(self, batch_size=64, new_size=None): """ Augment photos in data/photos/ Args: batch_size: Size of each augmentation batch. new_size: tuple, new image size. Returns: None """ LOGGER.info(f'Started augmentation with {self.workers} workers') LOGGER.info(f'Total images to augment: {self.total_images}') LOGGER.info(f'Session assigned id: {self.session_id}') with ThreadPoolExecutor(max_workers=self.workers) as executor: while self.image_paths_copy: current_batch, current_paths = self.load_batch( new_size, batch_size) future_augmentations = { executor.submit(self.augment_image, image, path): path for image, path in zip(current_batch, current_paths) } for future_augmented in as_completed(future_augmentations): future_augmented.result() LOGGER.info(f'Augmentation completed') augmentation_frame = pd.DataFrame(self.augmentation_data, columns=self.mapping.columns) saving_path = os.path.join('output', 'data', f'augmented_data_plus_original.csv') combined = pd.concat([self.mapping, augmentation_frame]) for item in ['bx', 'by', 'bw', 'bh']: combined = combined.drop(combined[combined[item] > 1].index) combined.to_csv(saving_path, index=False) LOGGER.info(f'Saved old + augmented labels to {saving_path}') adjusted_combined = adjust_non_voc_csv(saving_path, self.image_folder, self.image_width, self.image_height) adjusted_saving_path = saving_path.replace('augmented', 'adjusted_aug') adjusted_combined.to_csv(adjusted_saving_path, index=False) LOGGER.info( f'Saved old + augmented (adjusted) labels to {adjusted_saving_path}' ) return adjusted_combined
def get_dataset_next(dataset): try: return next(dataset) except tf.errors.UnknownError as e: # sometimes encountered when reading from google drive LOGGER.error(f'Error occurred during reading from dataset\n{e}')
def train( self, epochs, batch_size, learning_rate, new_anchors_conf=None, new_dataset_conf=None, dataset_name=None, weights=None, evaluate=True, merge_evaluation=True, evaluation_workers=8, shuffle_buffer=512, min_overlaps=None, display_stats=True, plot_stats=True, save_figs=True, clear_outputs=False, n_epoch_eval=None, ): """ Train on the dataset. Args: epochs: Number of training epochs. batch_size: Training batch size. learning_rate: non-negative value. new_anchors_conf: A dictionary containing anchor generation configuration. new_dataset_conf: A dictionary containing dataset generation configuration. dataset_name: Name of the dataset for model checkpoints. weights: .tf or .weights file evaluate: If False, the trained model will not be evaluated after training. merge_evaluation: If False, training and validation maps will be calculated separately. evaluation_workers: Parallel predictions. shuffle_buffer: Buffer size for shuffling datasets. min_overlaps: a float value between 0 and 1, or a dictionary containing each class in self.class_names mapped to its minimum overlap display_stats: If True and evaluate=True, evaluation statistics will be displayed. plot_stats: If True, Precision and recall curves as well as comparative bar charts will be plotted save_figs: If True and plot_stats=True, figures will be saved clear_outputs: If True, old outputs will be cleared n_epoch_eval: Conduct evaluation every n epoch. Returns: history object, pandas DataFrame with statistics, mAP score. """ min_overlaps = min_overlaps or 0.5 if clear_outputs: self.clear_outputs() activate_gpu() LOGGER.info(f'Starting training ...') if new_anchors_conf: LOGGER.info(f'Generating new anchors ...') self.generate_new_anchors(new_anchors_conf) self.create_models(reverse_v4=True) if weights: self.load_weights(weights) if new_dataset_conf: self.create_new_dataset(new_dataset_conf) self.check_tf_records() training_dataset = self.initialize_dataset(self.train_tf_record, batch_size, shuffle_buffer) valid_dataset = self.initialize_dataset(self.valid_tf_record, batch_size, shuffle_buffer) optimizer = tf.keras.optimizers.Adam(learning_rate) loss = [ calculate_loss(self.anchors[mask], self.classes, self.iou_threshold) for mask in self.masks ] self.training_model.compile(optimizer=optimizer, loss=loss) checkpoint_path = get_abs_path( 'models', f'{dataset_name or "trained"}_model.tf') callbacks = self.create_callbacks(checkpoint_path) if n_epoch_eval: mid_train_eval = MidTrainingEvaluator( self.input_shape, self.model_configuration, self.classes_file, self.train_tf_record, self.valid_tf_record, self.anchors, self.masks, self.max_boxes, self.iou_threshold, self.score_threshold, n_epoch_eval, merge_evaluation, evaluation_workers, shuffle_buffer, min_overlaps, display_stats, plot_stats, save_figs, checkpoint_path, self.image_folder, ) callbacks.append(mid_train_eval) history = self.training_model.fit( training_dataset, epochs=epochs, callbacks=callbacks, validation_data=valid_dataset, ) LOGGER.info('Training complete') if evaluate: evaluations = self.evaluate( checkpoint_path, merge_evaluation, evaluation_workers, shuffle_buffer, min_overlaps, display_stats, plot_stats, save_figs, ) return evaluations, history return history
def evaluate( self, weights_file, merge, workers, shuffle_buffer, min_overlaps, display_stats=True, plot_stats=True, save_figs=True, ): """ Evaluate on training and validation datasets. Args: weights_file: Path to trained .tf file. merge: If False, training and validation datasets will be evaluated separately. workers: Parallel predictions. shuffle_buffer: Buffer size for shuffling datasets. min_overlaps: a float value between 0 and 1, or a dictionary containing each class in self.class_names mapped to its minimum overlap display_stats: If True evaluation statistics will be printed. plot_stats: If True, evaluation statistics will be plotted including precision and recall curves and mAP save_figs: If True, resulting plots will be save to output folder. Returns: stats, map_score. """ LOGGER.info('Starting evaluation ...') evaluator = Evaluator( self.input_shape, self.model_configuration, self.train_tf_record, self.valid_tf_record, self.classes_file, self.anchors, self.masks, self.max_boxes, self.iou_threshold, self.score_threshold, ) predictions = evaluator.make_predictions(weights_file, merge, workers, shuffle_buffer) if isinstance(predictions, tuple): training_predictions, valid_predictions = predictions if any([training_predictions.empty, valid_predictions.empty]): LOGGER.info('Aborting evaluations, no detections found') return training_actual = pd.read_csv( get_abs_path('data', 'tfrecords', 'training_data.csv', verify=True)) valid_actual = pd.read_csv( get_abs_path('data', 'tfrecords', 'test_data.csv', verify=True)) training_stats, training_map = evaluator.calculate_map( training_predictions, training_actual, min_overlaps, display_stats, 'Train', save_figs, plot_stats, ) valid_stats, valid_map = evaluator.calculate_map( valid_predictions, valid_actual, min_overlaps, display_stats, 'Valid', save_figs, plot_stats, ) return training_stats, training_map, valid_stats, valid_map actual_data = pd.read_csv( get_abs_path('data', 'tfrecords', 'full_data.csv', verify=True)) if predictions.empty: LOGGER.info('Aborting evaluations, no detections found') return stats, map_score = evaluator.calculate_map( predictions, actual_data, min_overlaps, display_stats, save_figs=save_figs, plot_results=plot_stats, ) return stats, map_score