Spaces:
Runtime error
Runtime error
| 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): | |
| self.sess.close() |