# ---
# jupyter:
#   jupytext:
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.14.4
#   kernelspec:
#     display_name: Python [conda env:dgs]
#     language: python
#     name: conda-env-dgs-py
# ---

import networkx as nx

with open("../../../flicker30k_denotation_graph/flicker30k/graph/node.idx", "r") as f:
    node_idx = f.read()

node_idx = node_idx.split("\n")

with open("../../../flicker30k_denotation_graph/flicker30k/graph/node-cap.map", "r") as f:
    node_caption = f.read().split("\n")

with open("../../../flicker30k_denotation_graph/flicker30k/graph/node-img.map", "r") as f:
    node_image = f.read().split("\n")

len(node_idx)

node_idx[0]

len(node_image)

len(node_caption)

from tqdm import tqdm

graph = nx.DiGraph()
for idx, images, captions in tqdm(zip(node_idx, node_image, node_caption), total=len(node_idx)):
    if idx == "":
        continue
    i, text = idx.split("\t")
    images = images.split("\t")
    assert images[0] == i
    images = images[1:]
    # remove the caption information from the images and keep only the images
    images = list(set(map(lambda x: x.split("#")[0], images)))
    captions = captions.split("\t")
    assert captions[0] == i
    captions=captions[1:]
    graph.add_node(int(i), text=text, images=sorted(images), captions=sorted(captions))

with open("../../../flicker30k_denotation_graph/flicker30k/graph/node-tree.txt", "r") as f:
    node_tree = f.read()

node_tree = node_tree.split("\n")

for n in tqdm(node_tree):
    if n == "":
        continue
    child, edge, parent, *rewrite_rules = n.split("\t")
    child = int(child)
    parent = int(parent)
    assert child in graph
    assert parent in graph
    graph.add_edge(parent, child, edge_type=edge, rewrite_rules=rewrite_rules)

graph.size()

graph.in_degree(0)

graph.get_edge_data(0,54)

leaf_nodes = [node for node in tqdm(graph.nodes()) if graph.in_degree(node)!=0 and graph.out_degree(node)==0]

len(leaf_nodes)

len(graph.nodes)

leaf_nodes[0]

leaf_nodes[-1]

list(graph.in_edges(2642008))

subgraph = graph.edge_subgraph([(2501206, 2642008), (352, 2642008), (6902, 2642008), (1357, 2642008), (3566, 2642008)])

subgraph

list(subgraph.nodes)

list(subgraph.edges)

subgraph.subgraph([352, 2501206])

type(subgraph.nodes[352])

subgraph.edges[(2501206, 2642008)]

a = graph.successors(0)

a = list(a)

len(a)

import matplotlib.pyplot as plt
# %matplotlib inline

StarcoderdataJupyterScriptsDedupFiltered