Add renderer, material module, dataset generation, configuration file

Author: STASYA00

dataset/dataset.py

@@ -0,0 +1,57 @@
+import bpy, bmesh
+from datetime import datetime
+from math import radians
+import numpy as np
+import os
+import random
+import sys
+
+file_dir = os.path.dirname(__file__)
+sys.path.append(file_dir)
+
+from annotation import Annotation
+from blender_utils import extrude, gancio, get_min_max
+from dataset_config import *
+from generator import BuildingFactory
+from material import MaterialFactory
+from module import *
+from shp2obj import Collection, deselect_all
+
+
+class Dataset:
+	def __init__(self):
+		self.name = 'Building_dataset_{}_{}_{}'.format(datetime.now().year,
+		                                               datetime.now().month,
+		                                               datetime.now().day)
+		self.size = SIZE
+		self.json = Annotation()
+		self.factory = BuildingFactory()
+		self.material_factory = MaterialFactory()
+
+	def populate(self):
+		for i in range(self.size):
+			building = self.factory.produce()
+			building.make()
+			if use_materials:
+				_monomaterial = np.random.random() < MATERIAL_PROB
+				mat = self.material_factory.produce()
+				print(mat.name)
+				for v in building.volumes:
+					if not _monomaterial:
+						mat = self.material_factory.produce()
+					v.apply(mat)
+			self.json.add(building, '{}.png'.format(i), '{}.obj'.format(i))
+			building.save(filename=str(i))
+			building.demolish()
+
+
+	def write(self):
+		self.json.write(self.name + '.json')
+
+
+if __name__ == '__main__':
+	d = Dataset()
+	d.populate()
+	d.write()
+
+

dataset/dataset_config.py

@@ -0,0 +1,28 @@
+MIN_HEIGHT = 3.0
+MIN_WIDTH = 6.0
+MIN_LENGTH = 6.0
+
+MAX_HEIGHT = 30.0
+MAX_WIDTH = 30.0
+MAX_LENGTH = 30.0
+
+MAX_VOLUMES = 4
+
+# BUILDINGS = ['Patio', 'L', 'C', 'Single', 'Skyscraper', 'Closedpatio', 'Equalpatio']
+# Choose building typologies to be produced
+BUILDINGS = ['Patio', 'L', 'C', 'Single', 'Skyscraper', 'Closedpatio', 'Equalpatio']
+
+SIZE = 2  # dataset size
+
+use_materials = True  # apply materials to the facades of the buildings, bool
+
+MATERIAL_PROB = 0.7  # Probability of all the volumes of one building to be of the same material
+
+use_modules = True
+MODULES = ['window']
+
+MODEL_SAVE = 'Models'
+IMG_SAVE = 'Images'
+MASK_SAVE = 'Masks'
+
+ENGINE = 'CYCLES'

dataset/material.py

@@ -0,0 +1,115 @@
+import bpy
+import numpy as np
+import os
+import sys
+
+file_dir = os.path.dirname(__file__)
+sys.path.append(file_dir)
+file_dir = file_dir.replace('\\', '/').replace('\r', '/r').replace('\n', '/n').\
+	replace('\t', '/t')
+
+
+class Material:
+	"""
+	Class that represents a material object in Blender.
+	"""
+	def __init__(self, name, filename='material'):
+		self.name = name.lower().capitalize()  # name of the material and its texture folder
+		self.filename = filename
+		self._path = file_dir + '/Textures/{}.blend'.format(self.filename)
+		self._add = '\\Material\\'
+		self.value = self._load()  # loads material into the scene
+		self._update_nodes()  # loads the textures to the material
+
+	def _load(self):
+		try:
+			return bpy.data.materials[self.name].copy()
+		except KeyError:
+			try:
+				# print(file_dir)
+				bpy.ops.wm.append(
+					filepath=self._path + self._add + self.filename,
+					filename='{}'.format(self.filename),
+					directory=self._path + self._add)
+				bpy.data.materials[-1].name = self.name
+				return bpy.data.materials[self.name]
+			except Exception as e:
+				print(repr(e))
+				print('Could not import {} from {}'.format(self.name, self._path))
+				raise KeyboardInterrupt()
+
+	def _load_maps(self, map_type):
+		"""
+		Function that uploads texture map into the Material tree nodes. Textures
+		 are taken from folder Texture/Material where Material corresponds to the
+		 name of the material.
+		:param map_type: type of the map to upload, str, one of 'Diffuse', 'Normal',
+		'Roughness', 'Displacement'
+		:return: texture map, bpy image object
+		"""
+		assert map_type in ['Diffuse', 'Normal', 'Roughness', 'Displacement'], \
+			"Unknown map type, expected one of: 'Diffuse', 'Normal'," \
+			"'Roughness', 'Displacement'"
+		try:
+			bpy.ops.image.open(filepath=file_dir + '/Textures/{}/{}.png'.
+			                   format(self.name, map_type.capitalize()))
+			return bpy.data.images[-2]
+		except Exception as e:
+			print('Failed to load {} texture of {}'.format(map_type, self.name))
+			print(repr(e))
+
+	def _update_nodes(self):
+		"""
+		Function that updates the nodes of the material tree with the material
+		textures.
+		:return:
+		"""
+		self.value.node_tree.nodes['Diffuse_texture'].image = self._load_maps('Diffuse')
+		self.value.node_tree.nodes['Normal_texture'].image = self._load_maps('Normal')
+		self.value.node_tree.nodes['Displacement_texture'].image = self._load_maps('Displacement')
+
+
+class MaskMaterial(Material):
+	def __init__(self, name='mask', filename='mask', color=(1.0, 1.0, 1.0)):
+		self.filename = filename
+		self.color = color
+		Material.__init__(self, name, filename=self.filename)
+
+	def _update_nodes(self):
+		self.value.node_tree.nodes['RGB'].outputs[0].default_value[0] = self.color[0]
+		self.value.node_tree.nodes['RGB'].outputs[0].default_value[1] = self.color[1]
+		self.value.node_tree.nodes['RGB'].outputs[0].default_value[2] = self.color[2]
+
+
+class MaterialFactory:
+	"""
+	Class that produces materials based on the given name.
+	"""
+	def __init__(self):
+		self.materials = os.listdir('{}/Textures'.format(file_dir))
+		self.materials = [x for x in self.materials if os.path.isdir('{}/Textures/{}'.format(file_dir, x))]
+
+	def produce(self, name=None, color=None):
+		if name:
+			if name == 'mask':
+				if not color:
+					color = [1.0, 0.0, 0.0]
+				return MaskMaterial(name, color)
+			else:
+				name = name.lower().capitalize()
+				assert name in self.materials, "Unknown material {}, not in Textures folder".format(name)
+				return Material(name)
+		else:
+			return Material(np.random.choice(self.materials))
+
+
+if __name__ == '__main__':
+
+	from generator import *
+
+	material = MaterialFactory().produce('CONCrete')
+	f = CollectionFactory()
+	collection = f.produce(number=1)
+	building = ComposedBuilding(collection.collection)
+	building.make()
+	building.volumes[0].apply(material)

dataset/renderer.py

@@ -0,0 +1,75 @@
+import bpy, bmesh
+from math import radians
+import numpy as np
+import os
+import random
+import sys
+
+file_dir = os.path.dirname(__file__)
+sys.path.append(file_dir)
+
+from dataset_config import ENGINE
+
+
+class Renderer:
+	"""
+	Class that manages the scene rendering. Incomplete.
+	"""
+	def __init__(self):
+		self.engine = ENGINE
+		self._scene_name = bpy.data.scenes[-1].name
+
+	def render(self):
+		self._render()
+
+	def _render(self, filename='test'):
+		"""
+		Function that renders the scene.
+		:return:
+		"""
+		bpy.data.scenes[-1].render.engine = self.engine
+		bpy.ops.render.render()
+		if not IMG_SAVE in os.listdir(file_dir):
+			os.mkdir(file_dir + '/' + IMG_SAVE)
+		bpy.data.images["Render Result"].save_render(
+			'{}/{}.png'.format(IMG_SAVE, filename))
+
+	def _render_mask(self):
+		"""
+		Function that renders the scene as a one-channel mask.
+		:return:
+		"""
+		# update materials
+		bpy.data.scenes["Scene"].render.engine = 'BLENDER_RENDER'
+		bpy.ops.render.render()
+		if not MASK_SAVE in os.listdir(file_dir):
+			os.mkdir(file_dir + '/' + MASK_SAVE)
+		bpy.data.images["Render Result"].save_render(
+			'{}/{}.png'.format(MASK_SAVE, filename))
+		raise NotImplementedError
+
+	def _render_keypoints(self):
+		"""
+		Function that renders the scene as a one-channel mask of predefined
+		keypoints.
+		:return:
+		"""
+		raise NotImplementedError
+
+
+if __name__ == '__main__':
+
+	from generator import *
+	from material import MaterialFactory
+	from volume import CollectionFactory
+
+	f = CollectionFactory()
+	collection = f.produce(number=1)
+	building = ComposedBuilding(collection.collection)
+	building.make()
+	for v in building.volumes:
+		v.apply(MaterialFactory().produce())
+	r = Renderer()
+	r.render()
+
+