Spaces:
Runtime error
Runtime error
File size: 2,558 Bytes
d359351 419718d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 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:
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() |