From 97ff3c25ac3a2dde8c4248b01985931dad1d2776 Mon Sep 17 00:00:00 2001
From: Gregory Kielian <gkielian@google.com>
Date: Wed, 28 Aug 2024 22:45:23 -0700
Subject: [PATCH] Add numpy npy viewer and heatmap genertor

Creating a new folder for visualizations
---
 visualization_util/README.md              | 70 +++++++++++++++++++++++
 visualization_util/view_npy_as_heatmap.py | 54 +++++++++++++++++
 2 files changed, 124 insertions(+)
 create mode 100644 visualization_util/README.md
 create mode 100644 visualization_util/view_npy_as_heatmap.py

diff --git a/visualization_util/README.md b/visualization_util/README.md
new file mode 100644
index 0000000000..63716277ea
--- /dev/null
+++ b/visualization_util/README.md
@@ -0,0 +1,70 @@
+# `view_npy_heatmap.py`
+
+`view_npy_heatmap.py` is a Python script that allows you to visualize the contents of a `.npy` file as a heatmap using Seaborn. This is particularly useful for inspecting 2D arrays, such as matrices or images, stored in `.npy` files.
+
+## Features
+
+- **Load and visualize 2D numpy arrays** stored in `.npy` files.
+- **Customizable colormap** to adjust the heatmap's appearance.
+- **Option to save** the heatmap as an image file.
+- **Simple command-line interface** for easy usage.
+
+## Prerequisites
+
+Before using the script, make sure you have the following Python packages installed:
+
+- `numpy`
+- `seaborn`
+- `matplotlib`
+- `argparse`
+
+You can install these packages using pip:
+
+```bash
+pip install numpy seaborn matplotlib
+```
+
+## Usage
+
+### Basic Usage
+
+To visualize a `.npy` file as a heatmap, run the script with the path to the file:
+
+```bash
+python view_npy_heatmap.py path/to/yourfile.npy
+```
+
+### Customizing the Colormap
+
+You can specify a different colormap using the `--cmap` option. For example, to use the `plasma` colormap:
+
+```bash
+python view_npy_heatmap.py path/to/yourfile.npy --cmap plasma
+```
+
+### Saving the Heatmap
+
+If you want to save the heatmap as an image file instead of displaying it, use the `--save` option:
+
+```bash
+python view_npy_heatmap.py path/to/yourfile.npy --save output.png
+```
+
+This will save the heatmap to `output.png`.
+
+### Full Example
+
+To view a heatmap with the `inferno` colormap and save it as `heatmap.png`:
+
+```bash
+python view_npy_heatmap.py path/to/yourfile.npy --cmap inferno --save heatmap.png
+```
+
+## Supported Data Types
+
+This script currently supports 2D numpy arrays. If you attempt to load a 1D or
+higher-dimensional array, the script will display an error.
+
+## License
+
+This script is provided as-is under the MIT License. Feel free to modify and distribute it as needed.
diff --git a/visualization_util/view_npy_as_heatmap.py b/visualization_util/view_npy_as_heatmap.py
new file mode 100644
index 0000000000..64a8484030
--- /dev/null
+++ b/visualization_util/view_npy_as_heatmap.py
@@ -0,0 +1,54 @@
+import numpy as np
+import seaborn as sns
+import matplotlib.pyplot as plt
+import argparse
+import os
+
+def load_npy_file(file_path):
+    try:
+        data = np.load(file_path)
+        return data
+    except Exception as e:
+        print(f"Error loading file: {e}")
+        return None
+
+def display_heatmap(data, cmap='viridis', save_path=None):
+    plt.figure(figsize=(10, 8))
+    sns.heatmap(data, cmap=cmap, annot=False, fmt="g")
+
+    plt.title("Numpy Array Heatmap")
+    plt.xlabel("Column Index")
+    plt.ylabel("Row Index")
+
+    if save_path:
+        plt.savefig(save_path)
+        print(f"Heatmap saved to {save_path}")
+    else:
+        plt.show()
+
+def main():
+    parser = argparse.ArgumentParser(description="View a .npy file as a heatmap using Seaborn.")
+    parser.add_argument("file_path", type=str, help="Path to the .npy file.")
+    parser.add_argument("--cmap", type=str, default="viridis", help="Colormap to use for the heatmap (default: viridis).")
+    parser.add_argument("--save", type=str, default=None, help="Path to save the heatmap image (optional).")
+
+    args = parser.parse_args()
+
+    if not os.path.exists(args.file_path):
+        print(f"File not found: {args.file_path}")
+        return
+
+    data = load_npy_file(args.file_path)
+
+    if data is None:
+        return
+
+    if data.ndim != 2:
+        print(f"Unsupported data shape: {data.shape}. Only 2D arrays can be visualized as a heatmap.")
+        return
+
+    display_heatmap(data, cmap=args.cmap, save_path=args.save)
+
+if __name__ == "__main__":
+    main()
+