From 45450b15cadf02cf822781d846bce6fac1156b89 Mon Sep 17 00:00:00 2001 From: SAM NIJIN Date: Sat, 27 Jul 2024 13:57:45 +0530 Subject: [PATCH] updated funtion calls for tensorflow 2.x --- deepbrain/extractor.py | 10 +++++----- 1 file changed, 5 insertions(+), 5 deletions(-) diff --git a/deepbrain/extractor.py b/deepbrain/extractor.py index fa380bc..217f596 100644 --- a/deepbrain/extractor.py +++ b/deepbrain/extractor.py @@ -3,8 +3,8 @@ from skimage.transform import resize import os -PB_FILE = os.path.join(os.path.dirname(__file__), "models", "extractor", "graph_v2.pb") -CHECKPOINT_DIR = os.path.join(os.path.dirname(__file__), "models", "extractor", "v2") +PB_FILE = os.path.join(os.path.dirname(__file__), "models", "graph.pb") +CHECKPOINT_DIR = os.path.join(os.path.dirname(__file__), "models") class Extractor: @@ -15,9 +15,9 @@ def __init__(self): def load_pb(self): graph = tf.Graph() - self.sess = tf.Session(graph=graph) - with tf.gfile.FastGFile(PB_FILE, 'rb') as f: - graph_def = tf.GraphDef() + self.sess = tf.compat.v1.Session(graph=graph) + with tf.io.gfile.GFile(PB_FILE, 'rb') as f: + graph_def = tf.compat.v1.GraphDef() graph_def.ParseFromString(f.read()) with self.sess.graph.as_default(): tf.import_graph_def(graph_def)