def viz()

in build_graph/build.py [0:0]

• 49 lines of code
• 12 McCabe index (conditional complexity)

    def viz(self, i, block_size=512, img_sz=100, nsz=50, out_sz=2048):

        if not self.args.viz:
            return
        viz_items = []


        # current frame
        frame = self.load_frame_viz(self.frames[self.state['frame']])
        frame = resize_tensor(frame, block_size)
        viz_items.append(frame)

        #-----------------------------------------------------------------#

        # visited state grid
        frame_imgs = {}
        nbhs = set(self.G.neighbors(self.state['viz_node'])) if self.state['viz_node'] in self.G else set()
        node_ref = []
        for n in self.G.nodes():

            last_visit = self.G.nodes[n]['members'][-1]
            frame_i = (last_visit['start'] + last_visit['stop'])//2
            img = self.load_frame_viz(self.frames[frame_i])
            img = transforms.ToPILImage()(img)

            if n==self.state['viz_node']:
                color = 'grey' if self.state['inactive'] else 'green'
                img = ImageOps.expand(img, border=15, fill=color)
            elif n in nbhs and not self.state['inactive']:
                img = ImageOps.expand(img, border=15, fill='blue')
            
            img = img.resize((img_sz, img_sz))

            node_ref.append(img)
            frame_imgs[n] = img

        N = 5
        node_ref = node_ref[:N**2]
        node_ref = node_ref + [Image.new("RGB", (img_sz, img_sz), "white")]*(N**2-len(node_ref))
        node_ref = [transforms.ToTensor()(img) for img in node_ref]
        grid = make_grid(node_ref, nrow=N)
        grid = resize_tensor(grid, block_size)
        viz_items.append(grid)

        #-----------------------------------------------------------------#

        # graph
        plt.clf()
        fig = plt.gcf()
        fig.set_size_inches(block_size/fig.get_dpi(), block_size/fig.get_dpi())
        plt.xlim(-1,1)
        plt.ylim(-1,1)

        self.G.pos = nx.spring_layout(self.G, scale=0.8, pos=self.G.pos, iterations=5)
        edges = self.G.edges()
        colors = ['black' if np.min(self.G[u][v]['dT']) < self.edge_delay else 'grey' for u, v in edges]
        nx.draw(self.G, self.G.pos, edgelist=edges, edge_color=colors, node_size=10)

        for n in self.G.nodes():
            x, y = self.G.pos[n]
            x, y = int((x+1)/2*block_size), int((y+1)/2*block_size)
            img = frame_imgs[n].resize((nsz, nsz))
            fig.figimage(img, x-nsz//2, y-nsz//2, zorder=10)
        graph = fig2tensor(fig)
        graph = resize_tensor(graph, block_size)
        viz_items.append(graph)

        #-----------------------------------------------------------------#

        viz = make_grid(viz_items, nrow=3)
        util.show_wait(viz, sz=out_sz, T=0)