From ce1f9080eef7d86242f62bf1104907f746764d48 Mon Sep 17 00:00:00 2001 From: KCFindstr Date: Mon, 7 Aug 2023 10:21:11 -0700 Subject: [PATCH] Add data converter and movinet code to model garden (#2152) * Add model garden data converter and movinet code. * Add movinet and data converter CODEOWNERS. --- community-content/CODEOWNERS | 2 + .../model_oss/data_converter/common_lib.py | 623 ++++++++++++++++++ .../convert_coco_label_map_iod.py | 111 ++++ .../convert_imagenet_label_map.py | 86 +++ .../data_converter/data_converter_icn_lib.py | 199 ++++++ .../data_converter/data_converter_iod_lib.py | 430 ++++++++++++ .../data_converter/data_converter_isg_lib.py | 328 +++++++++ .../data_converter/data_converter_main.py | 166 +++++ .../data_converter/data_converter_vcn_lib.py | 289 ++++++++ .../data_converter/dockerfile/base.Dockerfile | 50 ++ .../data_converter/dockerfile/run.Dockerfile | 23 + .../movinet/dockerfile/base.Dockerfile | 64 ++ .../dockerfile/model_export.Dockerfile | 13 + .../movinet/dockerfile/serve.Dockerfile | 18 + .../movinet/dockerfile/train.Dockerfile | 18 + .../model_oss/movinet/serving/serving_main.py | 142 ++++ .../movinet/serving/video_serving_lib.py | 462 +++++++++++++ .../model_oss/movinet/train.py | 210 ++++++ .../model_oss/pic2word/handler.py | 2 +- .../model_oss/util/commons.py | 79 +++ .../model_oss/util/constants.py | 9 + .../model_oss/util/fileutils.py | 169 ++++- .../model_oss/util/hypertune_utils.py | 22 + 23 files changed, 3513 insertions(+), 2 deletions(-) create mode 100644 community-content/vertex_vision_model_garden/model_oss/data_converter/common_lib.py create mode 100644 community-content/vertex_vision_model_garden/model_oss/data_converter/convert_coco_label_map_iod.py create mode 100644 community-content/vertex_vision_model_garden/model_oss/data_converter/convert_imagenet_label_map.py create mode 100644 community-content/vertex_vision_model_garden/model_oss/data_converter/data_converter_icn_lib.py create mode 100644 community-content/vertex_vision_model_garden/model_oss/data_converter/data_converter_iod_lib.py create mode 100644 community-content/vertex_vision_model_garden/model_oss/data_converter/data_converter_isg_lib.py create mode 100644 community-content/vertex_vision_model_garden/model_oss/data_converter/data_converter_main.py create mode 100644 community-content/vertex_vision_model_garden/model_oss/data_converter/data_converter_vcn_lib.py create mode 100644 community-content/vertex_vision_model_garden/model_oss/data_converter/dockerfile/base.Dockerfile create mode 100644 community-content/vertex_vision_model_garden/model_oss/data_converter/dockerfile/run.Dockerfile create mode 100644 community-content/vertex_vision_model_garden/model_oss/movinet/dockerfile/base.Dockerfile create mode 100644 community-content/vertex_vision_model_garden/model_oss/movinet/dockerfile/model_export.Dockerfile create mode 100644 community-content/vertex_vision_model_garden/model_oss/movinet/dockerfile/serve.Dockerfile create mode 100644 community-content/vertex_vision_model_garden/model_oss/movinet/dockerfile/train.Dockerfile create mode 100644 community-content/vertex_vision_model_garden/model_oss/movinet/serving/serving_main.py create mode 100644 community-content/vertex_vision_model_garden/model_oss/movinet/serving/video_serving_lib.py create mode 100644 community-content/vertex_vision_model_garden/model_oss/movinet/train.py create mode 100644 community-content/vertex_vision_model_garden/model_oss/util/commons.py create mode 100644 community-content/vertex_vision_model_garden/model_oss/util/hypertune_utils.py diff --git a/community-content/CODEOWNERS b/community-content/CODEOWNERS index 24447ed08..80ce20978 100644 --- a/community-content/CODEOWNERS +++ b/community-content/CODEOWNERS @@ -15,3 +15,5 @@ /vertex_vision_model_garden/model_oss/keras @dstnluong-google /vertex_vision_model_garden/model_oss/transformers @dstnluong-google /vertex_vision_model_garden/model_oss/pic2word @jismailyan-google +/vertex_vision_model_garden/model_oss/movinet @KCFindstr +/vertex_vision_model_garden/model_oss/data_converter @KCFindstr diff --git a/community-content/vertex_vision_model_garden/model_oss/data_converter/common_lib.py b/community-content/vertex_vision_model_garden/model_oss/data_converter/common_lib.py new file mode 100644 index 000000000..567e6c535 --- /dev/null +++ b/community-content/vertex_vision_model_garden/model_oss/data_converter/common_lib.py @@ -0,0 +1,623 @@ +"""Library with functions to use for data conversion.""" + +import json +import os +import random +from typing import Any, Callable, Dict, Iterable, List, Optional, Sequence, Tuple, Union +import uuid + +from absl import logging +import apache_beam as beam +import cv2 +import numpy as np +import pandas as pd +import PIL +from PIL import Image +import tensorflow as tf +import yaml + +from util import constants +from util import fileutils +from apache_beam.options import pipeline_options + +REFORMATTED_CSV_SUFFIX = '-reformatted.csv' + +LABEL_MAP_NAME = 'label_map.yaml' + +_SPLIT_RATIO_ERROR_THRESHOLD = 1e-5 +# Internal constant. Only for distinguishing rows without ML use. +ML_USE_UNASSIGNED = 'unassigned' +ALL_ML_USES = ( + constants.ML_USE_TRAINING, + constants.ML_USE_VALIDATION, + constants.ML_USE_TEST, + ML_USE_UNASSIGNED, +) +COLUMN_NAME_ML_USE = 'ml_use' +COLUMN_NAME_GCS_FILE_PATH = 'gcs_file_path' +COLUMN_NAME_LABEL = 'label' +COLUMN_NAME_START_SEC = 'start_sec' +COLUMN_NAME_END_SEC = 'end_sec' +# Output filenames +TRAIN_TFRECORD_NAME = 'train.tfrecord' +VALIDATION_TFRECORD_NAME = 'val.tfrecord' +TEST_TFRECORD_NAME = 'test.tfrecord' +# Jsonl keys +JSON_GCS_URI_KEY = 'imageGcsUri' +JSON_RESOURCE_LABEL_KEY = 'dataItemResourceLabels' +JSON_ML_USE_KEY = 'aiplatform.googleapis.com/ml_use' +# I/O parameters +READ_CHUNK_SIZE = 1024 * 1024 * 1024 # 1GB + + +class WriteToTFRecord(beam.DoFn): + """DoFn to write TF examples to sharded TF record files.""" + + def __init__( + self, + output_prefix: str, + num_shards: int, + convert_fn: Callable[[Dict[str, Any]], tf.train.Example], + ): + self.output_prefix = output_prefix + self.num_shards = num_shards + self.writer: list[tf.io.TFRecordWriter] = [] + self.sharded_files: list[str] = [] + self.convert_fn = convert_fn + self.success_counter = beam.metrics.Metrics.counter( + self.__class__.__name__, 'Success' + ) + self.failure_counter = beam.metrics.Metrics.counter( + self.__class__.__name__, 'Failure' + ) + + def start_bundle(self): + logging.info('Start writing TF Record to %s.', self.output_prefix) + unique_str = uuid.uuid4().hex + for i in range(self.num_shards): + uri = f'{self.output_prefix}-{i}-{unique_str}' + self.sharded_files.append(uri) + self.writer.append(tf.io.TFRecordWriter(uri)) + + def process(self, data: Dict[str, Any]) -> Iterable[Tuple[int, str]]: + try: + example = self.convert_fn(data) + data = example.SerializeToString() + idx = hash(data) % self.num_shards + self.writer[idx].write(data) + self.success_counter.inc() + yield (idx, self.sharded_files[idx]) + # pylint: disable-next=broad-exception-caught + except Exception as err: + logging.error('Failed to process %s', data) + logging.exception(err) + self.failure_counter.inc() + + def finish_bundle(self): + logging.info('Finish writing TF Record to %s.', self.output_prefix) + for writer in self.writer: + writer.close() + self.writer = [] + + +def convert_to_feature( + value: Union[List[Union[int, float, bytes]], int, float, bytes], + value_type: Optional[str] = None, +) -> tf.train.Feature: + """Converts the given python object to a tf.train.Feature. + + This is copied from tensorflow_models/official/vision/data/tfrecord_lib.py. + + Args: + value: int, float, bytes or a list of them. + value_type: optional, if specified, forces the feature to be of the given + type. Otherwise, type is inferred automatically. Can be one of ['bytes', + 'int64', 'float', 'bytes_list', 'int64_list', 'float_list'] + + Returns: + feature: A tf.train.Feature object. + """ + + if value_type is None: + element = value[0] if isinstance(value, list) else value + + if isinstance(element, bytes): + value_type = 'bytes' + + elif isinstance(element, (int, np.integer)): + value_type = 'int64' + + elif isinstance(element, (float, np.floating)): + value_type = 'float' + + else: + raise ValueError( + 'Cannot convert type {} to feature'.format(type(element)) + ) + + if isinstance(value, list): + value_type = value_type + '_list' + + if value_type == 'int64': + return tf.train.Feature(int64_list=tf.train.Int64List(value=[value])) + + elif value_type == 'int64_list': + value = np.asarray(value).astype(np.int64).reshape(-1) + return tf.train.Feature(int64_list=tf.train.Int64List(value=value)) + + elif value_type == 'float': + return tf.train.Feature(float_list=tf.train.FloatList(value=[value])) + + elif value_type == 'float_list': + value = np.asarray(value).astype(np.float32).reshape(-1) + return tf.train.Feature(float_list=tf.train.FloatList(value=value)) + + elif value_type == 'bytes': + return tf.train.Feature(bytes_list=tf.train.BytesList(value=[value])) + + elif value_type == 'bytes_list': + return tf.train.Feature(bytes_list=tf.train.BytesList(value=value)) + + else: + raise ValueError('Unknown value_type parameter - {}'.format(value_type)) + + +def convert_to_string_feature( + value: str, encoding: str = 'utf-8' +) -> tf.train.Feature: + """Returns a bytes_list from an encoded string.""" + return convert_to_feature(value.encode(encoding)) + + +def convert_to_list_string_feature( + lst: list[str], encoding: str = 'utf-8' +) -> tf.train.Feature: + """Returns a bytes_list from a list of encoded strings.""" + return convert_to_feature([value.encode(encoding) for value in lst]) + + +def create_ml_use_array_with_split( + total_size: int, + split_ratio: Sequence[float], +) -> list[str]: + """Create randomized list of 'training', 'validation', 'test'. + + The list of will be of length total_size with ratios according to train_size, + validation_size, and test_size. + + Args: + total_size: Length of sequence to return + split_ratio: Proportions to split into 'training', 'validation', and 'test' + + Returns: + List containing 'training', 'validation', and 'test' + """ + train_size, validation_size, _ = split_ratio + num_train = round(train_size * total_size) + num_validation = round(validation_size * total_size) + num_test = total_size - num_train - num_validation + ml_use_row = ( + [constants.ML_USE_TRAINING] * num_train + + [constants.ML_USE_VALIDATION] * num_validation + + [constants.ML_USE_TEST] * num_test + ) + random.shuffle(ml_use_row) + return ml_use_row + + +def format_ml_use_column(df: pd.DataFrame): + df[COLUMN_NAME_ML_USE].replace( + # We need to support non-standard ML uses other than documented ones, + # since they are used by some existing datasets. + [r'(?i)^train(ing)?$', r'(?i)^test$', r'(?i)^validat(ion|e)$'], + [ + constants.ML_USE_TRAINING, + constants.ML_USE_TEST, + constants.ML_USE_VALIDATION, + ], + inplace=True, + regex=True, + ) + + +def insert_missing_ml_use(df: pd.DataFrame) -> None: + """For every row that does not have ml_use as the first column, insert a column containing 'unassigned' to the front. + + Args: + df: The DataFrame to process. The first column should be 'ml_use'. + """ + df[COLUMN_NAME_ML_USE].fillna(ML_USE_UNASSIGNED, inplace=True) + rows_to_fill = ~df[COLUMN_NAME_ML_USE].isin(ALL_ML_USES) + df.loc[rows_to_fill] = df[rows_to_fill].shift( + axis=1, fill_value=ML_USE_UNASSIGNED + ) + + +def replace_unassigned_ml_use( + ml_uses: List[str], + split_ratio: Sequence[float], +): + """Replace `unassigned` in ml_uses with `training`, `validation`, and `test` with ratios according to split_ratio. + + Args: + ml_uses: List of ml_use string values. + split_ratio: Proportions to split into `training`, `validation`, and `test`. + """ + unassigned_indices = [ + i for i, ml_use in enumerate(ml_uses) if ml_use == ML_USE_UNASSIGNED + ] + ml_use_arr = create_ml_use_array_with_split( + len(unassigned_indices), split_ratio + ) + for unassigned_index, ml_use in zip(unassigned_indices, ml_use_arr): + ml_uses[unassigned_index] = ml_use + + +def merge_seq_into_dicts( + key: str, values: Sequence[Any], dicts: Sequence[Dict[Any, Any]] +): + """Merges a list of values into a list of dicts, inserted with the given key. + + Args: + key: Key to insert or overwrite in the dictionary. + values: A list of values to insert. + dicts: A list of dictionaries. Each value will be inserted into the + corresponding dictionary. The original value will be overwritten if the + key already existed. + + Raises: + ValueError: The values and dicts have different lengths. + """ + if len(values) != len(dicts): + raise ValueError( + f'Length of values and dicts must match, got {len(values)} and' + f' {len(dicts)}' + ) + for val, d in zip(values, dicts): + d[key] = val + + +def drop_invalid_rows(df: pd.DataFrame) -> int: + """Drops DataFrame rows missing the gcs_file_path column or the label column. + + Args: + df: The DataFrame to process in place. + + Returns: + The number of rows dropped. + """ + original_rows = df.shape[0] + df.dropna(subset=[COLUMN_NAME_GCS_FILE_PATH, COLUMN_NAME_LABEL], inplace=True) + dropped_num = original_rows - df.shape[0] + if dropped_num > 0: + df.reset_index(drop=True, inplace=True) + return dropped_num + + +def check_split_ratio(split_ratio: Sequence[float]): + """Checks if the give split ratio is valid. + + Args: + split_ratio: Proportions to split into 'training', 'validation', and 'test' + + Raises: + ValueError: Must have valid entries, correct length, and sum to 1. + """ + if len(split_ratio) != 3: + raise ValueError('split_ratio must contain exactly 3 values.') + if abs(sum(split_ratio) - 1) > _SPLIT_RATIO_ERROR_THRESHOLD: + raise ValueError('split_ratio must sum to 1.') + if not all([0 <= val <= 1 for val in split_ratio]): + raise ValueError('Entries of split_ratio must be in the range [0, 1].') + + +def check_num_shard(num_shard: Sequence[int]): + """Checks if the number of shards is valid. + + Args: + num_shard: The number of shards for each tfrecord. + + Raises: + ValueError: Must have valid entries and correct length. + """ + if len(num_shard) != 3: + raise ValueError('num_shard must contain exactly 3 values.') + if not all([val >= 1 for val in num_shard]): + raise ValueError('Shards must be at least 1.') + + +def create_label_map_yaml(meta_data_path: str, output_dir: str) -> None: + """Generate label_map.yaml from meta_data.yaml. + + Args: + meta_data_path: Path to a meta_data.yaml file. + output_dir: Directory to output label_map.yaml. + """ + tf.io.gfile.copy( + meta_data_path, os.path.join(output_dir, LABEL_MAP_NAME), overwrite=True + ) + + +def reformat_bbox( + bbox: Sequence[int], img_width: int, img_height: int +) -> Tuple[float, float, float, float]: + """Converts XYWH unnormalized bounding box with to a normalized XYXY bounding box. + + Args: + bbox: Relative bounding box with unnormalized coordinates as [x, y, width, + height]. + img_width: Image's pixel width. + img_height: Image's pixel height. + + Returns: + Absolute bounding box with normalized coordinates as + [xmin, ymin, xmax, ymax]. + """ + x, y, width, height = bbox + xmin = x / img_width + ymin = y / img_height + xmax = (x + width) / img_width + ymax = (y + height) / img_height + return xmin, ymin, xmax, ymax + + +def encode_image( + filepath: str, + output_shape: Optional[Sequence[int]] = None, + image_format: str = 'png', +) -> Tuple[bytes, Sequence[int]]: + """Encodes an image at the given path. + + Args: + filepath: Path to the image. + output_shape: The output shape of the image, (height, width). + image_format: The format of the output image. + + Returns: + The encoded image data in bytes and the shape of the image, (height, width). + + Raises: + IOError: The image file is corrupt. + """ + filepath = fileutils.force_gcs_fuse_path(filepath) + with open(filepath, 'rb') as f: + # If an output_shape is specified, resize the image and set data to the new + # bytes. + try: + img = Image.open(f) + except PIL.UnidentifiedImageError as e: + raise IOError(f'Failed to open {filepath}') from e + + try: + if output_shape is not None: + rgb_img = img.resize((output_shape[1], output_shape[0])).convert('RGB') + else: + rgb_img = img.convert('RGB') + rgb_img = np.array(rgb_img) + + _, data = cv2.imencode(f'.{image_format}', rgb_img) + data = data.tobytes() + return data, rgb_img.shape + except cv2.error as e: + raise IOError(f'Failed to encode {filepath}') from e + finally: + img.close() + + +def encode_video( + filepath: str, + start_sec: float, + end_sec: float, + output_fps: int = 5, + output_shape: Optional[Sequence[int]] = None, + image_format: str = 'jpg', +) -> Sequence[bytes]: + """Encodes a video clip at the given path with start and end timestamps. + + Args: + filepath: Path to the video. + start_sec: Start timestamp of the video clip in seconds. + end_sec: End timestamp of the video clip in seconds. + output_fps: The output frame rate per second. + output_shape: The output shape of each frame, (height, width). + image_format: The format of the encoded frames. + + Returns: + A list of the encoded frames data in bytes. + + Raises: + IOError if the video file is corrupt. + """ + filepath = fileutils.force_gcs_fuse_path(filepath) + video = None + + try: + video = cv2.VideoCapture(filepath) + frames = [] + frame_interval = 1 / output_fps + total_frames = video.get(cv2.CAP_PROP_FRAME_COUNT) + original_fps = video.get(cv2.CAP_PROP_FPS) + if not original_fps: + # 0 or None indicates the video is invalid + raise IOError(f'Failed to load {filepath}') + video_length = total_frames / original_fps + start_sec = max(start_sec, 0) + end_sec = min(end_sec, video_length) + for t in np.arange(start_sec, end_sec, frame_interval): + frame_idx = min(total_frames - 1, round(t * original_fps)) + video.set(cv2.CAP_PROP_POS_FRAMES, frame_idx) + ret, frame = video.read() + if not ret: + raise IOError(f'Failed to load {filepath} at frame {frame_idx}') + if output_shape is not None: + frame = cv2.resize(frame, (output_shape[1], output_shape[0])) + _, data = cv2.imencode(f'.{image_format}', frame) + frames.append(data.tobytes()) + except cv2.error as e: + raise IOError(f'Failed to load {filepath}') from e + finally: + if video: + video.release() + return frames + + +def create_label_map( + labels: Sequence[str], +) -> Tuple[Sequence[int], Dict[int, str]]: + """Creates a label map from a sequence of label strings. + + Args: + labels: The sequence of labels to create label map from. Must not contain + invalid values, which means data without labels should be filtered first. + + Returns: + The integer labels and the mapping from integers to the original strings. + """ + inverse_label_map: Dict[str, int] = dict() + num_labels = 0 + for label in labels: + if label not in inverse_label_map: + num_labels += 1 + inverse_label_map[label] = num_labels + int_labels = [inverse_label_map[label] for label in labels] + label_map = {value: key for key, value in inverse_label_map.items()} + return int_labels, label_map + + +def write_label_map(output_file: str, label_map: Dict[int, str]) -> None: + """Writes a label map to the output file, which can be a GCS uri.""" + with tf.io.gfile.GFile(output_file, 'w') as f: + yaml.dump({'label_map': label_map}, f) + + +def detectron_json_to_image_rows(input_json: str) -> list[Dict[str, Any]]: + """Converts a Detectron JSON file to a list of image rows. + + Args: + input_json: A path to a Detectron JSON or JSONL file. + + Returns: + A list of dictionaries, where each dictionary contains Detectron format + entry. + + Raises: + ValueError: If the input JSON is invalid. + """ + + image_rows = [] + with tf.io.gfile.GFile(input_json, 'r') as f: + for line in f: + json_data = json.loads(line) + if isinstance(json_data, dict): + image_rows.append(json_data) + elif isinstance(json_data, list): + image_rows.extend(json_data) + else: + raise ValueError( + 'The input JSON is invalid. Dict or list is expected, but got ' + f'{type(json_data)}.' + ) + return image_rows + + +def coco_json_to_image_rows( + input_json: str, +) -> List[Dict[str, Any]]: + """Converts a COCO JSON file to a list of image rows. + + Args: + input_json: A path to a COCO JSON or JSONL file. + + Returns: + A list of dictionaries, where each dictionary contains COCO format entry. + + Raises: + ValueError: If the input JSON is invalid. + """ + + with tf.io.gfile.GFile(input_json, 'r') as f: + coco_json = json.load(f) + if 'annotations' not in coco_json: + raise ValueError('"annotations" is not in the dataset.') + if 'images' not in coco_json: + raise ValueError('"images" is not in the dataset.') + + images = coco_json['images'] + return images + + +def partition_by_ml_use(element: Dict[str, Any], num_partitions: int) -> int: + """Beam partition function to split data by ml_use.""" + del num_partitions + try: + partition = ALL_ML_USES.index(element[COLUMN_NAME_ML_USE]) + except Exception as e: + raise ValueError(f'Invalid ML use: {element[COLUMN_NAME_ML_USE]}') from e + return partition + + +def run_beam_pipeline(pipeline: Any) -> None: + """Runs a beam pipeline. Works in both internal and docker environment.""" + options = pipeline_options.PipelineOptions([ + '--runner=FlinkRunner', + '--faster_copy', + '--max_parallelism', '8', + ]) + p = beam.Pipeline(options=options) + pipeline(p) + result = p.run() + result.wait_until_finish() + for counter in result.metrics().query()['counters']: + logging.info('%s counter: %s.', counter.key.metric.name, counter) + logging.info('Completing beam pipeline.') + + +def beam_convert_tfexamples( + root: beam.Pipeline, + data_list: Sequence[Dict[str, Any]], + convert_fn: Callable[[Dict[str, Any]], tf.train.Example], + output_dir: str, + num_shards: Sequence[int], +) -> None: + """Constructs beam pipelines to convert train, val, test TF Examples.""" + names = [TRAIN_TFRECORD_NAME, VALIDATION_TFRECORD_NAME, TEST_TFRECORD_NAME] + split_data = ( + root + | 'Create PCollection' >> beam.Create(data_list) + | 'Data split' >> beam.Partition(partition_by_ml_use, 3) + ) + for i in range(3): + ml_use: str = ALL_ML_USES[i] + num_shard = num_shards[i] + output_prefix = os.path.join(output_dir, names[i]) + _ = ( + split_data[i] + | f'Convert {ml_use} TF Examples' + >> beam.ParDo(WriteToTFRecord(output_prefix, num_shard, convert_fn)) + | f'Group {ml_use} TF Record files' >> beam.GroupBy(lambda x: x[0]) + | f'Merge {ml_use} TF Record files' + >> beam.Map(merge_tfrecords_func(output_prefix, num_shard)) + ) + + +def merge_tfrecords_func(output_prefix: str, num_shard: int) -> ...: + """Returns a function to merge sharded worker output into expected shards.""" + output_prefix = fileutils.force_gcs_fuse_path(output_prefix) + + def merge_tfrecords(worker_output: Tuple[int, Sequence[Tuple[int, str]]]): + idx = worker_output[0] + files: Sequence[str] = np.unique([x[1] for x in worker_output[1]]) + output_file = f'{output_prefix}-{idx:05d}-of-{num_shard:05d}' + with open(output_file, 'wb') as f: + for file in files: + logging.info('Merging %s.', file) + file = fileutils.force_gcs_fuse_path(file) + with open(file, 'rb') as fin: + while True: + data = fin.read(READ_CHUNK_SIZE) + if not data: + break + f.write(data) + os.remove(file) + + return merge_tfrecords diff --git a/community-content/vertex_vision_model_garden/model_oss/data_converter/convert_coco_label_map_iod.py b/community-content/vertex_vision_model_garden/model_oss/data_converter/convert_coco_label_map_iod.py new file mode 100644 index 000000000..c6b971724 --- /dev/null +++ b/community-content/vertex_vision_model_garden/model_oss/data_converter/convert_coco_label_map_iod.py @@ -0,0 +1,111 @@ +r"""Converts COCO labels as yamls for model garden playground (IOD). +""" + +import os +import urllib.request + +from absl import app +from absl import flags +import tensorflow as tf +import yaml + +from object_detection.utils import label_map_util + +_CONVERT_LABEL_TYPE_COCO_80 = 'coco_80' +_CONVERT_LABEL_TYPE_COCO_91 = 'coco_91' + +_CONVERT_LABEL_TYPE = flags.DEFINE_enum( + 'convert_label_type', + None, + [ + _CONVERT_LABEL_TYPE_COCO_80, + _CONVERT_LABEL_TYPE_COCO_91, + ], + 'Different types of label type conversion.', + required=True, +) + +_TEMPORARY_PATH = flags.DEFINE_string( + 'temporary_path', + None, + 'The tempory path.', + required=True, +) + +_OUTPUT_YAML_FILEPATH = flags.DEFINE_string( + 'output_yaml_filepath', + None, + 'The output yaml filepath.', + required=True, +) + + +def convert_coco_label_map_91( + output_yaml_filepath: str, +) -> None: + """Converts coco label map 91.""" + input_proto_filepath = 'https://raw.githubusercontent.com/tensorflow/models/master/research/object_detection/data/mscoco_label_map.pbtxt' + local_input_proto_filepath = os.path.join( + _TEMPORARY_PATH.value, 'mscoco_label_map.pbtxt' + ) + with open(local_input_proto_filepath, 'w') as writer: + contents = ( + urllib.request.urlopen(input_proto_filepath).read().decode('utf-8') + ) + writer.write(contents) + + label_map = label_map_util.load_labelmap(local_input_proto_filepath) + label_map_dict = label_map_util.get_label_map_dict( + label_map, use_display_name=True + ) + swapped_label_map_dict = {v: k for k, v in label_map_dict.items()} + print(swapped_label_map_dict) + + # Saves new label maps as yamls. + with tf.io.gfile.GFile(output_yaml_filepath, 'w') as writer: + writer.write(yaml.dump(swapped_label_map_dict)) + + +def convert_coco_label_map_80( + output_yaml_filepath: str, +) -> None: + """Converts coco label map 80.""" + # Loads label maps from texts. + input_text_filepath = 'https://gist.githubusercontent.com/AruniRC/7b3dadd004da04c80198557db5da4bda/raw/2f10965ace1e36c4a9dca76ead19b744f5eb7e88/ms_coco_classnames.txt' + local_input_text_filepath = os.path.join( + _TEMPORARY_PATH.value, 'ms_coco_classnames.txt' + ) + with open(local_input_text_filepath, 'w') as writer: + contents = ( + urllib.request.urlopen(input_text_filepath).read().decode('utf-8') + ) + writer.write(contents) + with open(local_input_text_filepath, 'r') as file: + content = file.read() + label_map = yaml.safe_load(content) + + # Removes background in label maps. + new_label_map = {} + for k, v in label_map.items(): + if k == 0: + continue + new_label_map[k - 1] = v + print(new_label_map) + # Saves new label maps as yamls. + with tf.io.gfile.GFile(output_yaml_filepath, 'w') as writer: + writer.write(yaml.dump(new_label_map)) + + +def main(_) -> None: + if _CONVERT_LABEL_TYPE.value == _CONVERT_LABEL_TYPE_COCO_80: + convert_coco_label_map_80(_OUTPUT_YAML_FILEPATH.value) + elif _CONVERT_LABEL_TYPE.value == _CONVERT_LABEL_TYPE_COCO_91: + convert_coco_label_map_91( + _OUTPUT_YAML_FILEPATH.value, + ) + else: + print('Not supported convert label type: ', _CONVERT_LABEL_TYPE.value) + + +if __name__ == '__main__': + app.run(main) diff --git a/community-content/vertex_vision_model_garden/model_oss/data_converter/convert_imagenet_label_map.py b/community-content/vertex_vision_model_garden/model_oss/data_converter/convert_imagenet_label_map.py new file mode 100644 index 000000000..408e742f3 --- /dev/null +++ b/community-content/vertex_vision_model_garden/model_oss/data_converter/convert_imagenet_label_map.py @@ -0,0 +1,86 @@ +r"""Converts ImageNet label texts as yamls for model garden playground. + +# ImageNet1K will have label maps with background. +""" + +import urllib.request +from absl import app +from absl import flags +import tensorflow as tf +import yaml + + +_INPUT_TEXT_FILEPATH = flags.DEFINE_string( + 'input_text_filepath', + None, + 'The input text filepath.', + required=True, +) + +_ADD_BACKGROUND_LABEL = flags.DEFINE_boolean( + 'add_background_label', + None, + 'Whether or not add background labels.', + required=True, +) + +_ADD_IDS = flags.DEFINE_boolean( + 'add_ids', + None, + 'Whether or not add ids.', + required=True, +) + +_OUTPUT_YAML_FILEPATH = flags.DEFINE_string( + 'output_yaml_filepath', + None, + 'The output yaml filepath.', + required=True, +) + + +def convert_imagenet_label_map_from_text_to_yaml( + input_text_filepath: str, + add_background_label: bool, + add_ids: bool, + output_yaml_filepath: str, +) -> None: + """Converts imagenet label map from text to yamls.""" + label_map = {} + + # Shifts all keys by 1, and add 0 as 'background'. + if add_background_label: + label_map = yaml.safe_load( + urllib.request.urlopen(input_text_filepath).read() + ) + new_label_map = {} + for key, value in label_map.items(): + new_label_map[key + 1] = value + new_label_map[0] = 'background' + label_map = new_label_map + + # Adds maps from id to each line. + if add_ids: + lines = urllib.request.urlopen(input_text_filepath).readlines() + current_id = 0 + for line in lines: + label_map[current_id] = line.decode('ascii').strip() + print(label_map[current_id]) + current_id += 1 + + # Saves new label maps as yamls. + with tf.io.gfile.GFile(output_yaml_filepath, 'w') as writer: + writer.write(yaml.dump(label_map)) + + +def main(_) -> None: + convert_imagenet_label_map_from_text_to_yaml( + _INPUT_TEXT_FILEPATH.value, + _ADD_BACKGROUND_LABEL.value, + _ADD_IDS.value, + _OUTPUT_YAML_FILEPATH.value, + ) + + +if __name__ == '__main__': + app.run(main) diff --git a/community-content/vertex_vision_model_garden/model_oss/data_converter/data_converter_icn_lib.py b/community-content/vertex_vision_model_garden/model_oss/data_converter/data_converter_icn_lib.py new file mode 100644 index 000000000..7b9475d97 --- /dev/null +++ b/community-content/vertex_vision_model_garden/model_oss/data_converter/data_converter_icn_lib.py @@ -0,0 +1,199 @@ +"""Converts ICN CSV/JSONL files to TFRecord with apache beam.""" + +import json +from os import path +from typing import Any, Dict, Sequence, Union, cast + +from absl import logging +import apache_beam as beam +import pandas as pd +import tensorflow as tf + +from data_converter import common_lib + + +_COLUMN_NAMES = [ + common_lib.COLUMN_NAME_ML_USE, + common_lib.COLUMN_NAME_GCS_FILE_PATH, + common_lib.COLUMN_NAME_LABEL, +] +_JSON_GCS_URI_KEY = 'imageGcsUri' +_JSON_CLASS_ANNOTATION_KEY = 'classificationAnnotation' +_JSON_RESOURCE_LABEL_KEY = 'dataItemResourceLabels' +_JSON_CLASS_NAME_KEY = 'displayName' +_JSON_ML_USE_KEY = 'aiplatform.googleapis.com/ml_use' + + +def build_tf_example(element: Dict[str, Union[str, int]]) -> tf.train.Example: + """Builds a TF Example from an image uri and label. + + Args: + element: A dict with the keys gcs_file_path and label. + + Returns: + The created TF Example. + """ + image_uri = cast(str, element[common_lib.COLUMN_NAME_GCS_FILE_PATH]) + label = cast(int, element[common_lib.COLUMN_NAME_LABEL]) + image_bytes, shape = common_lib.encode_image(image_uri, image_format='jpeg') + features = tf.train.Features( + feature={ + 'image/encoded': common_lib.convert_to_feature(image_bytes), + 'image/format': common_lib.convert_to_string_feature('jpeg'), + 'image/height': common_lib.convert_to_feature(shape[0]), + 'image/width': common_lib.convert_to_feature(shape[1]), + 'image/class/label': common_lib.convert_to_feature(label), + }, + ) + return tf.train.Example(features=features) + + +def _run_convert_pipeline( + output_dir: str, df: pd.DataFrame, num_shards: Sequence[int] +) -> None: + """Starts a Beam pipeline to write DataFrame as TF Records. + + Args: + output_dir: TF Records output directory. + df: DataFrame to convert from. + num_shards: Number of shards for train/validation/test TFRecord files. + """ + images_list = df.to_dict('records') + + def pipeline(root: beam.Pipeline): + common_lib.beam_convert_tfexamples( + root, + images_list, + build_tf_example, + output_dir, + num_shards, + ) + + common_lib.run_beam_pipeline(pipeline) + + +def _convert_df_to_tfrecord( + df: pd.DataFrame, + output_dir: str, + split_ratio: Sequence[float], + num_shard: Sequence[int], +) -> None: + """Converts a DataFrame into three separate tfrecords for training, validation, and testing into output_dir. + + Args: + df: DataFrame to convert. + output_dir: The directory to save TFRecords and label_map.yaml. + split_ratio: List specifying the training, validation, and testing splits + for unassigned TFRecords. + num_shard: Number of shards for train/validation/test TFRecord files. + """ + # Replaces ml_use with common_lib string constants for consistency. + common_lib.format_ml_use_column(df) + common_lib.insert_missing_ml_use(df) + + # Ignores invalid rows. + dropped_row_num = common_lib.drop_invalid_rows(df) + if dropped_row_num > 0: + logging.warning('Ignored %d invalid rows.', dropped_row_num) + + common_lib.replace_unassigned_ml_use( + df[common_lib.COLUMN_NAME_ML_USE], split_ratio + ) + + # Converts labels to integers as required by training. + new_labels, label_map = common_lib.create_label_map( + df[common_lib.COLUMN_NAME_LABEL] + ) + df[common_lib.COLUMN_NAME_LABEL] = new_labels + label_map_path = path.join(output_dir, common_lib.LABEL_MAP_NAME) + logging.info('Writing label map to %s.', label_map_path) + common_lib.write_label_map(label_map_path, label_map) + + _run_convert_pipeline(output_dir, df, num_shard) + + +def convert_csv_to_tfrecord( + input_csv: str, + output_dir: str, + split_ratio: Sequence[float] = (0.8, 0.1, 0.1), + num_shard: Sequence[int] = (10, 10, 10), +) -> None: + """Parses input_csv file into three separate tfrecords for training, validation, and testing into output_dir. + + The csv format is shown in + https://cloud.google.com/vertex-ai/docs/image-data/classification/prepare-data#csv. + + If an ml_use column is not provided, one will be created. + + label_map.yaml containing the label map will be placed in output_dir. + + Args: + input_csv: Name of the csv file. + output_dir: The directory to save TFRecords and label_map.yaml. + split_ratio: List specifying the training, validation, and testing splits + for unassigned TFRecords. + num_shard: Number of shards for train/validation/test TFRecord files. + """ + with tf.io.gfile.GFile(input_csv, 'r') as f: + df: pd.DataFrame = pd.read_csv( + f, header=None, names=_COLUMN_NAMES, on_bad_lines='warn' + ) + + _convert_df_to_tfrecord(df, output_dir, split_ratio, num_shard) + + +def convert_jsonl_to_tfrecord( + input_jsonl: str, + output_dir: str, + split_ratio: Sequence[float] = (0.8, 0.1, 0.1), + num_shard: Sequence[int] = (10, 10, 10), +) -> None: + """Parses input_jsonl file into three separate tfrecords for training, validation, and testing into output_dir. + + The JSONL format is shown in + https://cloud.google.com/vertex-ai/docs/image-data/classification/prepare-data#json-lines. + + If an ml_use column is not provided, one will be created. + + label_map.yaml containing the label map will be placed in output_dir. + + Args: + input_jsonl: Name of the JSONL file. + output_dir: The directory to save TFRecords and label_map.yaml. + split_ratio: List specifying the training, validation, and testing splits + for unassigned TFRecords. + num_shard: Number of shards for train/validation/test TFRecord files. + """ + df_rows = [] + with tf.io.gfile.GFile(input_jsonl, 'r') as f: + lines = f.read().rstrip().splitlines() + + for i, line in enumerate(lines, 1): + try: + item: Dict[str, Any] = json.loads(line) + + gcs_uri = item.get(_JSON_GCS_URI_KEY) + label = item.get(_JSON_CLASS_ANNOTATION_KEY, {}).get(_JSON_CLASS_NAME_KEY) + if not gcs_uri or not label: + logging.warning('Invalid JSON at line %d, skipped.', i) + continue + + ml_use = item.get(_JSON_RESOURCE_LABEL_KEY, {}).get( + _JSON_ML_USE_KEY, common_lib.ML_USE_UNASSIGNED + ) + except (json.JSONDecodeError, AttributeError): + logging.warning('Invalid JSON at line %d, skipped.', i) + continue + + df_rows.append([ml_use, gcs_uri, label]) + + df = pd.DataFrame( + data=df_rows, + columns=[ + common_lib.COLUMN_NAME_ML_USE, + common_lib.COLUMN_NAME_GCS_FILE_PATH, + common_lib.COLUMN_NAME_LABEL, + ], + ) + + _convert_df_to_tfrecord(df, output_dir, split_ratio, num_shard) diff --git a/community-content/vertex_vision_model_garden/model_oss/data_converter/data_converter_iod_lib.py b/community-content/vertex_vision_model_garden/model_oss/data_converter/data_converter_iod_lib.py new file mode 100644 index 000000000..17ce35d7e --- /dev/null +++ b/community-content/vertex_vision_model_garden/model_oss/data_converter/data_converter_iod_lib.py @@ -0,0 +1,430 @@ +"""Converts IOD dataset files to TFRecord with apache beam.""" + +import collections +import json +from os import path +from typing import Any, Dict, Sequence + +from absl import logging +import apache_beam as beam +import pandas as pd +import tensorflow as tf + +from data_converter import common_lib +from util import constants + +COLUMN_NAME_LABEL_INT = 'label_int' +_COLUMN_NAME_XMIN = 'X_MIN' +_COLUMN_NAME_YMIN = 'Y_MIN' +_COLUMN_NAME_XMAX = 'X_MAX' +_COLUMN_NAME_YMAX = 'Y_MAX' +COLUMN_NAMES = [ + common_lib.COLUMN_NAME_ML_USE, + common_lib.COLUMN_NAME_GCS_FILE_PATH, + common_lib.COLUMN_NAME_LABEL, + _COLUMN_NAME_XMIN, + _COLUMN_NAME_YMIN, + 'XMAX_NOT_USED', + 'YMIN_NOT_USED', + _COLUMN_NAME_XMAX, + _COLUMN_NAME_YMAX, + 'XMIN_NOT_USED', + 'YMAX_NOT_USED', +] +_BOUNDING_BOX_COLUMNS = [ + _COLUMN_NAME_XMIN, + _COLUMN_NAME_YMIN, + _COLUMN_NAME_XMAX, + _COLUMN_NAME_YMAX, +] +_JSON_BBOX_ANNOTATIONS_KEY = 'boundingBoxAnnotations' +_JSON_DISPLAY_NAME_KEY = 'displayName' +_JSON_X_MIN_KEY = 'xMin' +_JSON_X_MAX_KEY = 'xMax' +_JSON_Y_MIN_KEY = 'yMin' +_JSON_Y_MAX_KEY = 'yMax' + + +def build_tf_example(image_row: Dict[str, Any]) -> tf.train.Example: + """Builds a TF Example from an image row. + + Args: + image_row: A dictionary containing information about the image, such as its + GCS uri, labels, and bounding box coordinates. + + Returns: + A tf.train.Example containing the encoded image and optionally a + bounding box and label. + """ + image_uri = image_row[common_lib.COLUMN_NAME_GCS_FILE_PATH] + image_bytes, shape = common_lib.encode_image(image_uri, image_format='jpeg') + feature = { + 'image/encoded': common_lib.convert_to_feature(image_bytes), + 'image/format': common_lib.convert_to_string_feature('jpeg'), + 'image/height': common_lib.convert_to_feature(shape[0]), + 'image/width': common_lib.convert_to_feature(shape[1]), + 'image/source_id': common_lib.convert_to_string_feature(image_uri), + 'image/object/bbox/xmin': common_lib.convert_to_feature( + image_row[_COLUMN_NAME_XMIN] + ), + 'image/object/bbox/ymin': common_lib.convert_to_feature( + image_row[_COLUMN_NAME_YMIN] + ), + 'image/object/bbox/xmax': common_lib.convert_to_feature( + image_row[_COLUMN_NAME_XMAX] + ), + 'image/object/bbox/ymax': common_lib.convert_to_feature( + image_row[_COLUMN_NAME_YMAX] + ), + 'image/object/class/text': common_lib.convert_to_list_string_feature( + image_row[common_lib.COLUMN_NAME_LABEL] + ), + 'image/object/class/label': common_lib.convert_to_feature( + image_row[COLUMN_NAME_LABEL_INT] + ), + } + return tf.train.Example(features=tf.train.Features(feature=feature)) + + +def _run_convert_pipeline( + output_dir: str, + image_rows: Sequence[Dict[str, Any]], + num_shards: Sequence[int], +) -> None: + """Starts a Beam pipeline to write DataFrame as TF Records. + + Args: + output_dir: TF Records output directory. + image_rows: Contains all necessary information to create a TF Example. + num_shards: Number of shards for train/validation/test TFRecord files. + """ + + def pipeline(root: beam.Pipeline): + common_lib.beam_convert_tfexamples( + root, + image_rows, + build_tf_example, + output_dir, + num_shards, + ) + + common_lib.run_beam_pipeline(pipeline) + + +def _convert_df_to_tfrecord( + df: pd.DataFrame, + output_dir: str, + split_ratio: Sequence[float], + num_shard: Sequence[int], +) -> None: + """Converts a DataFrame into three separate tfrecords for training, validation, and testing into output_dir. + + Args: + df: DataFrame to convert. + output_dir: The directory to save TFRecords and label_map.yaml. + split_ratio: List specifying the training, validation, and testing splits + for unassigned TFRecords. + num_shard: Number of shards for train/validation/test TFRecord files. + """ + # Replaces ml_use with common_lib string constants for consistency. + common_lib.format_ml_use_column(df) + common_lib.insert_missing_ml_use(df) + + # Specify bounding box columns to be numeric. + df[_BOUNDING_BOX_COLUMNS] = df[_BOUNDING_BOX_COLUMNS].apply(pd.to_numeric) + + # Ignores invalid rows. + dropped_row_num = common_lib.drop_invalid_rows(df) + dropped_row_num += drop_rows_without_bbox(df) + if dropped_row_num > 0: + logging.warning('Ignored %d invalid rows.', dropped_row_num) + + # Converts labels to integers as required by training. + int_labels, label_map = common_lib.create_label_map( + df[common_lib.COLUMN_NAME_LABEL] + ) + df[COLUMN_NAME_LABEL_INT] = int_labels + label_map_path = path.join(output_dir, common_lib.LABEL_MAP_NAME) + logging.info('Writing label map to %s.', label_map_path) + common_lib.write_label_map(label_map_path, label_map) + + image_rows = _condense_bounding_boxes(df.to_dict(orient='records')) + ml_uses = [row[common_lib.COLUMN_NAME_ML_USE] for row in image_rows] + common_lib.replace_unassigned_ml_use(ml_uses, split_ratio) + common_lib.merge_seq_into_dicts( + common_lib.COLUMN_NAME_ML_USE, ml_uses, image_rows + ) + + _run_convert_pipeline(output_dir, image_rows, num_shard) + + +def _condense_bounding_boxes( + image_rows: Sequence[Dict[str, Any]] +) -> Sequence[Dict[str, Any]]: + """Gather all the bounding boxes in an image and put them in the same dictionary. + + Args: + image_rows: List of dictionaries, each containing information about the + image, such as its GCS uri, labels, and bounding box coordinates. + + Returns: + List of dictionaries such that each contains all the bounding boxes for a + given gcs_file_path. + + Raises: + RuntimeError: This is raised when the input data contains images that have + annotations in different ml_use classes. + """ + output = {} + for image_row in image_rows: + ml_use = image_row[common_lib.COLUMN_NAME_ML_USE] + gcs_file_path = image_row[common_lib.COLUMN_NAME_GCS_FILE_PATH] + label = image_row[common_lib.COLUMN_NAME_LABEL] + xmin = image_row[_COLUMN_NAME_XMIN] + ymin = image_row[_COLUMN_NAME_YMIN] + xmax = image_row[_COLUMN_NAME_XMAX] + ymax = image_row[_COLUMN_NAME_YMAX] + label_int = image_row[COLUMN_NAME_LABEL_INT] + if gcs_file_path in output: + d = output[gcs_file_path] + if ml_use != common_lib.ML_USE_UNASSIGNED: + if d[common_lib.COLUMN_NAME_ML_USE] == common_lib.ML_USE_UNASSIGNED: + d[common_lib.COLUMN_NAME_ML_USE] = ml_use + elif ml_use != d[common_lib.COLUMN_NAME_ML_USE]: + raise RuntimeError( + f'Image {gcs_file_path} can only be placed in one of' + f' training/validation/test. It is currently in {ml_use} and' + f' {d[common_lib.COLUMN_NAME_ML_USE]}.' + ) + d[common_lib.COLUMN_NAME_LABEL].append(label) + d[_COLUMN_NAME_XMIN].append(xmin) + d[_COLUMN_NAME_YMIN].append(ymin) + d[_COLUMN_NAME_XMAX].append(xmax) + d[_COLUMN_NAME_YMAX].append(ymax) + d[COLUMN_NAME_LABEL_INT].append(label_int) + else: + output[gcs_file_path] = { + common_lib.COLUMN_NAME_ML_USE: ml_use, + common_lib.COLUMN_NAME_GCS_FILE_PATH: gcs_file_path, + common_lib.COLUMN_NAME_LABEL: [label], + _COLUMN_NAME_XMIN: [xmin], + _COLUMN_NAME_YMIN: [ymin], + _COLUMN_NAME_XMAX: [xmax], + _COLUMN_NAME_YMAX: [ymax], + COLUMN_NAME_LABEL_INT: [label_int], + } + return list(output.values()) + + +def convert_csv_to_tfrecord( + input_csv: str, + output_dir: str, + split_ratio: Sequence[float] = (0.8, 0.1, 0.1), + num_shard: Sequence[int] = (10, 10, 10), +) -> None: + """Parses input_csv file into three separate tfrecords for training, validation, and testing into output_dir. + + The csv format is shown in + https://cloud.google.com/vertex-ai/docs/image-data/object-detection/prepare-data#csv. + + If an ml_use column is not provided, one will be created. + + label_map.yaml containing the label map will be placed in output_dir. + + Args: + input_csv: Name of the csv file. + output_dir: The directory to save TFRecords and label_map.yaml. + split_ratio: List specifying the train, validation, and test splits for + unassigned TFRecords. + num_shard: Number of shards for train/validation/test TFRecord files. + """ + with tf.io.gfile.GFile(input_csv, 'r') as f: + df: pd.DataFrame = pd.read_csv( + f, header=None, names=COLUMN_NAMES, on_bad_lines='warn' + ) + + _convert_df_to_tfrecord(df, output_dir, split_ratio, num_shard) + + +def drop_rows_without_bbox(df: pd.DataFrame) -> int: + """Drops DataFrame rows without bounding_boxes. + + Args: + df: The DataFrame to process in place. + + Returns: + The number of rows dropped. + """ + invalid_rows = df.index[~(df[_BOUNDING_BOX_COLUMNS].notnull().all(axis=1))] + dropped_num = len(invalid_rows) + if dropped_num > 0: + invalid_df = df.loc[invalid_rows].to_dict(orient='records') + for entry in invalid_df: + logging.warning('Skipping entry due to missing bounding box: %s.', entry) + df.drop(invalid_rows, inplace=True) + df.reset_index(drop=True, inplace=True) + return dropped_num + + +def convert_coco_json_categories_to_label_map( + categories: Sequence[Dict[str, Any]] +) -> Dict[int, str]: + return {category['id']: category['name'] for category in categories} + + +def convert_coco_json_to_tfrecord( + input_coco_json: str, + output_dir: str, + split_ratio: Sequence[float] = (0.8, 0.1, 0.1), + num_shard: Sequence[int] = (10, 10, 10), +) -> None: + """Parses input_csv file into three separate tfrecords for training, validation, and testing into output_dir. + + The COCO json format is shown here: https://cocodataset.org/#format-data. + + label_map.yaml containing the label map will be placed in output_dir. + + Args: + input_coco_json: Name of coco json file. + output_dir: The directory to save TFRecords and label_map.yaml. + split_ratio: List specifying the train, validation, and test splits for + dataset. + num_shard: Number of shards for train/validation/test TFRecord files. + """ + with tf.io.gfile.GFile(input_coco_json, 'r') as f: + coco_json = json.load(f) + # Writes label map from coco json categories. + label_map = convert_coco_json_categories_to_label_map( + coco_json[constants.COCO_JSON_CATEGORIES] + ) + label_map_path = path.join(output_dir, common_lib.LABEL_MAP_NAME) + logging.info('Writes label map to %s.', label_map_path) + common_lib.write_label_map(label_map_path, label_map) + + img_to_anns = collections.defaultdict(list) + imgs = {} + if constants.COCO_JSON_ANNOTATIONS in coco_json: + for ann in coco_json[constants.COCO_JSON_ANNOTATIONS]: + img_to_anns[ann[constants.COCO_JSON_ANNOTATION_IMAGE_ID]].append(ann) + + if constants.COCO_JSON_IMAGES in coco_json: + for img in coco_json[constants.COCO_JSON_IMAGES]: + imgs[img[constants.COCO_JSON_IMAGE_ID]] = img + + df_rows = [] + + for image_id, annotations in img_to_anns.items(): + img = imgs[image_id] + for ann in annotations: + xmin, ymin, xmax, ymax = common_lib.reformat_bbox( + ann[constants.COCO_ANNOTATION_BBOX], + img[constants.COCO_JSON_IMAGE_WIDTH], + img[constants.COCO_JSON_IMAGE_HEIGHT], + ) + df_rows.append([ + common_lib.ML_USE_UNASSIGNED, + img[constants.COCO_JSON_IMAGE_COCO_URL], + label_map[ann[constants.COCO_JSON_ANNOTATION_CATEGORY_ID]], + xmin, + ymin, + xmax, + ymin, + xmax, + ymax, + xmin, + ymax, + ann[constants.COCO_JSON_ANNOTATION_CATEGORY_ID], + ]) + df = pd.DataFrame( + data=df_rows, + columns=COLUMN_NAMES + [COLUMN_NAME_LABEL_INT], + ) + + # Replaces ml_use with common_lib string constants for consistency. + common_lib.format_ml_use_column(df) + common_lib.insert_missing_ml_use(df) + + # Species bounding box columns to be numeric. + df[_BOUNDING_BOX_COLUMNS] = df[_BOUNDING_BOX_COLUMNS].apply(pd.to_numeric) + + # Ignores invalid rows. + dropped_row_num = common_lib.drop_invalid_rows(df) + dropped_row_num += drop_rows_without_bbox(df) + if dropped_row_num > 0: + logging.warning('Ignored %d invalid rows.', dropped_row_num) + + image_rows = _condense_bounding_boxes(df.to_dict(orient='records')) + ml_uses = [row[common_lib.COLUMN_NAME_ML_USE] for row in image_rows] + common_lib.replace_unassigned_ml_use(ml_uses, split_ratio) + common_lib.merge_seq_into_dicts( + common_lib.COLUMN_NAME_ML_USE, ml_uses, image_rows + ) + + _run_convert_pipeline(output_dir, image_rows, num_shard) + + +def convert_jsonl_to_tfrecord( + input_jsonl: str, + output_dir: str, + split_ratio: Sequence[float] = (0.8, 0.1, 0.1), + num_shard: Sequence[int] = (10, 10, 10), +) -> None: + """Parses input_jsonl file into three separate tfrecords for training, validation, and testing into output_dir. + + The JSONL format is shown in + https://cloud.google.com/vertex-ai/docs/image-data/object-detection/prepare-data#json-lines. + + If an ml_use column is not provided, one will be created. + + label_map.yaml containing the label map will be placed in output_dir. + + Args: + input_jsonl: Name of the JSONL file. + output_dir: The directory to save TFRecords and label_map.yaml. + split_ratio: List specifying the training, validation, and testing splits + for unassigned TFRecords. + num_shard: Number of shards for train/validation/test TFRecord files. + """ + df_rows = [] + with tf.io.gfile.GFile(input_jsonl, 'r') as f: + lines = f.read().rstrip().splitlines() + + for i, line in enumerate(lines, start=1): + try: + item: Dict[str, Any] = json.loads(line) + except (json.JSONDecodeError, AttributeError): + logging.warning('Invalid JSON at line %d skipped.', i) + continue + + gcs_uri = item.get(common_lib.JSON_GCS_URI_KEY) + if not gcs_uri: + logging.warning( + 'Invalid JSON at line %d skipped. Missing gcs_uri_key.', i + ) + continue + ml_use = item.get(common_lib.JSON_RESOURCE_LABEL_KEY, {}).get( + common_lib.JSON_ML_USE_KEY, common_lib.ML_USE_UNASSIGNED + ) + + for bbox in item.get(_JSON_BBOX_ANNOTATIONS_KEY, []): + label = bbox.get(_JSON_DISPLAY_NAME_KEY) + xmin = bbox.get(_JSON_X_MIN_KEY) + ymin = bbox.get(_JSON_Y_MIN_KEY) + xmax = bbox.get(_JSON_X_MAX_KEY) + ymax = bbox.get(_JSON_Y_MAX_KEY) + + df_rows.append([ml_use, gcs_uri, label, xmin, ymin, xmax, ymax]) + + df = pd.DataFrame( + data=df_rows, + columns=[ + common_lib.COLUMN_NAME_ML_USE, + common_lib.COLUMN_NAME_GCS_FILE_PATH, + common_lib.COLUMN_NAME_LABEL, + _COLUMN_NAME_XMIN, + _COLUMN_NAME_YMIN, + _COLUMN_NAME_XMAX, + _COLUMN_NAME_YMAX, + ], + ) + _convert_df_to_tfrecord(df, output_dir, split_ratio, num_shard) diff --git a/community-content/vertex_vision_model_garden/model_oss/data_converter/data_converter_isg_lib.py b/community-content/vertex_vision_model_garden/model_oss/data_converter/data_converter_isg_lib.py new file mode 100644 index 000000000..bbadf6844 --- /dev/null +++ b/community-content/vertex_vision_model_garden/model_oss/data_converter/data_converter_isg_lib.py @@ -0,0 +1,328 @@ +"""Python script to convert different file formats for ISG to tfrecords.""" + +import hashlib +import os +from typing import Any, Dict, Iterator, List, Optional, Tuple, Union + +from absl import logging +import apache_beam as beam +from apache_beam.io import tfrecordio +import cv2 +import numpy as np +from pycocotools import coco +import tensorflow as tf +import yaml + +from data_converter import common_lib +from util import constants +from util import fileutils + +_IMAGE_FORMAT = 'PNG' + + +def build_tf_example( + image_info: dict[str, Union[str, int]], + segmentation_image: List[List[int]], + output_shape: Optional[Tuple[int, int]] = None, +) -> tf.train.Example: + """Encodes an image and its segmentation mask into a tf.train.Example. + + Args: + image_info: A dictionary containing information about the image, such as its + file name, height, and width. + segmentation_image: 2D image in list of lists having category ids. + output_shape: The desired output shape of the image. If None, the original + image shape will be used. + + Returns: + A tf.train.Example containing the encoded image and segmentation mask. + + Raises: + IOError: If image cannot be found in the path. + """ + file_name = image_info[constants.COCO_JSON_FILE_NAME] + height = int(image_info[constants.COCO_JSON_IMAGE_HEIGHT]) + width = int(image_info[constants.COCO_JSON_IMAGE_WIDTH]) + + segmentation_image = np.expand_dims( + np.asarray(segmentation_image, dtype=np.int32), axis=-1 + ) + _, encoded_seg = cv2.imencode(f'.{_IMAGE_FORMAT.lower()}', segmentation_image) + encoded_seg = encoded_seg.tobytes() + + encoded_img, _ = common_lib.encode_image( + image_info[constants.COCO_JSON_IMAGE_COCO_URL], + output_shape=output_shape, + image_format=_IMAGE_FORMAT.lower(), + ) + + key = hashlib.sha256(encoded_img).hexdigest() + + return tf.train.Example( + features=tf.train.Features( + feature={ + 'image/height': common_lib.convert_to_feature(height), + 'image/width': common_lib.convert_to_feature(width), + 'image/filename': common_lib.convert_to_string_feature(file_name), + 'image/sha256': common_lib.convert_to_string_feature(key), + 'image/encoded': common_lib.convert_to_feature(encoded_img), + 'image/format': common_lib.convert_to_string_feature( + _IMAGE_FORMAT + ), + 'image/segmentation/class/encoded': common_lib.convert_to_feature( + encoded_seg + ), + 'image/segmentation/class/format': ( + common_lib.convert_to_string_feature(_IMAGE_FORMAT) + ), + 'image/segmentation/class/height': common_lib.convert_to_feature( + height + ), + 'image/segmentation/class/width': common_lib.convert_to_feature( + width + ), + } + ) + ) + + +class AcquireTFExampleDoFn(beam.DoFn): + """Beam DoFn to build TF Examples from a single row of image_info data.""" + + # These tags will be used to tag the outputs of this DoFn. + output_tag_train = constants.ML_USE_TRAINING + output_tag_validation = constants.ML_USE_VALIDATION + output_tag_test = constants.ML_USE_TEST + + valid_ml_use_set = set( + [output_tag_train, output_tag_validation, output_tag_test] + ) + + def __init__(self, output_shape: Optional[Tuple[int, int]] = None): + self.acquired_examples_counter = beam.metrics.Metrics.counter( + self.__class__.__name__, 'Success' + ) + self.failure_counter = beam.metrics.Metrics.counter( + self.__class__.__name__, 'Failure' + ) + self.output_shape = output_shape + + def process( + self, + row: Tuple[str, Dict[str, Union[str, int]], List[List[int]]], + ) -> Iterator[tf.train.Example]: + ml_use, image_info, annotation_info = row + if ml_use not in self.valid_ml_use_set: + logging.warning('ml_use invalid: %s', ml_use) + self.failure_counter.inc() + return + + try: + tf_example = build_tf_example( + image_info, annotation_info, self.output_shape + ) + except IOError as e: + logging.warning('Failed to build TF Example: %s', e) + self.failure_counter.inc() + else: + self.acquired_examples_counter.inc() + yield beam.pvalue.TaggedOutput(ml_use, tf_example) + + +def _define_data_conversion_pipeline( + root: beam.Pipeline, + ml_use_rows: List[str], + image_rows: List[Dict[str, Union[str, int]]], + segmentation_rows: List[List[List[int]]], + output_dir: str, + output_shape: Optional[Tuple[int, int]], + num_shard_list: List[int], +): + """Define a data conversion pipeline. + + Args: + root: A Beam pipeline. + ml_use_rows: List containing the ml_use. + image_rows: List of dictionaries containing information about the image, + such as its file name, height, and width. + segmentation_rows: List of 2D images of integers representing segmentation + masks. + output_dir: Directory where the output TFRecords will be written. + output_shape: Desired output shape of the image. If None, the original image + shape will be used. + num_shard_list: Number of shards to write to each output TFRecord. + + Returns: + A Beam pipeline. + """ + train, validation, test = ( + root + | 'Load ml use and image rows to beam' + >> beam.Create(zip(ml_use_rows, image_rows, segmentation_rows)) + | 'Build TF Examples' + >> beam.ParDo(AcquireTFExampleDoFn(output_shape)).with_outputs( + AcquireTFExampleDoFn.output_tag_train, + AcquireTFExampleDoFn.output_tag_validation, + AcquireTFExampleDoFn.output_tag_test, + ) + ) + + # Save each split to TFRecord. + _ = train | 'Save train split to TFRecord' >> tfrecordio.WriteToTFRecord( + os.path.join(output_dir, common_lib.TRAIN_TFRECORD_NAME), + coder=beam.coders.ProtoCoder(tf.train.Example), + num_shards=num_shard_list[0], + ) + _ = ( + validation + | 'Save validation split to TFRecord' + >> tfrecordio.WriteToTFRecord( + os.path.join(output_dir, common_lib.VALIDATION_TFRECORD_NAME), + coder=beam.coders.ProtoCoder(tf.train.Example), + num_shards=num_shard_list[1], + ) + ) + _ = test | 'Save test split to TFRecord' >> tfrecordio.WriteToTFRecord( + os.path.join(output_dir, common_lib.TEST_TFRECORD_NAME), + coder=beam.coders.ProtoCoder(tf.train.Example), + num_shards=num_shard_list[2], + ) + + +def _image_info_to_segmentation_image( + img: Dict[str, Any], + coco_dataset: coco.COCO, + label_id_by_category_id: Dict[int, int], +) -> List[List[int]]: + """Convert image information to a segmentation image. + + Args: + img: The image information. + coco_dataset: The COCO dataset. + label_id_by_category_id: The mapping from label id used for training to + category_id defined in dataset. + + Returns: + The segmentation image. + + Raises: + ValueError: If the mask size does not match the image or if a pixel has + multiple labels. + """ + seg_img = np.zeros( + shape=( + img[constants.COCO_JSON_IMAGE_HEIGHT], + img[constants.COCO_JSON_IMAGE_WIDTH], + ), + dtype=np.int32, + ) + for ann in coco_dataset.imgToAnns[img[constants.COCO_JSON_IMAGE_ID]]: + new_category_id = ann[constants.COCO_JSON_ANNOTATION_CATEGORY_ID] + binary_mask = coco_dataset.annToMask(ann) + if seg_img.shape != binary_mask.shape: + raise ValueError( + 'Binary mask does not have the same shape as image. image_id:' + f' {img["id"]}' + ) + boolean_mask = binary_mask == 1 + if (seg_img[boolean_mask] != 0).any(): + raise ValueError( + 'Error: Some pixels have more than one label in image_id:' + f' {img["id"]}.' + ) + seg_img[boolean_mask] = label_id_by_category_id[new_category_id] + + return seg_img.tolist() + + +def get_input_rows( + coco_dataset: coco.COCO, + split_ratio: List[float], + label_id_by_category_id: Dict[int, int], +) -> Tuple[List[str], List[Dict[str, Union[str, int]]], List[List[List[int]]]]: + """Get input rows for training and validation. + + Args: + coco_dataset: The COCO dataset. + split_ratio: The split ratio for training and validation. + label_id_by_category_id: The mapping from label id used for training to + category_id defined in dataset. + + Returns: + - A list of ml_use strings. + - A list of image informations. + - A list of segmentation images for the corresponding images. + """ + image_rows = coco_dataset.dataset[constants.COCO_JSON_IMAGES] + + segmentation_rows = [ + _image_info_to_segmentation_image( + img, coco_dataset, label_id_by_category_id + ) + for img in image_rows + ] + + ml_use_rows = common_lib.create_ml_use_array_with_split( + len(image_rows), split_ratio + ) + return ml_use_rows, image_rows, segmentation_rows + + +def beam_build_tfrecord_from_coco_json( + input_json: str, + output_dir: str, + split_ratio: List[float], + num_shard_list: List[int], + output_shape: Optional[Tuple[int, int]] = None, +) -> None: + """Builds TFRecord files from COCO dataset. + + The output file names are `_TRAIN_TFRECORD_NAME`, `_VALIDATION_TFRECORD_NAME`, + and `_TEST_TFRECORD_NAME`. + + Args: + input_json: Path to a COCO JSON or JSONL file. + output_dir: Directory to output the TFRecord files. + split_ratio: List of how to split entries to train, validation, and test + TFRecords. + num_shard_list: List of the number of shards for each TFRecord file. + output_shape: The desired output shape of the image. If None, the original + image shape will be used. + """ + # `coco` cannot access gcs uri. Use gcsfuse, it is faster. + input_json = fileutils.force_gcs_fuse_path(input_json) + coco_dataset = coco.COCO(input_json) + + label_map = {} + label_id_by_category_id = {} + for idx, category in enumerate( + coco_dataset.dataset[constants.COCO_JSON_CATEGORIES], start=1 + ): + label_map[idx] = category[constants.COCO_JSON_CATEGORY_NAME] + label_id_by_category_id[category[constants.COCO_JSON_CATEGORY_ID]] = idx + label_map_path = os.path.join(output_dir, common_lib.LABEL_MAP_NAME) + logging.info('Writing label map to %s.', label_map_path) + common_lib.write_label_map(label_map_path, label_map) + + with tf.io.gfile.GFile( + os.path.join(output_dir, 'label_id_by_category_id.yaml'), 'w' + ) as f: + yaml.dump(label_id_by_category_id, f) + + ml_use_rows, image_rows, segmentation_rows = get_input_rows( + coco_dataset, split_ratio, label_id_by_category_id + ) + + def pipeline(root): + _define_data_conversion_pipeline( + root, + ml_use_rows, + image_rows, + segmentation_rows, + output_dir, + output_shape, + num_shard_list, + ) + + logging.info('Beginning beam pipeline to acquire tfrecords.') + common_lib.run_beam_pipeline(pipeline) diff --git a/community-content/vertex_vision_model_garden/model_oss/data_converter/data_converter_main.py b/community-content/vertex_vision_model_garden/model_oss/data_converter/data_converter_main.py new file mode 100644 index 000000000..fd7560bea --- /dev/null +++ b/community-content/vertex_vision_model_garden/model_oss/data_converter/data_converter_main.py @@ -0,0 +1,166 @@ +r"""Python script to convert user input data to training docker format. + + +Note: the training format is designed to be tfrecord as in the design doc. +If there are training efficiency issues for pytorch algorithms, we will also +support pytorch formats as well. +""" + +from absl import app +from absl import flags +from absl import logging +from data_converter import common_lib +from data_converter import data_converter_icn_lib +from data_converter import data_converter_iod_lib +from data_converter import data_converter_isg_lib +from data_converter import data_converter_vcn_lib +from util import constants + + +_INPUT_FILE_PATH = flags.DEFINE_string( + 'input_file_path', + None, + 'Input file path.', + required=True, +) +_INPUT_FILE_TYPE = flags.DEFINE_enum( + 'input_file_type', + None, + [ + constants.INPUT_FILE_TYPE_CSV, + constants.INPUT_FILE_TYPE_JSONL, + constants.INPUT_FILE_TYPE_COCO_JSON, + ], + 'Input file type.', + required=True, +) +_OBJECTIVE = flags.DEFINE_enum( + 'objective', + None, + [ + constants.OBJECTIVE_IMAGE_CLASSIFICATION, + constants.OBJECTIVE_IMAGE_OBJECT_DETECTION, + constants.OBJECTIVE_IMAGE_SEGMENTATION, + constants.OBJECTIVE_VIDEO_CLASSIFICATION, + ], + 'The objective of this training job.', + required=True, +) +_OUTPUT_DIR = flags.DEFINE_string( + 'output_dir', + None, + 'The output directory for converted data and label map files.', + required=True, +) +_SPLIT_RATIO = flags.DEFINE_list( + 'split_ratio', + '0.8,0.1,0.1', + 'Proportion of data to split into train/validation/test.', +) +_NUM_SHARD = flags.DEFINE_list( + 'num_shard', '10,10,10', 'The number of shards for train/validation/test.' +) +_OUTPUT_FPS = flags.DEFINE_integer( + 'output_fps', 5, 'For videos only. The output frames rate per second.' +) + + +def main(_) -> None: + logging.info( + ( + 'Start data converter on: %s (type: %s) with split: %s for %s' + ' (shard=%s), and output to %s.' + ), + _INPUT_FILE_PATH.value, + _INPUT_FILE_TYPE.value, + _SPLIT_RATIO.value, + _OBJECTIVE.value, + _NUM_SHARD.value, + _OUTPUT_DIR.value, + ) + split_ratio = list(map(float, _SPLIT_RATIO.value)) + num_shard = list(map(int, _NUM_SHARD.value)) + common_lib.check_split_ratio(split_ratio) + common_lib.check_num_shard(num_shard) + if ( + _OBJECTIVE.value == constants.OBJECTIVE_IMAGE_OBJECT_DETECTION + and _INPUT_FILE_TYPE.value == constants.INPUT_FILE_TYPE_CSV + ): + data_converter_iod_lib.convert_csv_to_tfrecord( + _INPUT_FILE_PATH.value, + _OUTPUT_DIR.value, + split_ratio, + num_shard, + ) + elif ( + _OBJECTIVE.value == constants.OBJECTIVE_IMAGE_OBJECT_DETECTION + and _INPUT_FILE_TYPE.value == constants.INPUT_FILE_TYPE_JSONL + ): + data_converter_iod_lib.convert_jsonl_to_tfrecord( + _INPUT_FILE_PATH.value, _OUTPUT_DIR.value, split_ratio, num_shard + ) + elif ( + _OBJECTIVE.value == constants.OBJECTIVE_IMAGE_OBJECT_DETECTION + and _INPUT_FILE_TYPE.value == constants.INPUT_FILE_TYPE_COCO_JSON + ): + data_converter_iod_lib.convert_coco_json_to_tfrecord( + _INPUT_FILE_PATH.value, + _OUTPUT_DIR.value, + split_ratio, + num_shard, + ) + elif _OBJECTIVE.value == constants.OBJECTIVE_IMAGE_SEGMENTATION: + data_converter_isg_lib.beam_build_tfrecord_from_coco_json( + _INPUT_FILE_PATH.value, + _OUTPUT_DIR.value, + split_ratio, + num_shard, + ) + elif ( + _OBJECTIVE.value == constants.OBJECTIVE_IMAGE_CLASSIFICATION + and _INPUT_FILE_TYPE.value == constants.INPUT_FILE_TYPE_CSV + ): + data_converter_icn_lib.convert_csv_to_tfrecord( + _INPUT_FILE_PATH.value, + _OUTPUT_DIR.value, + split_ratio, + num_shard, + ) + elif ( + _OBJECTIVE.value == constants.OBJECTIVE_IMAGE_CLASSIFICATION + and _INPUT_FILE_TYPE.value == constants.INPUT_FILE_TYPE_JSONL + ): + data_converter_icn_lib.convert_jsonl_to_tfrecord( + _INPUT_FILE_PATH.value, _OUTPUT_DIR.value, split_ratio, num_shard + ) + elif ( + _OBJECTIVE.value == constants.OBJECTIVE_VIDEO_CLASSIFICATION + and _INPUT_FILE_TYPE.value == constants.INPUT_FILE_TYPE_CSV + ): + data_converter_vcn_lib.convert_csv_to_tfrecord( + _INPUT_FILE_PATH.value, + _OUTPUT_DIR.value, + _OUTPUT_FPS.value, + split_ratio, + num_shard, + ) + elif ( + _OBJECTIVE.value == constants.OBJECTIVE_VIDEO_CLASSIFICATION + and _INPUT_FILE_TYPE.value == constants.INPUT_FILE_TYPE_JSONL + ): + data_converter_vcn_lib.convert_jsonl_to_tfrecord( + _INPUT_FILE_PATH.value, + _OUTPUT_DIR.value, + _OUTPUT_FPS.value, + split_ratio, + num_shard, + ) + else: + raise NotImplementedError( + f'File format {_INPUT_FILE_TYPE.value} is not supported for' + f' {_OBJECTIVE.value}.' + ) + + +if __name__ == '__main__': + app.run(main) diff --git a/community-content/vertex_vision_model_garden/model_oss/data_converter/data_converter_vcn_lib.py b/community-content/vertex_vision_model_garden/model_oss/data_converter/data_converter_vcn_lib.py new file mode 100644 index 000000000..c791b212d --- /dev/null +++ b/community-content/vertex_vision_model_garden/model_oss/data_converter/data_converter_vcn_lib.py @@ -0,0 +1,289 @@ +"""Converts VCN CSV/JSONL files to TFRecord with apache beam.""" + +import json +from os import path +from typing import Any, Dict, Iterator, Sequence, Union, cast + +from absl import logging +import apache_beam as beam +from apache_beam.io import tfrecordio +import numpy as np +import pandas as pd +import tensorflow as tf + +from data_converter import common_lib +from util import constants + + +_COLUMN_NAMES = [ + common_lib.COLUMN_NAME_ML_USE, + common_lib.COLUMN_NAME_GCS_FILE_PATH, + common_lib.COLUMN_NAME_LABEL, + common_lib.COLUMN_NAME_START_SEC, + common_lib.COLUMN_NAME_END_SEC, +] +_JSON_GCS_URI_KEY = 'videoGcsUri' +_JSON_CLASS_ANNOTATION_KEY = 'timeSegmentAnnotations' +_JSON_CLASS_NAME_KEY = 'displayName' +_JSON_START_TIME_KEY = 'startTime' +_JSON_END_TIME_KEY = 'endTime' +_JSON_RESOURCE_LABEL_KEY = 'dataItemResourceLabels' +_JSON_ML_USE_KEY = 'aiplatform.googleapis.com/ml_use' + + +def build_tf_example( + video_uri: str, + label: int, + start_sec: float, + end_sec: float, + output_fps: int, +) -> tf.train.SequenceExample: + """Builds a TF Example from a video clip. + + Args: + video_uri: GCS URI to the video file. + label: Class label as an integer. + start_sec: Start timestamp of the video clip in seconds. + end_sec: End timestamp of the video clip in seconds. + output_fps: The output frame rate per second. + + Returns: + The created TF Example. + """ + frame_bytes = common_lib.encode_video( + video_uri, start_sec, end_sec, output_fps, image_format='jpg' + ) + seq_example = tf.train.SequenceExample() + seq_example.context.feature['clip/label/index'].int64_list.value[:] = [label] + for frame in frame_bytes: + seq_example.feature_lists.feature_list.get_or_create( + 'image/encoded' + ).feature.add().bytes_list.value[:] = [frame] + + return seq_example + + +class AcquireTFExampleDoFn(beam.DoFn): + """Beam DoFn to build TF Examples from a DataFrame row dict for VCN.""" + + def __init__(self, output_fps: int): + self._success_counter = beam.metrics.Metrics.counter( + self.__class__.__name__, 'Success' + ) + self._failure_counter = beam.metrics.Metrics.counter( + self.__class__.__name__, 'Failure' + ) + self._output_fps = output_fps + + def process( + self, element: Dict[str, Union[float, int, str]] + ) -> Iterator[tf.train.SequenceExample]: + ml_use: str = cast(str, element[common_lib.COLUMN_NAME_ML_USE]) + video_uri: str = cast(str, element[common_lib.COLUMN_NAME_GCS_FILE_PATH]) + + try: + label: int = int(element[common_lib.COLUMN_NAME_LABEL]) + start_sec: float = float(element[common_lib.COLUMN_NAME_START_SEC]) + end_sec: float = float(element[common_lib.COLUMN_NAME_END_SEC]) + + tf_example = build_tf_example( + video_uri, + label, + start_sec, + end_sec, + self._output_fps, + ) + self._success_counter.inc() + yield beam.pvalue.TaggedOutput(ml_use, tf_example) + except (ValueError, IOError) as err: + logging.error('Failed to process %s', video_uri) + logging.exception(err) + self._failure_counter.inc() + + +def _run_convert_pipeline( + output_dir: str, + df: pd.DataFrame, + num_shards: Sequence[int], + output_fps: int, +) -> None: + """Starts a Beam pipeline to write DataFrame as TF Records. + + Args: + output_dir: TF Records output directory. + df: DataFrame to convert from. + num_shards: Number of shards for train/validation/test TFRecord files. + output_fps: The output frame rate per second. + """ + clip_list = df.to_dict('records') + + def pipeline(root): + train, val, test = ( + root + | 'Create PCollection' >> beam.Create(clip_list) + | 'Convert to TF Example' + >> beam.ParDo(AcquireTFExampleDoFn(output_fps)).with_outputs( + constants.ML_USE_TRAINING, + constants.ML_USE_VALIDATION, + constants.ML_USE_TEST, + ) + ) + _ = train | 'Save train TF Record' >> tfrecordio.WriteToTFRecord( + path.join(output_dir, common_lib.TRAIN_TFRECORD_NAME), + coder=beam.coders.ProtoCoder(tf.train.Example), + num_shards=num_shards[0], + ) + _ = val | 'Save val TF Record' >> tfrecordio.WriteToTFRecord( + path.join(output_dir, common_lib.VALIDATION_TFRECORD_NAME), + coder=beam.coders.ProtoCoder(tf.train.Example), + num_shards=num_shards[1], + ) + _ = test | 'Save test TF Record' >> tfrecordio.WriteToTFRecord( + path.join(output_dir, common_lib.TEST_TFRECORD_NAME), + coder=beam.coders.ProtoCoder(tf.train.Example), + num_shards=num_shards[2], + ) + + common_lib.run_beam_pipeline(pipeline) + + +def _convert_df_to_tfrecord( + df: pd.DataFrame, + output_dir: str, + split_ratio: Sequence[float], + num_shard: Sequence[int], + output_fps: int, +) -> None: + """Converts a DataFrame into three separate tfrecords for training, validation, and testing into output_dir. + + Args: + df: DataFrame to convert. + output_dir: The directory to save TFRecords and label_map.yaml. + split_ratio: List specifying the training, validation, and testing splits + for unassigned TFRecords. + num_shard: Number of shards for train/validation/test TFRecord files. + output_fps: The output frame rate per second. + """ + # Replaces ml_use with common_lib string constants for consistency. + common_lib.format_ml_use_column(df) + common_lib.insert_missing_ml_use(df) + + # Ignores invalid rows. + dropped_row_num = common_lib.drop_invalid_rows(df) + if dropped_row_num > 0: + logging.warning('Ignored %d invalid rows.', dropped_row_num) + + common_lib.replace_unassigned_ml_use( + df[common_lib.COLUMN_NAME_ML_USE], split_ratio + ) + + # Converts labels to integers as required by training. + new_labels, label_map = common_lib.create_label_map( + df[common_lib.COLUMN_NAME_LABEL] + ) + df[common_lib.COLUMN_NAME_LABEL] = new_labels + label_map_path = path.join(output_dir, common_lib.LABEL_MAP_NAME) + logging.info('Writing label map to %s.', label_map_path) + common_lib.write_label_map(label_map_path, label_map) + + # Missing start / end times are treated as 0, inf, respectively. + df[common_lib.COLUMN_NAME_START_SEC].fillna(0, inplace=True) + df[common_lib.COLUMN_NAME_END_SEC].fillna(np.inf, inplace=True) + + _run_convert_pipeline(output_dir, df, num_shard, output_fps) + + +def convert_csv_to_tfrecord( + input_csv: str, + output_dir: str, + output_fps: int, + split_ratio: Sequence[float] = (0.8, 0.1, 0.1), + num_shard: Sequence[int] = (10, 10, 10), +) -> None: + """Parses input_csv file into three separate tfrecords for training, validation, and testing into output_dir. + + The csv format is shown in + https://cloud.google.com/vertex-ai/docs/video-data/classification/prepare-data#csv + + If an ml_use column is not provided, one will be created. + + label_map.yaml containing the label map will be placed in output_dir. + + Args: + input_csv: Name of the csv file. + output_dir: The directory to save TFRecords and label_map.yaml. + output_fps: The output frame rate per second. + split_ratio: List specifying the training, validation, and testing splits + for unassigned TFRecords. + num_shard: Number of shards for train/validation/test TFRecord files. + """ + with tf.io.gfile.GFile(input_csv, 'r') as f: + df: pd.DataFrame = pd.read_csv( + f, header=None, names=_COLUMN_NAMES, on_bad_lines='warn' + ) + + _convert_df_to_tfrecord(df, output_dir, split_ratio, num_shard, output_fps) + + +def convert_jsonl_to_tfrecord( + input_jsonl: str, + output_dir: str, + output_fps: int, + split_ratio: Sequence[float] = (0.8, 0.1, 0.1), + num_shard: Sequence[int] = (10, 10, 10), +) -> None: + """Parses input_jsonl file into three separate tfrecords for training, validation, and testing into output_dir. + + The JSONL format is shown in + https://cloud.google.com/vertex-ai/docs/video-data/classification/prepare-data#jsonl. + + If an ml_use column is not provided, one will be created. + + label_map.yaml containing the label map will be placed in output_dir. + + Args: + input_jsonl: Name of the JSONL file. + output_dir: The directory to save TFRecords and label_map.yaml. + output_fps: The output frame rate per second. + split_ratio: List specifying the training, validation, and testing splits + for unassigned TFRecords. + num_shard: Number of shards for train/validation/test TFRecord files. + """ + df_rows = [] + with tf.io.gfile.GFile(input_jsonl, 'r') as f: + lines = f.read().rstrip().splitlines() + + for i, line in enumerate(lines, 1): + try: + item: Dict[str, Any] = json.loads(line) + + gcs_uri = item.get(_JSON_GCS_URI_KEY) + if not gcs_uri: + logging.warning('Invalid JSON at line %d, skipped.', i) + continue + + annotations = item.get(_JSON_CLASS_ANNOTATION_KEY, []) + ml_use = item.get(_JSON_RESOURCE_LABEL_KEY, {}).get( + _JSON_ML_USE_KEY, common_lib.ML_USE_UNASSIGNED + ) + + for j, annotation in enumerate(annotations): + label = annotation.get(_JSON_CLASS_NAME_KEY) + if not label: + logging.warning('Invalid annotation #%d at line %d, skipped.', j, i) + continue + # The example in external documentation uses strings like "1.0s", so we + # need to remove the "s" suffix. + start_time = annotation.get(_JSON_START_TIME_KEY, '0').removesuffix('s') + end_time = annotation.get(_JSON_END_TIME_KEY, 'inf').removesuffix('s') + df_rows.append([ml_use, gcs_uri, label, start_time, end_time]) + except (json.JSONDecodeError, AttributeError): + logging.warning('Invalid JSON at line %d, skipped.', i) + continue + + df = pd.DataFrame( + data=df_rows, + columns=_COLUMN_NAMES, + ) + + _convert_df_to_tfrecord(df, output_dir, split_ratio, num_shard, output_fps) diff --git a/community-content/vertex_vision_model_garden/model_oss/data_converter/dockerfile/base.Dockerfile b/community-content/vertex_vision_model_garden/model_oss/data_converter/dockerfile/base.Dockerfile new file mode 100644 index 000000000..d8cf82084 --- /dev/null +++ b/community-content/vertex_vision_model_garden/model_oss/data_converter/dockerfile/base.Dockerfile @@ -0,0 +1,50 @@ +FROM python:3.9 + +ENV DEBIAN_FRONTEND=noninteractive + +# Install basic libs. +RUN apt-get update && apt-get install -y --no-install-recommends \ + cmake \ + curl \ + wget \ + sudo \ + gnupg \ + python3-opencv \ + lsb-release \ + ca-certificates \ + build-essential \ + git \ + vim \ + screen \ + libportaudio2 \ + libusb-1.0-0-dev \ + openjdk-17-jre + +# Add gcsfuse distribution URL as a package source and import its public key. +RUN echo "deb https://packages.cloud.google.com/apt gcsfuse-`lsb_release -c -s` main" | sudo tee /etc/apt/sources.list.d/gcsfuse.list +RUN curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | sudo apt-key add - + +# Install gcsfuse. +RUN apt-get update && apt-get install -y --no-install-recommends gcsfuse + +# Install google cloud SDK. +RUN wget -q https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/google-cloud-sdk-359.0.0-linux-x86_64.tar.gz +RUN tar xzf google-cloud-sdk-359.0.0-linux-x86_64.tar.gz +RUN ./google-cloud-sdk/install.sh -q +# Make sure gsutil will use the default service account. +RUN echo '[GoogleCompute]\nservice_account = default' > /etc/boto.cfg + + +# Install required libs. +RUN pip install --upgrade pip +RUN pip install pyyaml==5.4.1 +RUN pip install pycocotools==2.0.6 +RUN pip install opencv-python-headless==4.7.0.72 +RUN pip install numpy==1.24.2 +RUN pip install pandas==1.5.3 +RUN pip install Pillow==9.4.0 +RUN pip install apache-beam[gcp]==2.45.0 +RUN pip install object-detection==0.0.3 +RUN pip install google-cloud-storage==1.42.3 +RUN pip install gcsfs==2021.10.1 +RUN pip install pylint==2.17.2 diff --git a/community-content/vertex_vision_model_garden/model_oss/data_converter/dockerfile/run.Dockerfile b/community-content/vertex_vision_model_garden/model_oss/data_converter/dockerfile/run.Dockerfile new file mode 100644 index 000000000..55dbcc122 --- /dev/null +++ b/community-content/vertex_vision_model_garden/model_oss/data_converter/dockerfile/run.Dockerfile @@ -0,0 +1,23 @@ +FROM gcr.io/automl-migration-test/automl-vision-data-converter-base:latest + +# Copy license. +RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE + +COPY model_oss/data_converter /automl_vision/data_converter +COPY model_oss/util /automl_vision/util + +WORKDIR /automl_vision +ENV PYTHONPATH "${PYTHONPATH}:/automl_vision" + +# Run pylint to validate code. +COPY .pylintrc /automl_vision/.pylintrc +RUN find . -type f -name "*.py" | xargs pylint --rcfile=./.pylintrc --errors-only + +ENTRYPOINT ["python3","data_converter/data_converter_main.py"] + +CMD ["--input_file_path=YOUR_INPUT_FILE",\ +"--input_file_type=csv",\ +"--objective=iod",\ +"--output_dir=YOUR_OUTPUT_DIR",\ +"--num_shard=10,10,10",\ +"--split_ratio=0.8,0.1,0.1"] diff --git a/community-content/vertex_vision_model_garden/model_oss/movinet/dockerfile/base.Dockerfile b/community-content/vertex_vision_model_garden/model_oss/movinet/dockerfile/base.Dockerfile new file mode 100644 index 000000000..871e5a472 --- /dev/null +++ b/community-content/vertex_vision_model_garden/model_oss/movinet/dockerfile/base.Dockerfile @@ -0,0 +1,64 @@ +FROM tensorflow/build:2.12-python3.9 + +ENV DEBIAN_FRONTEND=noninteractive + +# This is added to fix docker build error related to Nvidia key update. +RUN rm -f /etc/apt/sources.list.d/cuda.list +RUN curl https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add - + +# Install basic libs. +RUN apt-get update && apt-get install -y --no-install-recommends \ + cmake \ + curl \ + wget \ + sudo \ + gnupg \ + libsm6 \ + libxext6 \ + libxrender-dev \ + lsb-release \ + ca-certificates \ + build-essential \ + git \ + vim \ + libtcmalloc-minimal4 + + +# Install google cloud CLI. +RUN wget -q https://dl.google.com/dl/cloudsdk/channels/rapid/downloads/google-cloud-cli-430.0.0-linux-x86.tar.gz +RUN tar xzf google-cloud-cli-430.0.0-linux-x86.tar.gz +RUN ./google-cloud-sdk/install.sh -q +# Make sure gsutil will use the default service account. +RUN echo '[GoogleCompute]\nservice_account = default' > /etc/boto.cfg + + +# Install required libs. +RUN pip install --upgrade pip +RUN pip install cloud-tpu-client==0.10 +RUN pip install pyyaml==6.0 +RUN pip install fsspec==2023.4.0 +RUN pip install gcsfs==2023.4.0 +RUN pip install tf-models-official==2.12.0 +RUN pip install cloudml-hypertune==0.1.0.dev6 +RUN pip install pylint==2.17.3 + +# Installs Reduction Server NCCL plugin. +RUN echo "deb https://packages.cloud.google.com/apt google-fast-socket main" | tee /etc/apt/sources.list.d/google-fast-socket.list \ +&& curl -s -L https://packages.cloud.google.com/apt/doc/apt-key.gpg | apt-key add - \ +&& apt update && apt install -y google-reduction-server + +ENV PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=cpp + +# Lower the memory fragmentation, and speed up the training. +# https://github.com/tensorflow/tensorflow/issues/44176#issuecomment-783768033 +ENV LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libtcmalloc_minimal.so.4 + +# Enable userspace DNS cache +ENV GCS_RESOLVE_REFRESH_SECS=60 +ENV GCS_REQUEST_CONNECTION_TIMEOUT_SECS=300 +ENV GCS_METADATA_REQUEST_TIMEOUT_SECS=300 +ENV GCS_READ_REQUEST_TIMEOUT_SECS=300 +ENV GCS_WRITE_REQUEST_TIMEOUT_SECS=600 +# Each opened GCS file takes GCS_READ_CACHE_BLOCK_SIZE_MB of RAM, reduce the +# value from the default 64MB to 8MB to decrease memory footprint. +ENV GCS_READ_CACHE_BLOCK_SIZE_MB=8 diff --git a/community-content/vertex_vision_model_garden/model_oss/movinet/dockerfile/model_export.Dockerfile b/community-content/vertex_vision_model_garden/model_oss/movinet/dockerfile/model_export.Dockerfile new file mode 100644 index 000000000..7a9ee5f26 --- /dev/null +++ b/community-content/vertex_vision_model_garden/model_oss/movinet/dockerfile/model_export.Dockerfile @@ -0,0 +1,13 @@ +FROM gcr.io/automl-migration-test/movinet-base:latest + +# Copy license. +RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE + +RUN wget https://raw.githubusercontent.com/tensorflow/models/954dd73bffd43174bd3ca26a4a34abebe4147570/official/projects/movinet/tools/export_saved_model.py \ + -O /usr/local/lib/python3.9/dist-packages/official/projects/movinet/tools/export_saved_model.py + +WORKDIR /automl_vision + +ENV PYTHONPATH "${PYTHONPATH}:/automl_vision/util" + +ENTRYPOINT ["python3", "-m", "official.projects.movinet.tools.export_saved_model"] diff --git a/community-content/vertex_vision_model_garden/model_oss/movinet/dockerfile/serve.Dockerfile b/community-content/vertex_vision_model_garden/model_oss/movinet/dockerfile/serve.Dockerfile new file mode 100644 index 000000000..f0c184ded --- /dev/null +++ b/community-content/vertex_vision_model_garden/model_oss/movinet/dockerfile/serve.Dockerfile @@ -0,0 +1,18 @@ +FROM gcr.io/automl-migration-test/movinet-base:latest + +# Copy license. +RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE + +RUN pip install flask==2.3.2 +RUN pip install waitress==2.1.2 + +RUN mkdir -p /automl_vision/movinet/serving +COPY model_oss/movinet/serving /automl_vision/movinet/serving +COPY model_oss/util /automl_vision/util + +WORKDIR /automl_vision + +ENV PYTHONPATH "${PYTHONPATH}:/automl_vision/util" + +ENTRYPOINT ["flask", "--app", "movinet.serving.serving_main", "run"] +CMD ["--host=0.0.0.0", "--port=8501"] diff --git a/community-content/vertex_vision_model_garden/model_oss/movinet/dockerfile/train.Dockerfile b/community-content/vertex_vision_model_garden/model_oss/movinet/dockerfile/train.Dockerfile new file mode 100644 index 000000000..c8e356c3e --- /dev/null +++ b/community-content/vertex_vision_model_garden/model_oss/movinet/dockerfile/train.Dockerfile @@ -0,0 +1,18 @@ +FROM gcr.io/automl-migration-test/movinet-base:latest + +# Copy license. +RUN wget https://raw.githubusercontent.com/GoogleCloudPlatform/vertex-ai-samples/main/LICENSE + +RUN mkdir -p /automl_vision/movinet +COPY model_oss/movinet/*.py /automl_vision/movinet/ +COPY model_oss/util /automl_vision/util + +WORKDIR /automl_vision + +ENV PYTHONPATH "${PYTHONPATH}:/automl_vision/util" + +# Run pylint to validate code. +COPY .pylintrc /automl_vision/.pylintrc +RUN find . -type f -name "*.py" | xargs pylint --rcfile=./.pylintrc --errors-only + +ENTRYPOINT ["python3","movinet/train.py"] diff --git a/community-content/vertex_vision_model_garden/model_oss/movinet/serving/serving_main.py b/community-content/vertex_vision_model_garden/model_oss/movinet/serving/serving_main.py new file mode 100644 index 000000000..00043f277 --- /dev/null +++ b/community-content/vertex_vision_model_garden/model_oss/movinet/serving/serving_main.py @@ -0,0 +1,142 @@ +"""Main executable for MoViNet online / batch predictions.""" + +from collections.abc import Sequence +import json +import os + +from absl import app +from absl import logging +import flask +import tensorflow as tf +import waitress + +from movinet.serving import video_serving_lib +from util import constants + + +flask_app = flask.Flask(__name__) +logging.set_verbosity(logging.INFO) + +movinet_model = None + +_BATCH_SIZE = int(os.environ.get('BATCH_SIZE', '1')) +_NUM_FRAMES = int(os.environ.get('NUM_FRAMES', '32')) +_FPS = float(os.environ.get('FPS', '5')) +_OVERLAP_FRAMES = int(os.environ.get('OVERLAP_FRAMES', '24')) +_OBJECTIVE = os.environ.get( + 'OBJECTIVE', constants.OBJECTIVE_VIDEO_CLASSIFICATION +).lower() + +# VAR parameters. +_CONFIDENCE_THRESHOLD = float(os.environ.get('CONFIDENCE_THRESHOLD', '0.5')) +_MIN_GAP_TIME = float(os.environ.get('MIN_GAP_TIME', '1.5')) + + +def load_movinet_model() -> None: + model_path = os.environ.get('MODEL_PATH') + + if not model_path: + raise app.UsageError('Missing MODEL_PATH environment variable.') + + # We just reload the weights of the fine-tuned diffusion model. + logging.info('Initialize finetuned models from: %s', model_path) + global movinet_model + movinet_model = tf.saved_model.load(model_path) + + +load_movinet_model() + + +def error(message: str) -> str: + """Returns a JSON representing an error response.""" + return json.dumps({ + 'success': False, + 'error': message, + }) + + +# The health check route is required for docker deployment in google cloud. +@flask_app.route('/ping') +def ping() -> flask.Response: + """Health checks.""" + return flask.Response(status=200) + + +# The return should be `Response` for docker deployment in google cloud. +@flask_app.route('/predict', methods=['GET', 'POST']) +def predict_model() -> flask.Response: + """Predictions.""" + if flask.request.method == 'POST': + contents = flask.request.get_json(force=True) + + logging.info('The input contents are: %s', contents) + instances = contents.get('instances', []) + + try: + predictions = [] + for instance in instances: + executor = video_serving_lib.parse_request(instance) + prediction = executor.get_prediction( + movinet_model, + _BATCH_SIZE, + _FPS, + _NUM_FRAMES, + _OVERLAP_FRAMES, + _OBJECTIVE, + ) + if _OBJECTIVE == constants.OBJECTIVE_VIDEO_CLASSIFICATION: + prediction = video_serving_lib.postprocess_vcn(prediction) + elif _OBJECTIVE == constants.OBJECTIVE_VIDEO_ACTION_RECOGNITION: + prediction = video_serving_lib.postprocess_var( + executor.windows, prediction, _CONFIDENCE_THRESHOLD, _MIN_GAP_TIME + ) + predictions.append(prediction) + except ValueError as e: + return flask.Response( + error(str(e)), status=500, mimetype='application/json' + ) + + return flask.Response( + response=json.dumps({ + 'success': True, + 'predictions': predictions, + }), + status=200, + mimetype='application/json', + ) + else: + return flask.Response( + response=json.dumps({ + 'success': True, + 'isalive': movinet_model is not None, + }), + status=200, + mimetype='application/json', + ) + + +def main(argv: Sequence[str]) -> None: + if len(argv) > 1: + raise app.UsageError('Too many command-line arguments.') + # This is used when running locally only. When deploying to Google App + # Engine, a webserver process such as Gunicorn will serve the app. + # # Debug deployment. + # flask_app.run(host='0.0.0.0', port=8501, debug=True) + # Prod deployment. + if _OBJECTIVE not in [ + constants.OBJECTIVE_VIDEO_CLASSIFICATION, + constants.OBJECTIVE_VIDEO_ACTION_RECOGNITION, + ]: + raise app.UsageError('Objective must be vcn or var.') + logging.info( + 'Env: batch_size: %s, num_frames: %s, fps: %s, overlap_frames: %s', + _BATCH_SIZE, + _NUM_FRAMES, + _FPS, + _OVERLAP_FRAMES, + ) + waitress.serve(flask_app, host='0.0.0.0', port=8501) + + +if __name__ == '__main__': + app.run(main) diff --git a/community-content/vertex_vision_model_garden/model_oss/movinet/serving/video_serving_lib.py b/community-content/vertex_vision_model_garden/model_oss/movinet/serving/video_serving_lib.py new file mode 100644 index 000000000..17230a5c0 --- /dev/null +++ b/community-content/vertex_vision_model_garden/model_oss/movinet/serving/video_serving_lib.py @@ -0,0 +1,462 @@ +"""Lib for handling video prediction requests. + +The VCN inference algorithm is as follows: +1. Find all video frames within the given clip according to the sampling FPS. +2. Create possibly overlapping sliding windows according to the num_frames and +overlap_frames parameters. The last window might have a larger overlap if it +doesn't exactly fit. +3. Run model inference on each sliding window and compute softmax to obtain +probabilities. +4. Average the probabilities over all sliding windows. + +The VAR inference algorithm is very similar to VCN, with a few differences: +1. The last sliding window is discarded if it does not exactly fit. +2. Instead of averaging, the postprocessing consists of temporal nonmaximal +suppression and removing background and low-confidence labels. +""" + +from __future__ import annotations + +import dataclasses +import os +from typing import Any, Dict, Optional, Sequence, Union, cast + +from absl import logging +import cv2 +import numpy as np +import tensorflow as tf + +from util import constants +from util import fileutils + + +_JSON_LABEL_KEY = 'label' +_JSON_GCS_URI_KEY = 'content' +_JSON_CONFIDENCE_KEY = 'confidence' +_JSON_START_TIME_KEY = 'timeSegmentStart' +_JSON_END_TIME_KEY = 'timeSegmentEnd' +_BACKGROUND_LABEL = 0 +_JSON_REQUIRED_KEYS = [ + _JSON_GCS_URI_KEY, + _JSON_START_TIME_KEY, + _JSON_END_TIME_KEY, +] +_IMAGE_WIDTH = int(os.environ.get('IMAGE_WIDTH', '172')) +_IMAGE_HEIGHT = int(os.environ.get('IMAGE_HEIGHT', '172')) + + +@dataclasses.dataclass +class DetectionOutput: + timestamp: float + label: int + confidence: float + + def to_json_obj(self) -> Dict[str, Union[int, float]]: + """Encodes self as a dict for JSON serialization.""" + return { + _JSON_LABEL_KEY: self.label, + _JSON_START_TIME_KEY: self.timestamp, + _JSON_END_TIME_KEY: self.timestamp, + _JSON_CONFIDENCE_KEY: self.confidence, + } + + +def create_detection_output( + timestamp: float, predictions: np.ndarray +) -> DetectionOutput: + label = np.argmax(predictions).item() + confidence: float = predictions[label].item() + return DetectionOutput(timestamp, label, confidence) + + +class SlidingWindow: + """Represents a sliding window with start / end timestamps.""" + + def __init__(self, fps: float, frames: Sequence[int]): + if not frames: + raise ValueError('Sliding window cannot be empty.') + self.frames = frames + self.start_time = frames[0] / fps + self.end_time = frames[-1] / fps + self.frame_data: list[Optional[np.ndarray]] = [] + self.clear_frame_data() + + def load_cache_from(self, other: SlidingWindow) -> int: + """Loads cache from another sliding window if possible.""" + cache_count = 0 + for i, frame in enumerate(self.frames): + try: + other_idx = other.frames.index(frame) + self.frame_data[i] = other.frame_data[other_idx] + cache_count += 1 + except ValueError: + # Cache miss. + pass + return cache_count + + def load_frames(self, video: Any) -> Sequence[np.ndarray]: + """Loads frames of this sliding window from a video.""" + for i, frame in enumerate(self.frames): + if self.frame_data[i] is None: + video.set(cv2.CAP_PROP_POS_FRAMES, frame) + ret, frame = video.read() + if not ret: + raise IOError(f'Failed to read video at frame {frame}.') + self.frame_data[i] = cv2.resize(frame, (_IMAGE_WIDTH, _IMAGE_HEIGHT)) + return cast(Sequence[np.ndarray], self.frame_data) + + def clear_frame_data(self) -> None: + """Clears frame data of this sliding window to reduce memory usage.""" + self.frame_data: list[Optional[np.ndarray]] = [None] * len(self) + + def __len__(self) -> int: + return len(self.frames) + + @property + def middle_timestamp(self) -> float: + return (self.start_time + self.end_time) / 2 + + +def _get_sliding_windows( + frames: Sequence[int], + original_fps: float, + window_size: int, + overlap: int, + flush_last_window: bool, +) -> Sequence[SlidingWindow]: + """Computes a list of sliding windows from frames. + + Args: + frames: A list of frame indices. + original_fps: Frames per second of the original video. + window_size: Number of frames in a single window. + overlap: Number of overlapping frames in adjacent windows. + flush_last_window: Where to flush the last window if there are not enough + frames left. + + Returns: + A list of sliding windows, each has a list of frame indices. The last two + windows might have a larger overlap if the last window does not exactly fit + and flush_last_window is set to True. + + Raises: + ValueError: Arguments are invalid. + """ + if window_size <= overlap: + raise ValueError(f'Window size {window_size} <= overlap {overlap}') + total_frames = len(frames) + windows: list[SlidingWindow] = [] + for i in range(0, total_frames, window_size - overlap): + if i == 0 or i + window_size <= total_frames: + windows.append(SlidingWindow(original_fps, frames[i : i + window_size])) + elif i + overlap < total_frames and flush_last_window: + # Some frames in this window are not covered by the previous window. + windows.append( + SlidingWindow( + original_fps, frames[total_frames - window_size : total_frames] + ) + ) + return windows + + +def _sample_frame_indices( + start_time: float, + end_time: float, + original_fps: float, + sample_fps: float, + max_frames: int, + padding_left: int = 0, + padding_right: int = 0, +) -> Sequence[int]: + """Samples frames from start_time to end_time by sample_fps. + + Args: + start_time: Start timestamp in seconds. + end_time: End timestamp in seconds. + original_fps: Frames per second of the original video. + sample_fps: Number of frames to sample per second. + max_frames: Total number of frames in the video. + padding_left: Padding to add to the start in frames. Padded frames will be + duplicates of the first frame. + padding_right: Padding to add to the end in frames. Padded frames will be + duplicates of the last frame. + + Returns: + A list of sampled frame indices. + """ + ret = [ + min(max_frames - 1, round(t * original_fps)) + for t in np.arange(start_time, end_time, 1 / sample_fps) + ] + if ret: + ret = [ret[0]] * padding_left + ret + [ret[-1]] * padding_right + return ret + + +class VideoPredictionExecutor: + """Represents a Video prediction request with a video clip.""" + + def __init__(self, gcs_uri: str, start_time: float, end_time: float): + self._gcs_uri = gcs_uri + self._start_time = start_time + self._end_time = end_time + self.windows: Sequence[SlidingWindow] = [] + self._last_window: SlidingWindow = None + + def _read_frames_from_window( + self, video: Any, new_window: SlidingWindow + ) -> Sequence[np.ndarray]: + """Reads video frames from the new window. + + Args: + video: Video loaded with cv2. + new_window: A list of sorted frame indices in the new window. + + Returns: + Frame data from the video as a list of numpy arrays. + + Raises: + IOError: Failed to read video. + """ + # Caches frames as much as possible. + if self._last_window is not None: + cache_count = new_window.load_cache_from(self._last_window) + logging.info('Cached %d frames.', cache_count) + self._last_window.clear_frame_data() + self._last_window = new_window + return new_window.load_frames(video) + + def _predict( + self, model: Any, video: Any, batched_windows: Sequence[SlidingWindow] + ) -> np.ndarray: + """Run model inference on specific frames of a video. + + Args: + model: MoViNet model. + video: Video loaded with cv2. + batched_windows: A batch of sliding windows to predict. Each element is an + integer frame index. Must have equal number of frames in each window. + + Returns: + Prediction results. + + Raises: + ValueError: Batched windows are not sorted, or do not have equal number of + frames in each window. + IOError: Failed to read video. + """ + if any( + ( + len(window) != len(batched_windows[0]) + for window in batched_windows[1:] + ) + ): + raise ValueError( + 'Batched windows do not have equal number of frames in each window.' + ) + batch = [] + logging.info('Loading video frames...') + for window in batched_windows: + logging.info('Predict frames: %s', window.frames) + frames = self._read_frames_from_window(video, window) + batch.append(frames) + input_tensor = tf.convert_to_tensor(batch, dtype=tf.float32) / 255.0 + logging.info('Predict: Input tensor shape %s', input_tensor.shape) + predictions = model({'image': input_tensor}) + logging.info('Running softmax on predictions...') + predictions = tf.nn.softmax(predictions, axis=1) + return predictions.numpy() + + def get_prediction( + self, + model: Any, + batch_size: int, + fps: float, + num_frames: int, + overlap_frames: int, + objective: str, + ) -> Sequence[np.ndarray]: + """Predicts the video clip with the model. + + Args: + model: The loaded MoViNet model. + batch_size: Batch size for prediction. + fps: Video sampling FPS. + num_frames: Number of frames in a single predictions. If the model is + exported with a fixed input shape, this must match its num_frames + dimension. + overlap_frames: Number of overlapping frames of consecutive sliding + windows. + objective: A string `vcn` or `var`. + + Returns: + A list of floats as the prediction response. + + Raises: + IOError: The video fails to load. + ValueError: Some arguments are invalid. + """ + if objective not in [ + constants.OBJECTIVE_VIDEO_CLASSIFICATION, + constants.OBJECTIVE_VIDEO_ACTION_RECOGNITION, + ]: + raise ValueError(f'{objective} objective is not supported.') + + # cv2 expects a local path so we need to download the video from GCS. + local_file_path = fileutils.generate_tmp_path( + os.path.splitext(self._gcs_uri)[1] + ) + logging.info('Downloading %s to %s...', self._gcs_uri, local_file_path) + fileutils.download_gcs_file_to_local(self._gcs_uri, local_file_path) + logging.info('Download %s complete.', self._gcs_uri) + + # Loads video. + video = cv2.VideoCapture(local_file_path) + total_frames = video.get(cv2.CAP_PROP_FRAME_COUNT) + original_fps = video.get(cv2.CAP_PROP_FPS) + if not original_fps: + # 0 or None indicates the video is invalid. + raise IOError(f'Failed to load {self._gcs_uri}.') + video_length = total_frames / original_fps + self._start_time = max(0, self._start_time) + self._end_time = min(video_length, self._end_time) + padding = ( + (num_frames // 2) + if objective == constants.OBJECTIVE_VIDEO_ACTION_RECOGNITION + else 0 + ) + + # Computes sliding windows. + frame_indices = _sample_frame_indices( + self._start_time, + self._end_time, + original_fps, + fps, + total_frames, + padding, + padding, + ) + logging.info('Frame indices: %s', frame_indices) + self.windows = _get_sliding_windows( + frame_indices, + original_fps, + num_frames, + overlap_frames, + objective != 'var', + ) + if not self.windows: + raise ValueError( + f'No sliding windows found from {self._start_time} to' + f' {self._end_time}.' + ) + self._last_window = None + + # Runs inference. + predictions = [] + for i in range(0, len(self.windows), batch_size): + predictions.extend( + self._predict(model, video, self.windows[i : i + batch_size]) + ) + return predictions + + +def parse_request(req_json: Any) -> VideoPredictionExecutor: + """Parses VideoPredictionExecutor from request JSON object. + + Args: + req_json: Request JSON object. + + Returns: + Parsed VideoPredictionExecutor. + + Raises: + ValueError: Request JSON object is invalid. + """ + for key in _JSON_REQUIRED_KEYS: + if key not in req_json: + raise ValueError(f'{key} not found in {req_json}.') + gcs_uri = req_json[_JSON_GCS_URI_KEY] + start_time = float(req_json[_JSON_START_TIME_KEY].removesuffix('s')) + end_time = float(req_json[_JSON_END_TIME_KEY].removesuffix('s')) + return VideoPredictionExecutor(gcs_uri, start_time, end_time) + + +def postprocess_vcn(predictions: Sequence[np.ndarray]) -> Sequence[float]: + """Aggregates VCN predictions of sliding windows.""" + return np.mean(predictions, axis=0).tolist() + + +def temporal_nonmaximal_suppression( + detections: Sequence[DetectionOutput], min_gap_time: float +) -> Sequence[DetectionOutput]: + """Nonmaximal suppression for key frame detection. + + For consecutive packets of the same label within a pre-defined duration, we + only keep the one with the highest confidence score. Such duration can be + determined by performing data analysis on users' dataset. + + Args: + detections: A list of DetectionOutputs. + min_gap_time: Minimum time between consecutive key frames of the same label + in seconds. + + Returns: + DetectionOutput after nonmaximal suppression sorted in ascending timestamps. + """ + max_label = max([detection.label for detection in detections]) + prev_detections: list[Optional[DetectionOutput]] = [None] * (max_label + 1) + ret: list[DetectionOutput] = [] + by_time = lambda x: x.timestamp + for detection in sorted(detections, key=by_time): + prev_detection = prev_detections[detection.label] + prev_detections[detection.label] = detection + if not prev_detection: + continue + if detection.timestamp - prev_detection.timestamp > min_gap_time: + ret.append(prev_detection) + continue + detection.confidence = max(detection.confidence, prev_detection.confidence) + ret.extend((d for d in prev_detections if d is not None)) + return sorted(ret, key=by_time) + + +def postprocess_var( + windows: Sequence[SlidingWindow], + predictions: Sequence[np.ndarray], + confidence_threshold: float, + min_gap_time: float, +) -> Sequence[Dict[str, Any]]: + """Generates a list of detected keyframes from sliding window predictions. + + Args: + windows: Sliding windows. + predictions: A list of predictions of sliding windows. + confidence_threshold: Only probabilities greater than this threshold will + contribute to the final result. + min_gap_time: Minimum time between consecutive key frames of the same label + in seconds. Used in temporal nonmaximal suppression. + + Returns: + A sequence of dictionaries, each item has the following keys: + - label: Integer label of the detection result. + - timeSegmentStart: Start timestamp in seconds. + - timeSegmentEnd: End timestamp in seconds. Always equals timeSegmentStart. + """ + if len(windows) != len(predictions): + raise ValueError('Mismatched # of windows with # of predictions.') + + # Creates detection results from windows, filtering out the background label. + detections = [ + create_detection_output(window.middle_timestamp, predictions[i]) + for i, window in enumerate(windows) + ] + + # Temporal nonmaximal suppression. + detections = temporal_nonmaximal_suppression(detections, min_gap_time) + + # Filters out ones with low confidence and the background label. + return [ + x.to_json_obj() + for x in detections + if x.label != _BACKGROUND_LABEL and x.confidence > confidence_threshold + ] diff --git a/community-content/vertex_vision_model_garden/model_oss/movinet/train.py b/community-content/vertex_vision_model_garden/model_oss/movinet/train.py new file mode 100644 index 000000000..bc385ab43 --- /dev/null +++ b/community-content/vertex_vision_model_garden/model_oss/movinet/train.py @@ -0,0 +1,210 @@ +"""Main executable for MoViNet docker.""" + +import json +import os +from typing import Sequence, Any + +from absl import app +from absl import flags +from absl import logging +import gin +import hypertune +import tensorflow as tf + +from util import constants +from util import hypertune_utils +from official.common import distribute_utils +from official.common import flags as tfm_flags +from official.core import task_factory +from official.core import train_lib +from official.core import train_utils +from official.modeling import performance +# Import movinet libraries to register the backbone and model into tf.vision +# model garden factory. +# pylint: disable=unused-import +from official.projects.movinet.modeling import movinet +from official.projects.movinet.modeling import movinet_model +from official.vision import registry_imports +# pylint: enable=unused-import + + +FLAGS = flags.FLAGS + +_FILE_TYPE_TFRECORD = 'tfrecord' + +_LEARNING_RATE = flags.DEFINE_float( + 'learning_rate', None, 'The learning rate of this training job.' +) + +_NUM_CLASSES = flags.DEFINE_integer( + 'num_classes', None, 'The number of classes.' +) + +_INIT_CHECKPOINT = flags.DEFINE_string( + 'init_checkpoint', None, 'The initial checkpoint of this training job.' +) + +_INPUT_TRAIN_DATA_PATH = flags.DEFINE_string( + 'input_train_data_path', None, 'Input train data path.' +) + +_INPUT_VALIDATION_DATA_PATH = flags.DEFINE_string( + 'input_validation_data_path', None, 'Input validation data path.' +) + +_GLOBAL_BATCH_SIZE = flags.DEFINE_integer( + 'global_batch_size', None, 'Global batch size.' +) + +_PREFETCH_BUFFER_SIZE = flags.DEFINE_integer( + 'prefetch_buffer_size', None, 'Prefetch buffer size.' +) + +_SHUFFLE_BUFFER_SIZE = flags.DEFINE_integer( + 'shuffle_buffer_size', None, 'Shuffle buffer size.' +) + +_TRAIN_STEPS = flags.DEFINE_integer('train_steps', None, 'Train steps.') +_LOG_LEVEL = flags.DEFINE_enum( + 'log_level', + 'INFO', + ['FATAL', 'ERROR', 'WARNING', 'INFO', 'DEBUG'], + 'Log level.', +) + + +def parse_params() -> Any: + """Parses parameters.""" + gin.parse_config_files_and_bindings(FLAGS.gin_file, FLAGS.gin_params) + params = train_utils.parse_configuration(FLAGS, lock_return=False) + if _INIT_CHECKPOINT.value: + params.task.init_checkpoint = _INIT_CHECKPOINT.value + params.task.init_checkpoint_modules = 'backbone' + if _NUM_CLASSES.value: + params.task.model.num_classes = _NUM_CLASSES.value + params.task.train_data.num_classes = _NUM_CLASSES.value + params.task.validation_data.num_classes = _NUM_CLASSES.value + # If users set input train/validation data path, we assume the data are + # converted from data converter as tfrecord. Users can use tfds by writing + # their own config directly, and no need to override this parameter. + if _INPUT_TRAIN_DATA_PATH.value: + params.task.train_data.input_path = _INPUT_TRAIN_DATA_PATH.value + params.task.train_data.file_type = _FILE_TYPE_TFRECORD + params.task.train_data.tfds_name = '' + if _INPUT_VALIDATION_DATA_PATH.value: + params.task.validation_data.input_path = _INPUT_VALIDATION_DATA_PATH.value + params.task.validation_data.file_type = _FILE_TYPE_TFRECORD + params.task.validation_data.tfds_name = '' + if _GLOBAL_BATCH_SIZE.value: + params.task.train_data.global_batch_size = _GLOBAL_BATCH_SIZE.value + params.task.validation_data.global_batch_size = _GLOBAL_BATCH_SIZE.value + if _PREFETCH_BUFFER_SIZE.value: + params.task.train_data.prefetch_buffer_size = _PREFETCH_BUFFER_SIZE.value + params.task.validation_data.prefetch_buffer_size = ( + _PREFETCH_BUFFER_SIZE.value + ) + if _SHUFFLE_BUFFER_SIZE.value: + params.task.train_data.shuffle_buffer_size = _SHUFFLE_BUFFER_SIZE.value + if _TRAIN_STEPS.value: + params.trainer.train_steps = _TRAIN_STEPS.value + if _LEARNING_RATE.value: + logging.info('Updating learning_rate: %s', _LEARNING_RATE.value) + # Use `get` method of train_utils.hyperparams.OneOfConfig to get learning + # rate config. + learning_rate = params.trainer.optimizer_config.learning_rate.get() + if hasattr(learning_rate, 'initial_learning_rate'): + learning_rate.initial_learning_rate = _LEARNING_RATE.value + else: + logging.warning('Cannot set learning rate for %s', learning_rate) + # Set default params for best checkpoints. + params.trainer.best_checkpoint_export_subdir = constants.BEST_CKPT_DIRNAME + params.trainer.best_checkpoint_metric_comp = constants.BEST_CKPT_METRIC_COMP + params.trainer.best_checkpoint_eval_metric = ( + constants.VIDEO_CLASSIFICATION_BEST_EVAL_METRIC + ) + return params + + +def main(argv: Sequence[str]) -> None: + logging.set_verbosity(_LOG_LEVEL.value) + if len(argv) > 1: + raise app.UsageError('Too many command-line arguments.') + params = parse_params() + logging.info('The actual training parameters are:\n%s', params.as_dict()) + model_dir: str = os.path.join( + FLAGS.model_dir, + constants.TRIAL_PREFIX + hypertune_utils.get_trial_id_from_environment(), + ) + logging.info('model_dir: %s', model_dir) + + if 'train' in FLAGS.mode: + # Pure eval modes do not output yaml files. Otherwise continuous eval job + # may race against the train job for writing the same file. + train_utils.serialize_config(params, model_dir) + + # Sets mixed_precision policy. Using 'mixed_float16' or 'mixed_bfloat16' + # can have significant impact on model speeds by utilizing float16 in case of + # GPUs, and bfloat16 in the case of TPUs. loss_scale takes effect only when + # dtype is float16 + if params.runtime.mixed_precision_dtype: + performance.set_mixed_precision_policy(params.runtime.mixed_precision_dtype) + distribution_strategy = distribute_utils.get_distribution_strategy( + distribution_strategy=params.runtime.distribution_strategy, + all_reduce_alg=params.runtime.all_reduce_alg, + num_gpus=params.runtime.num_gpus, + tpu_address=params.runtime.tpu, + ) + + # Create task and run experiment. + with distribution_strategy.scope(): + task = task_factory.get_task(params.task, logging_dir=model_dir) + + train_lib.run_experiment( + distribution_strategy=distribution_strategy, + task=task, + mode=FLAGS.mode, + params=params, + model_dir=model_dir, + ) + + train_utils.save_gin_config(FLAGS.mode, model_dir) + + eval_metric_name = constants.VIDEO_CLASSIFICATION_BEST_EVAL_METRIC + + eval_filepath = os.path.join( + model_dir, constants.BEST_CKPT_DIRNAME, constants.BEST_CKPT_EVAL_FILENAME + ) + logging.info('Load eval metrics from: %s.', eval_filepath) + + with tf.io.gfile.GFile(eval_filepath, 'rb') as f: + eval_metric_results = json.load(f) + logging.info('eval metrics are: %s.', eval_metric_results) + if ( + eval_metric_name in eval_metric_results + and constants.BEST_CKPT_STEP_NAME in eval_metric_results + ): + hp_metric = eval_metric_results[eval_metric_name] + hp_step = int(eval_metric_results[constants.BEST_CKPT_STEP_NAME]) + hpt = hypertune.HyperTune() + hpt.report_hyperparameter_tuning_metric( + hyperparameter_metric_tag=constants.HP_METRIC_TAG, + metric_value=hp_metric, + global_step=hp_step, + ) + logging.info( + 'Send HP metric: %f and steps %d to hyperparameter tuning.', + hp_metric, + hp_step, + ) + else: + logging.info( + 'Either %s or %s is not included in the evaluation results: %s.', + eval_metric_name, + constants.BEST_CKPT_STEP_NAME, + eval_metric_results, + ) + + +if __name__ == '__main__': + tfm_flags.define_flags() + app.run(main) diff --git a/community-content/vertex_vision_model_garden/model_oss/pic2word/handler.py b/community-content/vertex_vision_model_garden/model_oss/pic2word/handler.py index 23f1a7d6c..0b33baeb7 100644 --- a/community-content/vertex_vision_model_garden/model_oss/pic2word/handler.py +++ b/community-content/vertex_vision_model_garden/model_oss/pic2word/handler.py @@ -15,7 +15,7 @@ import torch from torch.utils.data import DataLoader from ts.torch_handler.base_handler import BaseHandler -from google3.cloud.ml.applications.vision.model_garden.model_oss.util import fileutils +from util import fileutils # The COCO dataset is stored in a publicly accessible bucket. _COCO_STORAGE_DIR = "gs://pic2word-bucket/data/coco/" diff --git a/community-content/vertex_vision_model_garden/model_oss/util/commons.py b/community-content/vertex_vision_model_garden/model_oss/util/commons.py new file mode 100644 index 000000000..69f106f64 --- /dev/null +++ b/community-content/vertex_vision_model_garden/model_oss/util/commons.py @@ -0,0 +1,79 @@ +"""Common utility lib for prediction on images.""" + +from typing import Any, Dict, List + +import numpy as np +from PIL import Image +import tensorflow as tf +import yaml + +from util import image_format_converter + + +def get_prediction_instances(image: Image.Image) -> List[Dict[str, Any]]: + """Gets prediction instances. + + Args: + image: Image instance. + + Returns: + List[Dict[str, Any]]: List of prediction instances. + """ + instances = [{ + "encoded_image": {"b64": image_format_converter.image_to_base64(image)}, + }] + return instances + + +def get_label_map(label_map_yaml_filepath: str) -> Dict[str, Any]: + """Gets the label map from a YAML file. + + Args: + label_map_yaml_filepath: Filepath to the label map YAML file. + + Returns: + dict: Label map. + """ + with tf.io.gfile.GFile(label_map_yaml_filepath, "rb") as input_file: + label_map = yaml.safe_load(input_file.read()) + return label_map + + +def get_object_detection_endpoint_predictions( + detection_endpoint: ..., + input_image: np.ndarray, + detection_thresh: float = 0.2, +) -> np.ndarray: + """Gets endpoint predictions. + + Args: + detection_endpoint: image object detection endpoint. + input_image: Input image. + detection_thresh: Detection threshold. + + Returns: + Object detection predictions from endpoints. + """ + height, width, _ = input_image.shape + predictions = detection_endpoint.predict( + get_prediction_instances(Image.fromarray(input_image)) + ).predictions + detection_scores = np.array(predictions[0]["detection_scores"]) + detection_classes = np.array(predictions[0]["detection_classes"]) + detection_boxes = np.array( + [ + [b[1] * width, b[0] * height, b[3] * width, b[2] * height] + for b in predictions[0]["detection_boxes"] + ] + ) + thresh_indices = [ + x for x, val in enumerate(detection_scores) if val > detection_thresh + ] + preds_merge_conf = np.column_stack(( + detection_boxes[thresh_indices], + detection_scores[thresh_indices], + )) + preds_merge_cls = np.column_stack( + (preds_merge_conf, detection_classes[thresh_indices]) + ) + return preds_merge_cls diff --git a/community-content/vertex_vision_model_garden/model_oss/util/constants.py b/community-content/vertex_vision_model_garden/model_oss/util/constants.py index e9b910cfe..e8d2521e8 100644 --- a/community-content/vertex_vision_model_garden/model_oss/util/constants.py +++ b/community-content/vertex_vision_model_garden/model_oss/util/constants.py @@ -4,6 +4,8 @@ OBJECTIVE_IMAGE_CLASSIFICATION = 'icn' OBJECTIVE_IMAGE_OBJECT_DETECTION = 'iod' OBJECTIVE_IMAGE_SEGMENTATION = 'isg' +OBJECTIVE_VIDEO_CLASSIFICATION = 'vcn' +OBJECTIVE_VIDEO_ACTION_RECOGNITION = 'var' # Input file types. INPUT_FILE_TYPE_CSV = 'csv' @@ -61,4 +63,11 @@ GCSFUSE_URI_PREFIX = '/gcs/' LOCAL_EVALUATION_RESULT_DIR = '/tmp/evaluation_result_dir' LOCAL_MODEL_DIR = '/tmp/model_dir' +LOCAL_BASE_MODEL_DIR = '/tmp/base_model_dir' LOCAL_DATA_DIR = '/tmp/data' + +# PEFT finetuning constants. +TEXT_TO_IMAGE_LORA = 'text-to-image-lora' +SEQUENCE_CLASSIFICATION_LORA = 'sequence-classification-lora' +CAUSAL_LANGUAGE_MODELING_LORA = 'causal-language-modeling-lora' +INSTRUCT_LORA = 'instruct-lora' diff --git a/community-content/vertex_vision_model_garden/model_oss/util/fileutils.py b/community-content/vertex_vision_model_garden/model_oss/util/fileutils.py index 65d401c4f..eb4c23ea9 100644 --- a/community-content/vertex_vision_model_garden/model_oss/util/fileutils.py +++ b/community-content/vertex_vision_model_garden/model_oss/util/fileutils.py @@ -2,6 +2,10 @@ import glob import os +import pathlib +import shutil +from typing import Tuple +import uuid from absl import logging from google.cloud import storage @@ -9,6 +13,44 @@ from google.cloud import storage from util import constants +def generate_tmp_path(extension: str = '') -> str: + """Generates a temporary file path with UUID. + + Args: + extension: File extension, e.g. '.jpg', '.avi'. If not given, no extension + will be appended to the filename. + + Returns: + Generated file path. + """ + return os.path.join(constants.LOCAL_DATA_DIR, uuid.uuid1().hex) + extension + + +def force_gcs_fuse_path(gcs_uri: str) -> str: + """Converts gs:// uris to their /gcs/ equivalents. No-op for other uris.""" + if is_gcs_path(gcs_uri): + return ( + constants.GCSFUSE_URI_PREFIX + gcs_uri[len(constants.GCS_URI_PREFIX) :] + ) + else: + return gcs_uri + + +def download_gcs_file_to_local_dir(gcs_uri: str, local_dir: str): + """Download a gcs file to a local dir. + + Args: + gcs_uri: A string of file path on GCS. + local_dir: A string of local directory. + """ + if not is_gcs_path(gcs_uri): + raise ValueError( + f'{gcs_uri} is not a GCS path starting with {constants.GCS_URI_PREFIX}.' + ) + filename = os.path.basename(gcs_uri) + download_gcs_file_to_local(gcs_uri, os.path.join(local_dir, filename)) + + def download_gcs_file_to_local(gcs_uri: str, local_path: str): """Download a gcs file to a local path. @@ -16,7 +58,7 @@ def download_gcs_file_to_local(gcs_uri: str, local_path: str): gcs_uri: A string of file path on GCS. local_path: A string of local file path. """ - if not gcs_uri.startswith(constants.GCS_URI_PREFIX): + if not is_gcs_path(gcs_uri): raise ValueError( f'{gcs_uri} is not a GCS path starting with {constants.GCS_URI_PREFIX}.' ) @@ -38,6 +80,8 @@ def download_gcs_dir_to_local(gcs_dir: str, local_dir: str): gcs_dir: A string of directory path on GCS. local_dir: A string of local directory path. """ + if not is_gcs_path(gcs_dir): + raise ValueError(f'{gcs_dir} is not a GCS path starting with gs://.') bucket_name = gcs_dir.split('/')[2] prefix = gcs_dir[len(constants.GCS_URI_PREFIX + bucket_name) :].strip('/') client = storage.Client() @@ -77,3 +121,126 @@ def upload_local_dir_to_gcs(local_dir: str, gcs_dir: str): ) blob = bucket.blob(os.path.join(blob_dir, os.path.basename(local_file))) blob.upload_from_filename(local_file) + + +def upload_file_to_gcs_path( + source_path: str, + destination_uri: str, +): + """Uploads local files to GCS uri. + + After upload the destination_uri will contain the same data as the + source_path. + + Args: + source_path: Required. Path of the local data to copy to GCS. + destination_uri: Required. GCS URI where the data should be uploaded. + + Raises: + RuntimeError: When source_path does not exist. + GoogleCloudError: When the upload process fails. + """ + source_path_obj = pathlib.Path(source_path) + if not source_path_obj.exists(): + raise RuntimeError(f'Source path does not exist: {source_path}') + + storage_client = storage.Client() + source_file_path = source_path + destination_file_uri = destination_uri + logging.info('Uploading "%s" to "%s"', source_file_path, destination_file_uri) + destination_blob = storage.Blob.from_string( + destination_file_uri, client=storage_client + ) + destination_blob.upload_from_filename(filename=source_file_path) + + +def is_gcs_path(input_path: str) -> bool: + """Checks if the input path is a Google Cloud Storage (GCS) path. + + Args: + input_path: The input path to be checked. + + Returns: + True if the input path is a GCS path, False otherwise. + """ + return input_path.startswith(constants.GCS_URI_PREFIX) + + +def release_text_assets( + output_bucket: str, local_text_file_name: str, remote_text_file_name: str +) -> None: + """Releases text assets. + + Args: + output_bucket: gcs output bucket. + local_text_file_name: Local text file name. + remote_text_file_name: Remote text file name. + + Returns: + None + """ + remote_file_path = '{}/{}'.format(output_bucket, remote_text_file_name) + logging.info('Uploading "%s" to "%s"', local_text_file_name, remote_file_path) + upload_file_to_gcs_path(local_text_file_name, remote_file_path) + os.remove(local_text_file_name) + + +def upload_video_from_local_to_gcs( + output_bucket: str, + local_video_file_name: str, + remote_video_file_name: str, + temp_local_video_file_name: str, +) -> None: + """Uploads video from local to gcs buckent and releases video assets. + + Args: + output_bucket: GCS bucket address. + local_video_file_name: Local video file name. + remote_video_file_name: Remote video file name. + temp_local_video_file_name: Temporary local video file name. + + Returns: + None + """ + upload_file_to_gcs_path( + temp_local_video_file_name, + '{}/{}'.format(output_bucket, remote_video_file_name), + ) + shutil.rmtree(local_video_file_name, ignore_errors=True) + shutil.rmtree(temp_local_video_file_name, ignore_errors=True) + + +def download_video_from_gcs_to_local(video_file_path: str) -> Tuple[str, str]: + """Downloads video from gcs to local folders. + + Args: + video_file_path: Path to the video file. + + Returns: + Local and remote video file paths. + """ + _, local_video_file_name = os.path.split(video_file_path) + file_extension = os.path.splitext(video_file_path)[1] + remote_video_file_name = local_video_file_name.replace( + file_extension, '_overlay.mp4' + ) + local_file_path = generate_tmp_path(os.path.splitext(video_file_path)[1]) + logging.info('Downloading %s to %s...', video_file_path, local_file_path) + download_gcs_file_to_local(video_file_path, local_file_path) + return local_file_path, remote_video_file_name + + +def get_output_video_file(video_output_file_path: str) -> str: + """Gets the output video file name for writing video. + + Args: + video_output_file_path: Path to the video output file. + + Returns: + str: Local video output file path. + """ + file_extension = os.path.splitext(video_output_file_path)[1] + out_local_video_file_name = video_output_file_path.replace( + file_extension, '_overlay' + file_extension + ) + return out_local_video_file_name diff --git a/community-content/vertex_vision_model_garden/model_oss/util/hypertune_utils.py b/community-content/vertex_vision_model_garden/model_oss/util/hypertune_utils.py new file mode 100644 index 000000000..886ce161d --- /dev/null +++ b/community-content/vertex_vision_model_garden/model_oss/util/hypertune_utils.py @@ -0,0 +1,22 @@ +"""Utility functions for Vertex Hyperparameter Tuning Jobs.""" + +import os + +from absl import logging + + +_ENVIRONMENT_VARIABLE_FOR_TRIAL_ID = 'CLOUD_ML_TRIAL_ID' + + +def get_trial_id_from_environment() -> str: + """Gets the trial id from environment variable. + + Returns: + The trial id from environement or '0' if not found. + """ + if _ENVIRONMENT_VARIABLE_FOR_TRIAL_ID not in os.environ: + logging.warning( + 'Environment variable %s not found, return 0 as default trial id.', + _ENVIRONMENT_VARIABLE_FOR_TRIAL_ID, + ) + return os.environ.get(_ENVIRONMENT_VARIABLE_FOR_TRIAL_ID, '0')