def GetSceneGraph(image_id, images=None, imageDataDir='by-id/'):
if images is None:
images = utils.GetDataDir()
imageDataDir = utils.GetDataDir() + '/' + imageDataDir
if type(images) is str:
images = ListToDict(GetAllImageData(images))
fname = str(image_id) + '.json'
image = images[image_id]
data = json.load(open(imageDataDir + fname, 'r'))
scene_graph = ParseGraphLocal(data, image)
return scene_graph
def GetAllQAs(dataDir=None):
if dataDir is None:
dataDir = utils.GetDataDir()
dataFile = os.path.join(dataDir, 'question_answers.json')
imageData = GetAllImageData(dataDir)
imageMap = {}
for d in imageData:
imageMap[d.id] = d
images = json.load(open(dataFile))
output = []
for image in images:
output.append(utils.ParseQA(image['qas'], imageMap))
return output
def GetAllRegionDescriptions(dataDir=None):
if dataDir is None:
dataDir = utils.GetDataDir()
dataFile = os.path.join(dataDir, 'region_descriptions.json')
imageData = GetAllImageData(dataDir)
imageMap = {}
for d in imageData:
imageMap[d.id] = d
images = json.load(open(dataFile))
output = []
for image in images:
output.append(
utils.ParseRegionDescriptions(image['regions'],
imageMap[image['id']]))
return output
def SaveSceneGraphsById(dataDir=None, imageDataDir='by-id/'):
if dataDir is None:
dataDir = utils.GetDataDir()
imageDataDir = dataDir + '/' + imageDataDir
if not os.path.exists(imageDataDir): os.mkdir(imageDataDir)
all_data = json.load(open(os.path.join(dataDir, 'scene_graphs.json')))
for sg_data in all_data:
img_fname = str(sg_data['image_id']) + '.json'
with open(os.path.join(imageDataDir, img_fname), 'w') as f:
json.dump(sg_data, f)
del all_data
gc.collect() # clear memory
def GetSceneGraphs(startIndex=0,
endIndex=-1,
dataDir=None,
imageDataDir='by-id/',
minRels=0,
maxRels=100):
if dataDir is None:
dataDir = utils.GetDataDir()
imageDataDir = dataDir + '/' + imageDataDir
images = ListToDict(GetAllImageData(dataDir))
scene_graphs = []
img_fnames = os.listdir(imageDataDir)
if (endIndex < 1): endIndex = len(img_fnames)
for fname in img_fnames[startIndex:endIndex]:
image_id = int(fname.split('.')[0])
scene_graph = GetSceneGraph(image_id, images, imageDataDir)
n_rels = len(scene_graph.relationships)
if (minRels <= n_rels <= maxRels):
scene_graphs.append(scene_graph)
return scene_graphs
def GetAllImageData(dataDir=None):
if dataDir is None:
dataDir = utils.GetDataDir()
dataFile = os.path.join(dataDir, 'image_data.json')
data = json.load(open(dataFile))
return [utils.ParseImageData(image) for image in data]