Update deepfloorplan_inference.py

deepfloorplan_inference.py

@@ -1,56 +1,58 @@
-import os
-import numpy as np
-import tensorflow.compat.v1 as tf
-tf.disable_v2_behavior()
-from PIL import Image
-import imageio
-from net import Network
-from utils.rgb_ind_convertor import ind2rgb, floorplan_fuse_map
-
-class DeepFloorPlanModel:
-    def __init__(self, model_dir='pretrained', input_size=(512, 512)):
-        self.input_size = input_size
-        self.model_dir = model_dir
-        self._build_graph()
-        self._load_weights()
-
-    def _build_graph(self):
-        tf.compat.v1.reset_default_graph()
-        self.sess = tf.compat.v1.Session()
-        self.x = tf.compat.v1.placeholder(shape=[1, self.input_size[0], self.input_size[1], 3], dtype=tf.float32, name='inputs')
-        self.network = Network()
-        logits1, logits2 = self.network.forward(self.x, init_with_pretrain_vgg=False)
-        self.rooms = self.network.convert_one_hot_to_image(logits1, act='softmax', dtype='int')
-        self.close_walls = self.network.convert_one_hot_to_image(logits2, act='softmax', dtype='int')
-        self.sess.run(tf.compat.v1.global_variables_initializer())
-        self.sess.run(tf.compat.v1.local_variables_initializer())
-        self.saver = tf.compat.v1.train.Saver()
-
-    def _load_weights(self):
-        ckpt = tf.train.latest_checkpoint(self.model_dir)
-        if ckpt is None:
-            raise FileNotFoundError(f"No checkpoint found in {self.model_dir}")
-        self.saver.restore(self.sess, ckpt)
-
-    def predict(self, image):
-        # Accepts a numpy array or PIL image, returns a numpy array (segmentation mask)
-        if isinstance(image, Image.Image):
-            image = np.array(image)
-        if image.shape[-1] == 4:
-            image = image[..., :3]
-        im_resized = np.array(Image.fromarray(image).resize(self.input_size, Image.BICUBIC)) / 255.0
-        im_resized = im_resized.astype(np.float32)
-        im_resized = np.reshape(im_resized, (1, self.input_size[0], self.input_size[1], 3))
-        out1, out2 = self.sess.run([self.rooms, self.close_walls], feed_dict={self.x: im_resized})
-        out1 = np.squeeze(out1)
-        out2 = np.squeeze(out2)
-        # Merge logic: set out1 pixels to 9/10 where out2==1/2
-        out1[out2==2] = 10
-        out1[out2==1] = 9
-        # Convert to RGB for visualization
-        out_rgb = ind2rgb(out1, color_map=floorplan_fuse_map)
-        out_rgb = out_rgb.astype(np.uint8)
-        return out_rgb
-
-    def close(self):
+import os
+import numpy as np
+import tensorflow.compat.v1 as tf
+tf.disable_v2_behavior()
+from PIL import Image
+import imageio
+from net import Network
+from utils.rgb_ind_convertor import ind2rgb, floorplan_fuse_map
+
+class DeepFloorPlanModel:
+    def __init__(self, model_dir='pretrained', input_size=(512, 512)):
+        self.input_size = input_size
+        self.model_dir = model_dir
+        self._build_graph()
+        self._load_weights()
+
+    def _build_graph(self):
+        tf.compat.v1.reset_default_graph()
+        self.sess = tf.compat.v1.Session()
+        self.x = tf.compat.v1.placeholder(shape=[1, self.input_size[0], self.input_size[1], 3], dtype=tf.float32, name='inputs')
+        self.network = Network()
+        logits1, logits2 = self.network.forward(self.x, init_with_pretrain_vgg=False)
+        self.rooms = self.network.convert_one_hot_to_image(logits1, act='softmax', dtype='int')
+        self.close_walls = self.network.convert_one_hot_to_image(logits2, act='softmax', dtype='int')
+        self.sess.run(tf.compat.v1.global_variables_initializer())
+        self.sess.run(tf.compat.v1.local_variables_initializer())
+        self.saver = tf.compat.v1.train.Saver()
+
+    def _load_weights(self):
+        ckpt = tf.train.latest_checkpoint(self.model_dir)
+        if ckpt is None:
+            print(f"[ERROR] No checkpoint found in {self.model_dir}")
+            raise FileNotFoundError(f"No checkpoint found in {self.model_dir}")
+        print(f"[INFO] Restoring model weights from {ckpt}")
+        self.saver.restore(self.sess, ckpt)
+
+    def predict(self, image):
+        # Accepts a numpy array or PIL image, returns a numpy array (segmentation mask)
+        if isinstance(image, Image.Image):
+            image = np.array(image)
+        if image.shape[-1] == 4:
+            image = image[..., :3]
+        im_resized = np.array(Image.fromarray(image).resize(self.input_size, Image.BICUBIC)) / 255.0
+        im_resized = im_resized.astype(np.float32)
+        im_resized = np.reshape(im_resized, (1, self.input_size[0], self.input_size[1], 3))
+        out1, out2 = self.sess.run([self.rooms, self.close_walls], feed_dict={self.x: im_resized})
+        out1 = np.squeeze(out1)
+        out2 = np.squeeze(out2)
+        # Merge logic: set out1 pixels to 9/10 where out2==1/2
+        out1[out2==2] = 10
+        out1[out2==1] = 9
+        # Convert to RGB for visualization
+        out_rgb = ind2rgb(out1, color_map=floorplan_fuse_map)
+        out_rgb = out_rgb.astype(np.uint8)
+        return out_rgb
+
+    def close(self):
         self.sess.close()