tools/py_stats/avm_stats/frame_visualizations.py - avm - Git at Google

 #!/usr/bin/env python3
 ## Copyright (c) 2023, Alliance for Open Media. All rights reserved
 ##
 ## This source code is subject to the terms of the BSD 3-Clause Clear License and the
 ## Alliance for Open Media Patent License 1.0. If the BSD 3-Clause Clear License was
 ## not distributed with this source code in the LICENSE file, you can obtain it
 ## at aomedia.org/license/software-license/bsd-3-c-c/.  If the Alliance for Open Media Patent
 ## License 1.0 was not distributed with this source code in the PATENTS file, you
 ## can obtain it at aomedia.org/license/patent-license/.
 ##
 """Matplotlib visualization library for AVM Frame protos.

 Defines a number of composable frame visualization layers that can be used to
 show some data of interest within a specific frame in a Colab notebook.

 The typical use case will start with some representation of the underlying pixel
 data, and then render annotation additional layers on top of it, which will get
 combined together into the final plot.

 An example use case:
 ```
 viz = (ReconstructionYuvLayer()
         .add(TransformBlockLayer())
         .add(PartitionTreeLayer())
         .add(SuperblockLayer())
 fig, ax = plt.subplots(1, 1)
 viz.show(frame, ax)
 ```

 This will render the reconstructed YUV as the base layer, and then draw on top
 the transform units, the partition tree, and the superblock boundaries.
 """

 from __future__ import annotations

 import abc

 from avm_stats import proto_helpers
 import matplotlib.pyplot as plt
 from mpl_toolkits.axes_grid1 import inset_locator
 import numpy as np


 class VisualizationLayer(metaclass=abc.ABCMeta):
   """Abstract base class for all visualization layers.

   Attributes:
     layers: list of additional layers that will be rendered on top of this one.
       The additional layers will be rendered bottom to top.
   """

   def __init__(self):
     self.layers = []

   def add(self, layer: VisualizationLayer) -> VisualizationLayer:
     """Adds another visualization layer to the list of layers.

     Args:
       layer: Additional layer to render.

     Returns:
       self, so that additional add() calls can be chained together.
     """
     self.layers.append(layer)
     return self

   @abc.abstractmethod
   def render(self, frame: proto_helpers.Frame, ax: plt.Axes):
     """Render this visualization to the Matplotlib axes.

     Args:
       frame: Frame used to get render data from, e.g. pixels or coding units.
       ax: Matplotlib axes to render to.
     """

   def show(self, frame: proto_helpers.Frame, ax: plt.Axes):
     self.render(frame, ax)
     for layer in self.layers:
       layer.render(frame, ax)


 class YuvPlaneLayer(VisualizationLayer):
   """Visualize a single YUV plane.

   Attributes:
     name: Name of the layer; will be displayed in the plot title.
     pixels_attribute: The attribute of the frames PlaneBuffer to plot pixels
       from.
     plane: Which YUV plane to plot.
   """

   def __init__(
       self,
       name: str,
       pixels_attribute: str,
       *,
       plane: proto_helpers.Plane = proto_helpers.Plane.Y,
   ):
     self.name = name
     self.pixels_attribute = pixels_attribute
     self.plane = plane
     super().__init__()

   def render(self, frame: proto_helpers.Frame, ax: plt.Axes):
     ax.title.set_text(
         f"{self.name} ({self.plane.name})\nframe={frame.frame_id}"
     )
     pixels = getattr(frame.pixels[self.plane], self.pixels_attribute)
     width = frame.width
     height = frame.height
     ax.imshow(
         pixels, cmap="gray", vmin=0, vmax=255, extent=[0, width, height, 0]
     )


 class OriginalYuvLayer(YuvPlaneLayer):

   def __init__(self, **kwargs):
     super().__init__("Original", "original", **kwargs)


 class PredictionYuvLayer(YuvPlaneLayer):

   def __init__(self, **kwargs):
     super().__init__("Prediction", "prediction", **kwargs)


 class PrefilteredYuvLayer(YuvPlaneLayer):

   def __init__(self, **kwargs):
     super().__init__("Pre-filtered Reconstruction", "pre_filtered", **kwargs)


 class ReconstructionYuvLayer(YuvPlaneLayer):

   def __init__(self, **kwargs):
     super().__init__("Reconstruction", "reconstruction", **kwargs)


 class YuvDeltaLayer(VisualizationLayer):

   def __init__(
       self,
       plane: proto_helpers.Plane = proto_helpers.Plane.Y,
       show_relative=False,
   ):
     self.plane = plane
     self.show_relative = show_relative
     super().__init__()

   @abc.abstractmethod
   def get_pixels(self, frame: proto_helpers.Frame) -> np.ndarray:
     pass

   @abc.abstractmethod
   def get_name(self) -> str:
     pass

   def render(self, frame: proto_helpers.Frame, ax: plt.Axes):
     pixels = self.get_pixels(frame)
     pixel_min = np.min(pixels)
     pixel_max = np.max(pixels)
     vmin, vmax = -255, 255
     annotation = "Relative" if self.show_relative else "Absolute"
     if self.show_relative:
       vmin, vmax = pixel_min, pixel_max
     ax.title.set_text(
         f"{self.get_name()} ({self.plane.name}) - {annotation}"
         f"\nframe={frame.frame_id}, min={pixel_min}, max={pixel_max}"
     )
     width = frame.width
     height = frame.height
     ax.imshow(
         pixels, cmap="gray", vmin=vmin, vmax=vmax, extent=[0, width, height, 0]
     )


 class ResidualYuvLayer(YuvDeltaLayer):

   def __init__(
       self,
       plane: proto_helpers.Plane = proto_helpers.Plane.Y,
       show_relative=False,
   ):
     super().__init__(plane, show_relative)

   def get_pixels(self, frame: proto_helpers.Frame) -> np.ndarray:
     return frame.pixels[int(self.plane)].residual

   def get_name(self) -> str:
     return "Residual"


 class FilterDeltaYuvLayer(YuvDeltaLayer):

   def __init__(
       self,
       plane: proto_helpers.Plane = proto_helpers.Plane.Y,
       show_relative=False,
   ):
     super().__init__(plane, show_relative)

   def get_pixels(self, frame: proto_helpers.Frame) -> np.ndarray:
     return frame.pixels[int(self.plane)].filter_delta

   def get_name(self) -> str:
     return "Filter Delta"


 class DistortionYuvLayer(YuvDeltaLayer):

   def __init__(
       self,
       plane: proto_helpers.Plane = proto_helpers.Plane.Y,
       show_relative=False,
   ):
     super().__init__(plane, show_relative)

   def get_pixels(self, frame: proto_helpers.Frame) -> np.ndarray:
     return frame.pixels[int(self.plane)].distortion

   def get_name(self) -> str:
     return "Distortion"


 class SuperblockLayer(VisualizationLayer):
   """Draw the boundaries of the superblocks."""

   def render(self, frame: proto_helpers.Frame, ax: plt.Axes):
     for superblock in frame.superblocks:
       rectangle = superblock.rect
       ax.add_patch(
           plt.Rectangle(
               (rectangle.left_x, rectangle.top_y),
               rectangle.width,
               rectangle.height,
               linewidth=2,
               edgecolor="blue",
               facecolor="none",
           )
       )


 class PartitionTreeLayer(VisualizationLayer):
   """Draw the boundaries of the partition tree leaf nodes."""

   def __init__(self, use_chroma=False):
     self.use_chroma = use_chroma
     super().__init__()

   def render(self, frame: proto_helpers.Frame, ax: plt.Axes):
     for superblock in frame.superblocks:
       for rectangle in superblock.get_partition_rects(self.use_chroma):
         ax.add_patch(
             plt.Rectangle(
                 (rectangle.left_x, rectangle.top_y),
                 rectangle.width,
                 rectangle.height,
                 linewidth=1,
                 edgecolor="red",
                 facecolor="none",
             )
         )


 class TransformBlockLayer(VisualizationLayer):
   """Draw the boundaries of all transform units."""

   def __init__(self, use_chroma=False):
     self.use_chroma = use_chroma
     super().__init__()

   def render(self, frame: proto_helpers.Frame, ax: plt.Axes):
     for superblock in frame.superblocks:
       for rectangle in superblock.get_transform_rects(
           use_chroma=self.use_chroma
       ):
         ax.add_patch(
             plt.Rectangle(
                 (rectangle.left_x, rectangle.top_y),
                 rectangle.width,
                 rectangle.height,
                 linestyle=":",
                 linewidth=1,
                 edgecolor="green",
                 facecolor="none",
             )
         )


 class MotionVectorLayer(VisualizationLayer):
   MOTION_VECTOR_PRECISION = 8

   def __init__(self, use_chroma=False):
     self.use_chroma = use_chroma
     super().__init__()

   def render(self, frame: proto_helpers.Frame, ax: plt.Axes):
     # Don't plot motion vectors for intra frames
     if frame.proto.frame_params.frame_type == 0:
       return
     for sb in frame.superblocks:
       for cu in sb.get_coding_units(use_chroma=self.use_chroma):
         for mv in cu.proto.prediction_mode.motion_vectors:
           if mv.ref_frame == -1:
             continue
           cx = cu.proto.position.x + cu.proto.size.width / 2
           cy = cu.proto.position.y + cu.proto.size.height / 2
           dx = mv.dx / self.MOTION_VECTOR_PRECISION
           dy = mv.dy / self.MOTION_VECTOR_PRECISION
           ax.add_patch(
               plt.Arrow(
                   cx,
                   cy,
                   dx,
                   dy,
                   linewidth=1,
                   edgecolor="green",
               )
           )


 class SuperblockAnnotationLayer(VisualizationLayer):

   def render(self, frame: proto_helpers.Frame, ax: plt.Axes):
     for sb in frame.superblocks:
       rect = sb.clipped_rect
       cx = rect.x + rect.width / 2
       cy = rect.y + rect.height / 2
       bits = sb.get_total_bits()
       qindex = sb.qindex
       ax.text(
           cx,
           cy,
           f"Bits: {bits:.1f}\nqindex={qindex}",
           verticalalignment="center",
           horizontalalignment="center",
           size="large",
       )


 class BitsHeatmapLayer(VisualizationLayer):

   def __init__(
       self,
       use_chroma: bool = False,
       filt: proto_helpers.SymbolFilter | None = None,
   ):
     self.use_chroma = use_chroma
     self.filt = filt
     super().__init__()

   def render(self, frame: proto_helpers.Frame, ax: plt.Axes):
     heatmap = np.zeros((frame.height, frame.width))
     # TODO(comc): allow min / max across multiple frames.
     min_bpp, max_bpp = 1e9, 0
     for sb in frame.superblocks:
       bits_by_cu = list(
           sb.get_bits_per_coding_unit(
               use_chroma=self.use_chroma, filt=self.filt
           )
       )
       for i, rect in enumerate(sb.get_partition_rects(self.use_chroma)):
         bits = bits_by_cu[i]
         bpp = bits / (rect.width * rect.height)
         min_bpp = min(bpp, min_bpp)
         max_bpp = max(bpp, max_bpp)
         heatmap[
             rect.top_y : rect.top_y + rect.height,
             rect.left_x : rect.left_x + rect.width,
         ] = bpp
     width = frame.width
     height = frame.height
     pos = ax.imshow(
         heatmap,
         cmap="YlOrRd",
         interpolation="none",
         vmin=min_bpp,
         vmax=max_bpp,
         extent=[0, width, height, 0],
     )
     axins = inset_locator.inset_axes(
         ax, width="100%", height="5%", loc="lower center", borderpad=-5
     )
     ax.get_figure().colorbar(pos, cax=axins, orientation="horizontal")


 class PredictionModeAnnotationLayer(VisualizationLayer):

   def __init__(self, use_chroma: bool = False):
     self.use_chroma = use_chroma
     super().__init__()

   def render(self, frame: proto_helpers.Frame, ax: plt.Axes):
     for sb in frame.superblocks:
       for cu in sb.get_coding_units(use_chroma=self.use_chroma):
         rect = cu.rect
         intra_mode = cu.get_prediction_mode()
         cx = rect.left_x + rect.width / 2
         cy = rect.top_y + rect.height / 2
         mode = intra_mode.removesuffix("_PRED")
         ax.text(cx, cy, mode, horizontalalignment="center", size="small")
	#!/usr/bin/env python3
	## Copyright (c) 2023, Alliance for Open Media. All rights reserved
	##
	## This source code is subject to the terms of the BSD 3-Clause Clear License and the
	## Alliance for Open Media Patent License 1.0. If the BSD 3-Clause Clear License was
	## not distributed with this source code in the LICENSE file, you can obtain it
	## at aomedia.org/license/software-license/bsd-3-c-c/. If the Alliance for Open Media Patent
	## License 1.0 was not distributed with this source code in the PATENTS file, you
	## can obtain it at aomedia.org/license/patent-license/.
	##
	"""Matplotlib visualization library for AVM Frame protos.

	Defines a number of composable frame visualization layers that can be used to
	show some data of interest within a specific frame in a Colab notebook.

	The typical use case will start with some representation of the underlying pixel
	data, and then render annotation additional layers on top of it, which will get
	combined together into the final plot.

	An example use case:
	```
	viz = (ReconstructionYuvLayer()
	.add(TransformBlockLayer())
	.add(PartitionTreeLayer())
	.add(SuperblockLayer())
	fig, ax = plt.subplots(1, 1)
	viz.show(frame, ax)
	```

	This will render the reconstructed YUV as the base layer, and then draw on top
	the transform units, the partition tree, and the superblock boundaries.
	"""

	from __future__ import annotations

	import abc

	from avm_stats import proto_helpers
	import matplotlib.pyplot as plt
	from mpl_toolkits.axes_grid1 import inset_locator
	import numpy as np


	class VisualizationLayer(metaclass=abc.ABCMeta):
	"""Abstract base class for all visualization layers.

	Attributes:
	layers: list of additional layers that will be rendered on top of this one.
	The additional layers will be rendered bottom to top.
	"""

	def __init__(self):
	self.layers = []

	def add(self, layer: VisualizationLayer) -> VisualizationLayer:
	"""Adds another visualization layer to the list of layers.

	Args:
	layer: Additional layer to render.

	Returns:
	self, so that additional add() calls can be chained together.
	"""
	self.layers.append(layer)
	return self

	@abc.abstractmethod
	def render(self, frame: proto_helpers.Frame, ax: plt.Axes):
	"""Render this visualization to the Matplotlib axes.

	Args:
	frame: Frame used to get render data from, e.g. pixels or coding units.
	ax: Matplotlib axes to render to.
	"""

	def show(self, frame: proto_helpers.Frame, ax: plt.Axes):
	self.render(frame, ax)
	for layer in self.layers:
	layer.render(frame, ax)


	class YuvPlaneLayer(VisualizationLayer):
	"""Visualize a single YUV plane.

	Attributes:
	name: Name of the layer; will be displayed in the plot title.
	pixels_attribute: The attribute of the frames PlaneBuffer to plot pixels
	from.
	plane: Which YUV plane to plot.
	"""

	def __init__(
	self,
	name: str,
	pixels_attribute: str,
	*,
	plane: proto_helpers.Plane = proto_helpers.Plane.Y,
	):
	self.name = name
	self.pixels_attribute = pixels_attribute
	self.plane = plane
	super().__init__()

	def render(self, frame: proto_helpers.Frame, ax: plt.Axes):
	ax.title.set_text(
	f"{self.name} ({self.plane.name})\nframe={frame.frame_id}"
	)
	pixels = getattr(frame.pixels[self.plane], self.pixels_attribute)
	width = frame.width
	height = frame.height
	ax.imshow(
	pixels, cmap="gray", vmin=0, vmax=255, extent=[0, width, height, 0]
	)


	class OriginalYuvLayer(YuvPlaneLayer):

	def __init__(self, **kwargs):
	super().__init__("Original", "original", **kwargs)


	class PredictionYuvLayer(YuvPlaneLayer):

	def __init__(self, **kwargs):
	super().__init__("Prediction", "prediction", **kwargs)


	class PrefilteredYuvLayer(YuvPlaneLayer):

	def __init__(self, **kwargs):
	super().__init__("Pre-filtered Reconstruction", "pre_filtered", **kwargs)


	class ReconstructionYuvLayer(YuvPlaneLayer):

	def __init__(self, **kwargs):
	super().__init__("Reconstruction", "reconstruction", **kwargs)


	class YuvDeltaLayer(VisualizationLayer):

	def __init__(
	self,
	plane: proto_helpers.Plane = proto_helpers.Plane.Y,
	show_relative=False,
	):
	self.plane = plane
	self.show_relative = show_relative
	super().__init__()

	@abc.abstractmethod
	def get_pixels(self, frame: proto_helpers.Frame) -> np.ndarray:
	pass

	@abc.abstractmethod
	def get_name(self) -> str:
	pass

	def render(self, frame: proto_helpers.Frame, ax: plt.Axes):
	pixels = self.get_pixels(frame)
	pixel_min = np.min(pixels)
	pixel_max = np.max(pixels)
	vmin, vmax = -255, 255
	annotation = "Relative" if self.show_relative else "Absolute"
	if self.show_relative:
	vmin, vmax = pixel_min, pixel_max
	ax.title.set_text(
	f"{self.get_name()} ({self.plane.name}) - {annotation}"
	f"\nframe={frame.frame_id}, min={pixel_min}, max={pixel_max}"
	)
	width = frame.width
	height = frame.height
	ax.imshow(
	pixels, cmap="gray", vmin=vmin, vmax=vmax, extent=[0, width, height, 0]
	)


	class ResidualYuvLayer(YuvDeltaLayer):

	def __init__(
	self,
	plane: proto_helpers.Plane = proto_helpers.Plane.Y,
	show_relative=False,
	):
	super().__init__(plane, show_relative)

	def get_pixels(self, frame: proto_helpers.Frame) -> np.ndarray:
	return frame.pixels[int(self.plane)].residual

	def get_name(self) -> str:
	return "Residual"


	class FilterDeltaYuvLayer(YuvDeltaLayer):

	def __init__(
	self,
	plane: proto_helpers.Plane = proto_helpers.Plane.Y,
	show_relative=False,
	):
	super().__init__(plane, show_relative)

	def get_pixels(self, frame: proto_helpers.Frame) -> np.ndarray:
	return frame.pixels[int(self.plane)].filter_delta

	def get_name(self) -> str:
	return "Filter Delta"


	class DistortionYuvLayer(YuvDeltaLayer):

	def __init__(
	self,
	plane: proto_helpers.Plane = proto_helpers.Plane.Y,
	show_relative=False,
	):
	super().__init__(plane, show_relative)

	def get_pixels(self, frame: proto_helpers.Frame) -> np.ndarray:
	return frame.pixels[int(self.plane)].distortion

	def get_name(self) -> str:
	return "Distortion"


	class SuperblockLayer(VisualizationLayer):
	"""Draw the boundaries of the superblocks."""

	def render(self, frame: proto_helpers.Frame, ax: plt.Axes):
	for superblock in frame.superblocks:
	rectangle = superblock.rect
	ax.add_patch(
	plt.Rectangle(
	(rectangle.left_x, rectangle.top_y),
	rectangle.width,
	rectangle.height,
	linewidth=2,
	edgecolor="blue",
	facecolor="none",
	)
	)


	class PartitionTreeLayer(VisualizationLayer):
	"""Draw the boundaries of the partition tree leaf nodes."""

	def __init__(self, use_chroma=False):
	self.use_chroma = use_chroma
	super().__init__()

	def render(self, frame: proto_helpers.Frame, ax: plt.Axes):
	for superblock in frame.superblocks:
	for rectangle in superblock.get_partition_rects(self.use_chroma):
	ax.add_patch(
	plt.Rectangle(
	(rectangle.left_x, rectangle.top_y),
	rectangle.width,
	rectangle.height,
	linewidth=1,
	edgecolor="red",
	facecolor="none",
	)
	)


	class TransformBlockLayer(VisualizationLayer):
	"""Draw the boundaries of all transform units."""

	def __init__(self, use_chroma=False):
	self.use_chroma = use_chroma
	super().__init__()

	def render(self, frame: proto_helpers.Frame, ax: plt.Axes):
	for superblock in frame.superblocks:
	for rectangle in superblock.get_transform_rects(
	use_chroma=self.use_chroma
	):
	ax.add_patch(
	plt.Rectangle(
	(rectangle.left_x, rectangle.top_y),
	rectangle.width,
	rectangle.height,
	linestyle=":",
	linewidth=1,
	edgecolor="green",
	facecolor="none",
	)
	)


	class MotionVectorLayer(VisualizationLayer):
	MOTION_VECTOR_PRECISION = 8

	def __init__(self, use_chroma=False):
	self.use_chroma = use_chroma
	super().__init__()

	def render(self, frame: proto_helpers.Frame, ax: plt.Axes):
	# Don't plot motion vectors for intra frames
	if frame.proto.frame_params.frame_type == 0:
	return
	for sb in frame.superblocks:
	for cu in sb.get_coding_units(use_chroma=self.use_chroma):
	for mv in cu.proto.prediction_mode.motion_vectors:
	if mv.ref_frame == -1:
	continue
	cx = cu.proto.position.x + cu.proto.size.width / 2
	cy = cu.proto.position.y + cu.proto.size.height / 2
	dx = mv.dx / self.MOTION_VECTOR_PRECISION
	dy = mv.dy / self.MOTION_VECTOR_PRECISION
	ax.add_patch(
	plt.Arrow(
	cx,
	cy,
	dx,
	dy,
	linewidth=1,
	edgecolor="green",
	)
	)


	class SuperblockAnnotationLayer(VisualizationLayer):

	def render(self, frame: proto_helpers.Frame, ax: plt.Axes):
	for sb in frame.superblocks:
	rect = sb.clipped_rect
	cx = rect.x + rect.width / 2
	cy = rect.y + rect.height / 2
	bits = sb.get_total_bits()
	qindex = sb.qindex
	ax.text(
	cx,
	cy,
	f"Bits: {bits:.1f}\nqindex={qindex}",
	verticalalignment="center",
	horizontalalignment="center",
	size="large",
	)


	class BitsHeatmapLayer(VisualizationLayer):

	def __init__(
	self,
	use_chroma: bool = False,
	filt: proto_helpers.SymbolFilter \| None = None,
	):
	self.use_chroma = use_chroma
	self.filt = filt
	super().__init__()

	def render(self, frame: proto_helpers.Frame, ax: plt.Axes):
	heatmap = np.zeros((frame.height, frame.width))
	# TODO(comc): allow min / max across multiple frames.
	min_bpp, max_bpp = 1e9, 0
	for sb in frame.superblocks:
	bits_by_cu = list(
	sb.get_bits_per_coding_unit(
	use_chroma=self.use_chroma, filt=self.filt
	)
	)
	for i, rect in enumerate(sb.get_partition_rects(self.use_chroma)):
	bits = bits_by_cu[i]
	bpp = bits / (rect.width * rect.height)
	min_bpp = min(bpp, min_bpp)
	max_bpp = max(bpp, max_bpp)
	heatmap[
	rect.top_y : rect.top_y + rect.height,
	rect.left_x : rect.left_x + rect.width,
	] = bpp
	width = frame.width
	height = frame.height
	pos = ax.imshow(
	heatmap,
	cmap="YlOrRd",
	interpolation="none",
	vmin=min_bpp,
	vmax=max_bpp,
	extent=[0, width, height, 0],
	)
	axins = inset_locator.inset_axes(
	ax, width="100%", height="5%", loc="lower center", borderpad=-5
	)
	ax.get_figure().colorbar(pos, cax=axins, orientation="horizontal")


	class PredictionModeAnnotationLayer(VisualizationLayer):

	def __init__(self, use_chroma: bool = False):
	self.use_chroma = use_chroma
	super().__init__()

	def render(self, frame: proto_helpers.Frame, ax: plt.Axes):
	for sb in frame.superblocks:
	for cu in sb.get_coding_units(use_chroma=self.use_chroma):
	rect = cu.rect
	intra_mode = cu.get_prediction_mode()
	cx = rect.left_x + rect.width / 2
	cy = rect.top_y + rect.height / 2
	mode = intra_mode.removesuffix("_PRED")
	ax.text(cx, cy, mode, horizontalalignment="center", size="small")