my_utils/plot.py

# import os, sys
# sys.path.append(os.path.join(os.path.dirname(__file__), '..', '..'))
# import matplotlib
# matplotlib.use('Agg')
# import numpy as np
# import networkx as nx
# from io import BytesIO
# from PIL import Image, ImageDraw
# from matplotlib import pyplot as plt
# from sklearn.decomposition import PCA
# from my_utils.refs import graph_color_ref

# def add_text(text, imgarr):
#     '''
#     Function to add text to image

#     Args:
#     - text (str): text to add
#     - imgarr (np.array): image array

#     Returns:
#     - img (np.array): image array with text
#     '''
#     img = Image.fromarray(imgarr)
#     I = ImageDraw.Draw(img)
#     I.text((10, 10), text, fill='black')
#     return np.asarray(img)

# def get_color(ref, n_nodes):
#     '''
#     Function to color the nodes

#     Args:
#     - ref (list): list of color reference
#     - n_nodes (int): number of nodes

#     Returns:
#     - colors (list): list of colors
#     '''
#     N = len(ref)
#     colors = []
#     for i in range(n_nodes):
#         colors.append(np.array([[int(i) for i in ref[i%N][4:-1].split(',')]]) / 255.)
#     return colors


# def make_grid(images, cols=5):
#     """
#     Arrange list of images into a N x cols grid.
    
#     Args:
#     - images (list): List of Numpy arrays representing the images.
#     - cols (int): Number of columns for the grid.
    
#     Returns:
#     - grid (numpy array): Numpy array representing the image grid.
#     """
#     # Determine the dimensions of each image
#     img_h, img_w, _ = images[0].shape
#     rows = len(images) // cols
    
#     # Initialize a blank canvas
#     grid = np.zeros((rows * img_h, cols * img_w, 3), dtype=images[0].dtype)
    
#     # Place each image onto the grid
#     for idx, img in enumerate(images):
#         y = (idx // cols) * img_h
#         x = (idx % cols) * img_w
#         grid[y: y + img_h, x: x + img_w] = img
    
#     return grid

# def viz_graph(info_dict, res=256):
#     '''
#     Function to plot the directed graph

#     Args:
#     - info_dict (dict): output json containing the graph information
#     - res (int): resolution of the image

#     Returns:
#     - img_arr (np.array): image array
#     '''
#     # build tree
#     tree = info_dict['diffuse_tree']
#     edges = []
#     for node in tree:
#         edges += [(node['id'], child) for child in node['children']]
#     G = nx.DiGraph()
#     G.add_edges_from(edges)

#     # plot tree
#     plt.figure(figsize=(res/100, res/100))

#     colors = get_color(graph_color_ref, len(tree))
#     pos = nx.nx_agraph.graphviz_layout(G, prog="twopi", args="")
#     node_order = sorted(G.nodes())
#     nx.draw(G, pos, node_color=colors, nodelist=node_order, edge_color='k', with_labels=False)
    
#     buf = BytesIO()
#     plt.savefig(buf, format="png", dpi=100)
#     buf.seek(0)
#     img = Image.open(buf)
#     img_arr = np.asarray(img)
#     buf.close()
#     plt.clf()
#     plt.close()
#     return img_arr[:, :, :3]

# def viz_patch_feat_pca(feat):
#     pca = PCA(n_components=3)
#     pca.fit(feat)
#     feat_pca = pca.transform(feat)
    
#     t = np.array(feat_pca)
#     t_min = t.min(axis=0, keepdims=True)
#     t_max = t.max(axis=0, keepdims=True)
#     normalized_t = (t - t_min) / (t_max - t_min)

#     array = (normalized_t * 255).astype(np.uint8)
#     img_array = array.reshape(16, 16, 3)
#     return img_array