uvq/uvq_utils.py

"""Utility functions of the UVQ model.

Copyright 2022 Google LLC

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.
"""


import csv
import os
import subprocess
import tempfile

import numpy as np
from PIL import Image
import tensorflow as tf

from tensorflow.compat.v1 import gfile, logging, saved_model
from tensorflow.compat.v1 import ConfigProto, Session, Graph


# Explicitly set input size
VIDEO_HEIGHT = 720
VIDEO_WIDTH = 1280
VIDEO_FPS = 5
VIDEO_CHANNEL = 3

# Output feature size
DIM_HEIGHT_FEATURE = 16
DIM_WIDTH_FEATURE = 16
DIM_CHANNEL_FEATURE = 100

# ContentNet specs
INPUT_HEIGHT_CONTENT = 496
INPUT_WIDTH_CONTENT = 496
INPUT_CHANNEL_CONTENT = 3
DIM_LABEL_CONTENT = 3862

def extend_array(rgb, total_len):
  """Extends the byte array (or truncates) to be total_len"""
  missing = total_len - len(rgb)
  if missing < 0:
    rgb = rgb[0 : total_len]
  else:
    rgb.extend(bytearray(missing))
  return rgb


def load_video(filepath, video_length, transpose=False):
  """Load input video."""
  # Rotate video if requested
  if transpose:
    transpose_param = 'transpose=1,'
  else:
    transpose_param = ''

  # Sample at constant frame rate, and save as RGB24 (RGBRGB...)
  fd, temp_filename = tempfile.mkstemp()
  fd_small, temp_filename_small = tempfile.mkstemp()
  cmd = (
      'ffmpeg  -i %s -filter_complex '
      ' "[0:v]%sscale=w=%d:h=%d:flags=bicubic:force_original_aspect_ratio=1,'
      'pad=%d:%d:(ow-iw)/2:(oh-ih)/2,format=rgb24,split=2[out1][tmp],[tmp]scale=%d:%d:flags=bilinear[out2]"'
      ' -map [out1] -r %d -f rawvideo -pix_fmt rgb24 -y %s'
      ' -map [out2] -r %d -f rawvideo -pix_fmt rgb24 -y %s'
  ) % (filepath, transpose_param, VIDEO_WIDTH, VIDEO_HEIGHT,
       VIDEO_WIDTH, VIDEO_HEIGHT, INPUT_WIDTH_CONTENT, INPUT_HEIGHT_CONTENT,
       VIDEO_FPS, temp_filename, VIDEO_FPS, temp_filename_small)

  try:
    logging.info('Run with cmd:% s\n', cmd)
    subprocess.check_output(cmd, stderr=subprocess.STDOUT, shell=True)
  except subprocess.CalledProcessError as error:
    logging.fatal('Run with cmd: %s \n terminated with return code %s\n%s',
                  cmd, str(error.returncode), error.output)
    raise error

  # For video, the entire video is divided into 1s chunks in 5 fps
  with gfile.Open(temp_filename, 'rb') as rgb_file, gfile.Open(temp_filename_small, 'rb') as rgb_file_small:
    single_frame_size = VIDEO_WIDTH * VIDEO_HEIGHT * VIDEO_CHANNEL
    full_decode_size = video_length * VIDEO_FPS * single_frame_size
    assert rgb_file.size() >= single_frame_size,  f'Decoding failed to output a single frame: {rgb_file.size()} < {single_frame_size}'
    if rgb_file.size() < full_decode_size:
      logging.warn('Decoding may be truncated: %d bytes (%d frames) < %d bytes (%d frames),'
                   ' or video length (%ds) may be too incorrect',
                   rgb_file.size(), rgb_file.size() / single_frame_size,
                   full_decode_size, full_decode_size / single_frame_size,
                   video_length)

    rgb = extend_array(bytearray(rgb_file.read()), full_decode_size)
    rgb_small = extend_array(bytearray(rgb_file_small.read()),
                             video_length * VIDEO_FPS * INPUT_WIDTH_CONTENT * INPUT_HEIGHT_CONTENT * VIDEO_CHANNEL)
    video = (np.reshape(np.frombuffer(rgb, 'uint8'),
                        (video_length, int(VIDEO_FPS), VIDEO_HEIGHT, VIDEO_WIDTH, 3))/255.0 - 0.5) *2
    video_resized = (np.reshape(np.frombuffer(rgb_small, 'uint8'),
                                (video_length, int(VIDEO_FPS), INPUT_HEIGHT_CONTENT, INPUT_WIDTH_CONTENT, 3))/255.0 - 0.5) *2

  # Delete temp files
  os.close(fd)
  os.remove(temp_filename)
  os.close(fd_small)
  os.remove(temp_filename_small)
  logging.info('Load %s done successfully.', filepath)
  return video, video_resized


def generate_content_feature(video, model_path, input_node, output_nodes):
  """Extract features from ContentNet."""
  with Session(graph=Graph(), config=ConfigProto(
      allow_soft_placement=True, log_device_placement=False)) as sess:
    saved_model.loader.load(
        sess, [saved_model.tag_constants.SERVING], model_path)

    label = np.ndarray((video.shape[0], DIM_LABEL_CONTENT), np.float32)
    feature = np.ndarray((video.shape[0], DIM_HEIGHT_FEATURE,
                          DIM_WIDTH_FEATURE, DIM_CHANNEL_FEATURE), np.float32)

    patch = np.ndarray((1, INPUT_HEIGHT_CONTENT, INPUT_WIDTH_CONTENT,
                        INPUT_CHANNEL_CONTENT), np.float32)

    for k in range(video.shape[0]):
      patch[0, :] = video[k, 0, :, :, :]
      patch_feature, patch_label = sess.run(output_nodes,
                                            feed_dict={input_node: patch})

      feature[k, :, :, :] = patch_feature
      label[k, :] = patch_label
    return feature, label


def generate_subnet_feature(video, model_path, input_width, input_height,
                            input_fps, feature_width, feature_height,
                            feature_channel, label_dim, input_node,
                            output_nodes):
  """Extract features from CompresionNet or DistortionNet."""
  with Session(graph=Graph(), config=ConfigProto(
      allow_soft_placement=True, log_device_placement=False)) as sess:
    saved_model.loader.load(
        sess, [saved_model.tag_constants.SERVING], model_path)

    num_h = int(VIDEO_HEIGHT / input_height)
    num_w = int(VIDEO_WIDTH / input_width)

    label = np.ndarray((video.shape[0], num_h, num_w, label_dim), np.float32)
    feature = np.ndarray((video.shape[0], num_h * feature_height,
                          num_w * feature_width, feature_channel), np.float32)

    if input_fps == 1:
      patch = np.ndarray((1, input_height, input_width,
                          video.shape[-1]), np.float32)
    else:
      patch = np.ndarray((1, input_fps, input_height, input_width,
                          video.shape[-1]), np.float32)

    for k in range(video.shape[0]):
      for j in range(num_h):
        for i in range(num_w):
          if input_fps == 1:
            patch[0, :] = video[k, 0, j * input_height:(j + 1) * input_height,
                                i * input_width:(i + 1) * input_width, :]
          else:
            patch[0, :] = video[k, :, j * input_height:(j + 1) * input_height,
                                i * input_width:(i + 1) * input_width, :]

          patch_feature, patch_label = sess.run(output_nodes,
                                                feed_dict={input_node: patch})

          feature[k, j * feature_height:(j + 1) * feature_height, i *
                  feature_width:(i + 1) * feature_width, :] = patch_feature

          label[k, j, i, :] = patch_label

    return feature, label


def generate_features(video_id, video_length, filepath, model_dir, output_dir,
                      transpose=False):
  """Generate features from input video."""
  video, video_resized = load_video(filepath, video_length, transpose)

  feature_compression, label_compression = generate_subnet_feature(
      video, '%s/compressionnet_baseline' % model_dir,
      320, 180, 5,  # input height, weight, fps,
      4, 4, 100, 1,  # feature map height, width, channels, and label_size
      'input_orig:0',
      ['feature_layer_orig:0', 'compress_level_orig:0'])

  feature_content, label_content = generate_content_feature(
      video_resized, '%s/contentnet_baseline' % model_dir,
      'map/TensorArrayV2Stack/TensorListStack:0',
      ['final_conv2d/Conv2D:0', 'class_confidence:0'])

  feature_distortion, label_distortion = generate_subnet_feature(
      video, '%s/distortionnet_baseline' % model_dir,
      640, 360, 1,  # input height, weight, fps,
      8, 8, 100, 26,  # feature map height, width, channels, and label_size
      'input_images:0',
      ['feature_map:0', 'dist_type_prediction/dist_type_predictions:0'])

  # Save features
  fd, temp = tempfile.mkstemp()

  feature_compression.astype('float32').tofile(temp)
  out_feature = '%s/%s_feature_compression.binary' % (output_dir, video_id)
  gfile.Copy(temp, out_feature, overwrite=True)

  feature_content.astype('float32').tofile(temp)
  out_feature = '%s/%s_feature_content.binary' % (output_dir, video_id)
  gfile.Copy(temp, out_feature, overwrite=True)

  feature_distortion.astype('float32').tofile(temp)
  out_feature = '%s/%s_feature_distortion.binary' % (output_dir, video_id)
  gfile.Copy(temp, out_feature, overwrite=True)

  # Feature labels
  np.savetxt(temp, label_compression.reshape(label_compression.shape[0], -1),
             fmt='%0.3f', delimiter=',')
  out_feature = '%s/%s_label_compression.csv' % (output_dir, video_id)
  gfile.Copy(temp, out_feature, overwrite=True)

  np.savetxt(temp, label_content.reshape(label_content.shape[0], -1),
             fmt='%0.3f', delimiter=',')
  out_feature = '%s/%s_label_content.csv' % (output_dir, video_id)
  gfile.Copy(temp, out_feature, overwrite=True)

  np.savetxt(temp, label_distortion.reshape(label_distortion.shape[0], -1),
             fmt='%0.3f', delimiter=',')
  out_feature = '%s/%s_label_distortion.csv' % (output_dir, video_id)
  gfile.Copy(temp, out_feature, overwrite=True)

  os.close(fd)
  os.remove(temp)


def load_features(video_id, dim_time, feature_dir):
  """Load pre-generated features."""
  input_compression_feature = '%s/%s_feature_compression.binary' % (
      feature_dir, video_id)
  with gfile.Open(input_compression_feature, 'rb') as input_file:
    s = input_file.read()
    with Session() as sess:
      feature_1d = tf.io.decode_raw(s, out_type=tf.float32)

      feature = tf.reshape(feature_1d,
                           [1, dim_time, DIM_HEIGHT_FEATURE,
                            DIM_WIDTH_FEATURE, DIM_CHANNEL_FEATURE])
      feature_compression = sess.run(feature)

  input_content_feature = '%s/%s_feature_content.binary' % (
      feature_dir, video_id)
  with gfile.Open(input_content_feature, 'rb') as input_file:
    s = input_file.read()
    with Session() as sess:
      feature_1d = tf.io.decode_raw(s, out_type=tf.float32)

      feature = tf.reshape(feature_1d,
                           [1, dim_time, DIM_HEIGHT_FEATURE,
                            DIM_WIDTH_FEATURE, DIM_CHANNEL_FEATURE])
      feature_content = sess.run(feature)

  input_distortion_feature = '%s/%s_feature_distortion.binary' % (
      feature_dir, video_id)
  with gfile.Open(input_distortion_feature, 'rb') as input_file:
    s = input_file.read()
    with Session() as sess:
      feature_1d = tf.io.decode_raw(s, out_type=tf.float32)

      feature = tf.reshape(feature_1d,
                           [1, dim_time, DIM_HEIGHT_FEATURE,
                            DIM_WIDTH_FEATURE, DIM_CHANNEL_FEATURE])
      feature_distortion = sess.run(feature)

  return feature_compression, feature_content, feature_distortion


def prediction(video_id, video_length, model_dir, feature_dir, output_dir):
  """Predict quality (MOS)."""
  trainset = 'ytugc20s'
  all_trainset_subindex = ['0', '1', '2', '3', '4']

  all_feature = ['compression',
                 'content',
                 'distortion',
                 'compression_content',
                 'compression_distortion',
                 'content_distortion',
                 'compression_content_distortion',
                 ]
  aggregation_model = 'avgpool'
  all_outputs = []
  for feature in all_feature:
    aggregated_mos = 0
    for trainset_subindex in all_trainset_subindex:
      model_name = '%s_%s_%s_%s' % (
          trainset, trainset_subindex, aggregation_model, feature)

      with Session(graph=Graph(), config=ConfigProto(
          allow_soft_placement=True, log_device_placement=False)) as sess:
        saved_model.loader.load(
            sess, [saved_model.tag_constants.SERVING],
            '%s/aggregationnet_baseline/%s' % (
                model_dir, model_name))

        [feature_compression, feature_content, feature_distortion
         ] = load_features(video_id, video_length, feature_dir)
        feature_compression = feature_compression[
            :, 0:video_length, :, :, :]
        feature_content = feature_content[
            :, 0:video_length, :, :, :]
        feature_distortion = feature_distortion[
            :, 0:video_length, :, :, :]

        pred_mos = sess.run(
            'Model/mos:0',
            feed_dict={'feature_compression:0': feature_compression,
                       'feature_content:0': feature_content,
                       'feature_distortion:0': feature_distortion,
                       })
        pred_mos = pred_mos[0][0]
        aggregated_mos += pred_mos

    aggregated_mos /= len(all_trainset_subindex)
    all_outputs.append([video_id, feature, aggregated_mos])

    fd, temp = tempfile.mkstemp()
    with gfile.Open(temp, 'w') as f:
      writer = csv.writer(f)
      writer.writerows(all_outputs)
    out_file = '%s/%s_uvq.csv' % (output_dir, video_id)
    gfile.Copy(temp, out_file, overwrite=True)
    os.close(fd)
    os.remove(temp)