Copyright 2022 The TensorFlow Authors.

python

#@title Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

MOSAIC with Model Garden

This tutorial demonstrates how to load the MOSAIC model trained using Tensorflow Model Garden library.

Tensorflow Model Garden contains a collection of state-of-the-art models, implemented with TensorFlow's high-level APIs. The implementations demonstrate the best practices for modeling, letting users to take full advantage of TensorFlow for their research and product development.

Install Necessary Dependencies

Import libraries

python

import matplotlib.pyplot as plt
import matplotlib.patches as mpatches

from PIL import Image
from six import BytesIO
from IPython import display
from urllib.request import urlopen

import numpy as np
import tensorflow as tf
import tensorflow_hub as hub

import absl.logging
absl.logging.set_verbosity(absl.logging.ERROR)

Load Pretrained Model

Note that you may also access the pretrained model provided on TFHub:

keras_layer = hub.KerasLayer(
  'https://tfhub.dev/google/mosaic/mobilenetmultiavgseg/2',
  signature='serving_default',
  output_key='logits')
model = tf.keras.Sequential([keras_layer])
model.build([None, IMAGE_HEIGHT, IMAGE_WIDTH, 3])

Download SavedModel

The model uses the implementation from the TensorFlow Model Garden GitHub repository, and achieves 77.24% mIoU on Cityscapes dataset with 19 classes.

python

! curl https://storage.googleapis.com/tf_model_garden/vision/mosaic/mosaic_mobilenet_multiavgseg_r1024_ebf64_gp_model.tar.gz --output model.tar.gz

python

# Extract savedmodel
! tar -xvf model.tar.gz

Load the SavedModel

python

# Load saved model
export_dir = "saved_model"
imported = tf.saved_model.load(export_dir)
model_fn = imported.signatures['serving_default']

Run Inference

python

# Defines helper function to download sample image
def load_image_into_numpy_array(path):
  """Load an image from file into a numpy array.

  Puts image into numpy array to feed into tensorflow graph.
  Note that by convention we put it into a numpy array with shape
  (height, width, channels), where channels=3 for RGB.

  Args:
    path: the file path to the image

  Returns:
    uint8 numpy array with shape (img_height, img_width, 3)
  """
  image = None
  if(path.startswith('http')):
    response = urlopen(path)
    image_data = response.read()
    image_data = BytesIO(image_data)
    image = Image.open(image_data)
  else:
    image_data = tf.io.gfile.GFile(path, 'rb').read()
    image = Image.open(BytesIO(image_data))

  (im_width, im_height) = image.size

  image = np.array(image.getdata()).reshape(
      (1, im_height, im_width, 3)).astype(np.uint8)

  return image

Download a Sample Image

python

image_path = "https://storage.googleapis.com/tf_model_garden/vision/mosaic/cityscape_sample.png"
image_array = load_image_into_numpy_array(image_path)
image_array.shape

python

Image.fromarray(image_array[0])

Run Inference on Sample Image

python

outputs = model_fn(inputs=image_array)

python

outputs['logits'].shape

Note that the output is in the shape of [batch_size, height, width, num_classes], which is the raw logits prediction output for each pixel.

python

detection_results = np.argmax(outputs['logits'][0], axis=-1)

python

fig, ax = plt.subplots(figsize=(12, 12))
ax.imshow(detection_results)

Model Training

Please check out the TF Model Garden MOSAIC implementation for model training.

Please check out the Image Classification Tutorial for fine-tuning models from the TensorFlow Model Garden package.