__init__.py 32 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034
  1. # ----------------------------------------------------------
  2. # File __init__.py
  3. # ----------------------------------------------------------
  4. # Imports
  5. # Check if this add-on is being reloaded
  6. if "bpy" in locals():
  7. # reloading .py files
  8. import bpy
  9. import importlib
  10. # from . import zs_renderscene as zsrs
  11. # importlib.reload(zsrs)
  12. # or if this is the first load of this add-on
  13. else:
  14. import bpy
  15. import json
  16. from pathlib import Path
  17. from dataclasses import dataclass
  18. from mathutils import Euler, Vector, Quaternion, Matrix
  19. import math
  20. import os
  21. import base64
  22. import numpy as np
  23. # from . import zs_renderscene as zsrs # noqa
  24. # Addon info
  25. bl_info = {
  26. "name": "ZS Stable Diffusion Connection V0.0.1",
  27. "author": "Sergiu <sergiu@zixelise.com>",
  28. "Version": (0, 0, 1),
  29. "blender": (4, 00, 0),
  30. "category": "Scene",
  31. "description": "Stable Diffusion Connection",
  32. }
  33. # -------------------------------------------------------------------
  34. # Functions
  35. # -------------------------------------------------------------------
  36. @dataclass
  37. class AssetData:
  38. path: str
  39. collection_name: str
  40. type: str
  41. product_asset = AssetData(
  42. path="sample_scene/01_Products", collection_name="01_Products", type="product"
  43. )
  44. element_asset = AssetData(
  45. path="sample_scene/02_Elements", collection_name="02_Elements", type="asset"
  46. )
  47. basic_shapes_asset = AssetData(
  48. path="sample_scene/03_BasicShapes",
  49. collection_name="03_BasicShapes",
  50. type="shape",
  51. )
  52. def convert_base64_string_to_object(base64_string):
  53. bytes = base64.b64decode(base64_string)
  54. string = bytes.decode("ascii")
  55. return json.loads(string)
  56. return string
  57. def load_scene():
  58. print("Loading Scene")
  59. # load scene data
  60. scene_data = load_scene_data()
  61. print(scene_data)
  62. # create parent collections
  63. create_parent_collections(product_asset.collection_name)
  64. create_parent_collections(element_asset.collection_name)
  65. create_parent_collections(basic_shapes_asset.collection_name)
  66. create_parent_collections("05_Cameras")
  67. # append products
  68. products_data = load_objects_data(scene_data, product_asset.type)
  69. for index, product in enumerate(products_data):
  70. append_objects(product, index, product_asset)
  71. # append elements
  72. elements_data = load_objects_data(scene_data, element_asset.type)
  73. for index, product in enumerate(elements_data):
  74. append_objects(product, index, element_asset)
  75. # append shapes
  76. basic_shapes_data = load_objects_data(scene_data, basic_shapes_asset.type)
  77. for index, product in enumerate(basic_shapes_data):
  78. append_objects(product, index, basic_shapes_asset)
  79. # set lighting
  80. set_environment(scene_data)
  81. # set camera
  82. create_cameras(scene_data)
  83. # rename outputs
  84. set_output_paths(
  85. "D://Git//ap-canvas-creation-module//03_blender//sd_blender//sample_scene//Renders",
  86. scene_data["project_id"],
  87. )
  88. # setup compositing
  89. set_cryptomatte_objects("01_Products", "mask_product")
  90. def invert_id_name(json_data):
  91. for obj in json_data["scene"]["objects"]:
  92. obj["name"], obj["id"] = obj["id"], obj["name"]
  93. return json_data
  94. def load_scene_data():
  95. # print("Loading Scene Data")
  96. # # load scene data
  97. # # to be replaced with actual data
  98. # # open scene_info.json
  99. # script_path = Path(__file__).resolve()
  100. # scene_data_path = script_path.parent / "sample_scene" / "scene_info.json"
  101. # with scene_data_path.open() as file:
  102. # scene_data = json.load(file)
  103. # print(scene_data)
  104. # return scene_data
  105. # load scene data
  106. print("Loading Scene Data")
  107. if bpy.context.scene.load_local_DB:
  108. loaded_scene_data = bpy.context.scene.config_string
  109. # check if loaded_scene_data is base64 encoded
  110. if loaded_scene_data.startswith("ey"):
  111. scene_data = convert_base64_string_to_object(loaded_scene_data)
  112. else:
  113. scene_data = json.loads(loaded_scene_data)
  114. else:
  115. scene_data = json.loads(bpy.context.scene.shot_info_ai)
  116. # invert_scene_data = invert_id_name(scene_data)
  117. return scene_data
  118. def load_objects_data(scene_data, object_type: str):
  119. print("Loading Assets Data")
  120. # load assets data
  121. objects_data = []
  122. for object in scene_data["scene"]["objects"]:
  123. if object["group_type"] == object_type:
  124. # get additional object data by id and combine with object data
  125. object_data = get_object_data_by_id(object["id"])
  126. if object_data:
  127. # temporary fix
  128. # object_data = get_object_data_by_id(object["name"])
  129. object.update(object_data)
  130. objects_data.append(object)
  131. else:
  132. print("Object not found in database", object["id"], object["name"])
  133. return objects_data
  134. # to be replaced with actual data
  135. def get_object_data_by_id(object_id: str):
  136. print("Getting Object Data")
  137. # open assets_database.json
  138. script_path = Path(__file__).resolve()
  139. assets_data_path = script_path.parent / "sample_scene" / "assets_database.json"
  140. with assets_data_path.open() as file:
  141. assets_data = json.load(file)
  142. # get object data by id
  143. for object in assets_data:
  144. if object["id"] == object_id:
  145. return object
  146. def get_hdri_data_by_id(object_id: str):
  147. print("Getting HDRI Data")
  148. # open assets_database.json
  149. script_path = Path(__file__).resolve()
  150. assets_data_path = script_path.parent / "sample_scene" / "lighting_database.json"
  151. with assets_data_path.open() as file:
  152. assets_data = json.load(file)
  153. # get object data by id
  154. for object in assets_data:
  155. if object["id"] == object_id:
  156. return object
  157. def append_objects(asset_info, index, asset_data: AssetData):
  158. print("Appending Objects")
  159. blendFileName = asset_info["name"]
  160. # visibleLayersJSONName = productInfo["Version"]
  161. # activeSwitchMaterials = json.dumps(productInfo["ActiveMaterials"])
  162. collectionName = blendFileName + "_" + str(index)
  163. append_active_layers(collectionName, asset_info, asset_data)
  164. # replace_switch_materials(shotInfo, productInfo["ActiveMaterials"])
  165. link_collection_to_collection(
  166. asset_data.collection_name, bpy.data.collections[collectionName]
  167. )
  168. def append_active_layers(newCollectionName, product_info, asset_data: AssetData):
  169. utility_collections = [
  170. "MaterialLibrary",
  171. "Animation_Controller",
  172. "Animation_Rig",
  173. "Animation_Target",
  174. ]
  175. # visible_layers = utility_collections + product_info["VisibleLayers"]
  176. visible_layers = utility_collections
  177. script_path = Path(__file__).resolve().parent
  178. filePath = str(
  179. script_path
  180. / asset_data.path
  181. / product_info["name"]
  182. / "BLEND"
  183. / (product_info["name"] + ".blend")
  184. )
  185. print(filePath)
  186. # delete all objects and collections from product collection
  187. create_parent_collections(newCollectionName)
  188. # append active collections
  189. with bpy.data.libraries.load(filePath) as (data_from, data_to):
  190. data_to.collections = []
  191. for name in data_from.collections:
  192. if name in visible_layers:
  193. data_to.collections.append(name)
  194. if "NonConfigurable" in name:
  195. data_to.collections.append(name)
  196. # link appended colections to newCollection
  197. for collection in data_to.collections:
  198. # try:
  199. # bpy.context.scene.collection.children.link(collection)
  200. # except:
  201. # print(collection)
  202. link_collection_to_collection(newCollectionName, collection)
  203. # hide utility collections
  204. for utilityCollectionName in utility_collections:
  205. if utilityCollectionName in collection.name:
  206. # rename utility collections
  207. collection.name = newCollectionName + "_" + utilityCollectionName
  208. hide_collection(collection)
  209. # need to redo this in the future
  210. if "Animation_Target" in collection.name:
  211. # set the x location
  212. collection.objects[0].location.x = product_info["properties"][
  213. "transform"
  214. ]["position"][0]
  215. # set the y location
  216. collection.objects[0].location.y = -product_info["properties"][
  217. "transform"
  218. ]["position"][2]
  219. # set the z location
  220. collection.objects[0].location.z = product_info["properties"][
  221. "transform"
  222. ]["position"][1]
  223. # collection.objects[0].location = product_info["properties"][
  224. # "transform"
  225. # ]["position"]
  226. # collection.objects[0].rotation_euler = product_info["properties"][
  227. # "transform"
  228. # ]["rotation"]
  229. rotation_in_degrees = product_info["properties"]["transform"][
  230. "rotation"
  231. ]
  232. rotation_in_degrees[0] = rotation_in_degrees[0] + 90
  233. # set object rotation in euler from radians
  234. collection.objects[0].rotation_euler = (
  235. math.radians(rotation_in_degrees[0]),
  236. math.radians(rotation_in_degrees[2]),
  237. math.radians(rotation_in_degrees[1]),
  238. )
  239. collection.objects[0].scale = product_info["properties"]["transform"][
  240. "scale"
  241. ]
  242. # make all objects in collection local
  243. for obj in bpy.data.collections[newCollectionName].all_objects:
  244. if obj.type == "MESH":
  245. obj.make_local()
  246. obj.data.make_local()
  247. # remove duplicated material slots
  248. mats = bpy.data.materials
  249. for obj in bpy.data.collections[newCollectionName].all_objects:
  250. for slot in obj.material_slots:
  251. part = slot.name.rpartition(".")
  252. if part[2].isnumeric() and part[0] in mats:
  253. slot.material = mats.get(part[0])
  254. def create_parent_collections(group_name: str):
  255. if collection_exists(group_name):
  256. remove_collection_and_objects(group_name)
  257. else:
  258. create_collection(group_name)
  259. def set_environment(scene_data):
  260. # Find the group named NG_Canvas_Background in the Blender world
  261. world = bpy.context.scene.world
  262. if world is not None:
  263. # Get the node group named NG_Canvas_Background
  264. node_group = world.node_tree.nodes.get("NG_Canvas_Background")
  265. if node_group is not None:
  266. # Set the group's properties from scene_data
  267. node_group.inputs[0].default_value = float(
  268. scene_data["scene"]["environment"]["lighting"]["rotation"]
  269. )
  270. node_group.inputs[1].default_value = float(
  271. scene_data["scene"]["environment"]["lighting"]["intensity"]
  272. )
  273. if scene_data["scene"]["environment"]["lighting"]["visible"] == True:
  274. node_group.inputs[2].default_value = 1.0
  275. else:
  276. node_group.inputs[2].default_value = 0.0
  277. node_group.inputs[3].default_value = (
  278. scene_data["scene"]["environment"]["background"]["color"]["r"],
  279. scene_data["scene"]["environment"]["background"]["color"]["g"],
  280. scene_data["scene"]["environment"]["background"]["color"]["b"],
  281. 1.0,
  282. )
  283. hdri_data = get_hdri_data_by_id(
  284. scene_data["scene"]["environment"]["lighting"]["id"]
  285. )
  286. # go inside the node group, and find the image texture node, and set the image to the one from the scene data
  287. for node in node_group.node_tree.nodes:
  288. if node.name == "environment_map":
  289. node.image = bpy.data.images.load(
  290. str(
  291. Path(__file__).resolve().parent
  292. / "sample_scene"
  293. / "05_Lighting"
  294. / "HDRI"
  295. / (hdri_data["name"] + ".exr")
  296. )
  297. )
  298. else:
  299. print("Group 'NG_Canvas_Background' not found")
  300. def calculate_focal_length(fov, film_height):
  301. # Convert FOV from degrees to radians
  302. fov_rad = math.radians(fov)
  303. # Calculate the focal length
  304. focal_length = (film_height / 2) / math.tan(fov_rad / 2)
  305. return focal_length
  306. def quaternion_multiply(q1, q2):
  307. """
  308. Multiplies two quaternions.
  309. q1 and q2 are arrays or lists of length 4.
  310. Returns the resulting quaternion as a NumPy array.
  311. """
  312. w1, x1, y1, z1 = q1
  313. w2, x2, y2, z2 = q2
  314. w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2
  315. x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2
  316. y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2
  317. z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2
  318. return np.array([w, x, y, z])
  319. def create_cameras(scene_data):
  320. # # Get the 05_Cameras collection, or create it if it doesn't exist
  321. collection_name = "05_Cameras"
  322. if collection_name in bpy.data.collections:
  323. collection = bpy.data.collections[collection_name]
  324. else:
  325. return
  326. # Assuming `scene_data` and `collection` are already defined
  327. for camera_data in scene_data["scene"]["cameras"]:
  328. # Create a new camera object
  329. bpy.ops.object.camera_add()
  330. # Get the newly created camera
  331. camera = bpy.context.object
  332. # Set the camera's name
  333. camera.name = camera_data["name"]
  334. # Set the camera's position
  335. position = camera_data["properties"]["transform"]["position"]
  336. camera.location.x = position[0]
  337. camera.location.y = -position[2]
  338. camera.location.z = position[1]
  339. # Set the camera's rotation
  340. # # euler
  341. # rotation_euler = camera_data["properties"]["transform"]["rotation"]
  342. # rotation_euler = Euler(
  343. # (
  344. # math.radians(rotation_euler[0] + 90),
  345. # math.radians(-rotation_euler[2]),
  346. # math.radians(rotation_euler[1]),
  347. # ),
  348. # "XYZ",
  349. # )
  350. # quaternion
  351. rotation_quat_data = camera_data["properties"]["transform"]["rotation"]
  352. # rotation_quat = (
  353. # rotation_quat[3],
  354. # rotation_quat[0],
  355. # rotation_quat[1],
  356. # rotation_quat[2],
  357. # )
  358. # new_quat = Quaternion((2**0.5 / 2, 2**0.5 / 2, 0.0, 0.0))
  359. # current_quat = Quaternion(rotation_quat)
  360. # result_quat = current_quat @ new_quat
  361. # rotation_quat = (result_quat.w, result_quat.x, result_quat.y, result_quat.z)
  362. rotation_quat = [
  363. rotation_quat_data[3],
  364. rotation_quat_data[0],
  365. rotation_quat_data[1],
  366. rotation_quat_data[2],
  367. ] # Example quaternion
  368. # new_quat = [2**0.5 / 2, 2**0.5 / 2, 0.0, 0.0]
  369. # result_quat = quaternion_multiply(rotation_quat, new_quat)
  370. # Example quaternion from GLTF: [x, y, z, w]
  371. gltf_quat = [0.0, 0.707, 0.0, 0.707]
  372. # Convert the GLTF quaternion to Blender space (Y-up to Z-up)
  373. converted_quat = Quaternion(
  374. [
  375. rotation_quat_data[3],
  376. rotation_quat_data[0],
  377. rotation_quat_data[1],
  378. rotation_quat_data[2],
  379. ]
  380. )
  381. # Define the camera correction quaternion (from GLTF file)
  382. camera_correction = Quaternion((2**0.5 / 2, 2**0.5 / 2, 0.0, 0.0))
  383. # Apply the camera correction to the converted quaternion
  384. corrected_quat = camera_correction @ converted_quat
  385. # Apply the corrected quaternion to the camera's rotation
  386. camera.rotation_mode = "QUATERNION"
  387. camera.rotation_quaternion = corrected_quat
  388. # camera.rotation_mode = "QUATERNION"
  389. # camera.rotation_quaternion = rotation_quat
  390. # Apply the local rotation to the camera
  391. # Set the camera's lens properties
  392. lens = camera_data["properties"]["lens"]
  393. type_mapping = {
  394. "PERSPECTIVE": "PERSP",
  395. "ORTHOGRAPHIC": "ORTHO",
  396. "PANORAMIC": "PANO",
  397. }
  398. camera.data.lens = lens["focalLength"]
  399. camera.data.type = type_mapping.get(lens["type"].upper(), "PERSP")
  400. camera.data.clip_start = lens["near"]
  401. camera.data.clip_end = lens["far"]
  402. # Add the camera to the 05_Cameras collection
  403. collection.objects.link(camera)
  404. bpy.context.scene.collection.objects.unlink(camera)
  405. # Set the camera as the active camera if "active" is true
  406. if camera_data["properties"]["active"]:
  407. bpy.context.scene.camera = camera
  408. def set_output_paths(base_path, project_name):
  409. # check if folder exist, if not create it
  410. folder_path = base_path + "//" + project_name
  411. if not os.path.exists(folder_path):
  412. os.makedirs(folder_path)
  413. # Get the current scene
  414. scene = bpy.context.scene
  415. # Check if the scene has a compositor node tree
  416. if scene.node_tree:
  417. # Iterate over all nodes in the node tree
  418. for node in scene.node_tree.nodes:
  419. # Check if the node is an output node
  420. if node.type == "OUTPUT_FILE":
  421. # Set the base path of the output node
  422. node.base_path = folder_path
  423. # Iterate over all file slots of the output node
  424. # for file_slot in node.file_slots:
  425. # # Set the path of the file slot
  426. # file_slot.path = project_name
  427. def set_cryptomatte_objects(collection_name, node_name):
  428. # Get all objects in the specified collection
  429. objects = bpy.data.collections[collection_name].all_objects
  430. # Convert the objects to a list
  431. object_names = [obj.name for obj in objects]
  432. print(object_names)
  433. # Get the current scene
  434. scene = bpy.context.scene
  435. # Check if the scene has a compositor node tree
  436. if scene.node_tree:
  437. # Iterate over all nodes in the node tree
  438. for node in scene.node_tree.nodes:
  439. # Check if the node is a Cryptomatte node with the specified name
  440. if node.type == "CRYPTOMATTE_V2" and node.name == node_name:
  441. # Set the Matte ID objects of the node
  442. node.matte_id = ",".join(object_names)
  443. # -------------------------------------------------------------------
  444. # Utilities
  445. # -------------------------------------------------------------------
  446. def remove_collection_and_objects(collection_name):
  447. oldObjects = list(bpy.data.collections[collection_name].all_objects)
  448. for obj in oldObjects:
  449. bpy.data.objects.remove(obj, do_unlink=True)
  450. old_collection = bpy.data.collections[collection_name]
  451. if old_collection is not None:
  452. old_collection_names = get_subcollection_names(old_collection)
  453. else:
  454. print("Collection not found.")
  455. # print line break
  456. print("-----------------------------------------------------------------")
  457. print(old_collection_names)
  458. print("-----------------------------------------------------------------")
  459. for old_collection_name in old_collection_names:
  460. for collection in bpy.data.collections:
  461. if collection.name == old_collection_name:
  462. bpy.data.collections.remove(collection)
  463. bpy.ops.outliner.orphans_purge(
  464. do_local_ids=True, do_linked_ids=True, do_recursive=True
  465. )
  466. def get_subcollection_names(collection):
  467. subcollection_names = []
  468. for child in collection.children:
  469. subcollection_names.append(child.name)
  470. subcollection_names.extend(get_subcollection_names(child))
  471. return subcollection_names
  472. def select_objects_in_collection(collection):
  473. """Recursively select all objects in the given collection and its subcollections."""
  474. for obj in collection.objects:
  475. obj.select_set(True)
  476. for subcollection in collection.children:
  477. select_objects_in_collection(subcollection)
  478. def link_collection_to_collection(parentCollectionName, childCollection):
  479. if bpy.context.scene.collection.children:
  480. parentCollection = bpy.context.scene.collection.children.get(
  481. parentCollectionName
  482. )
  483. # Add it to the main collection
  484. # childCollection = bpy.context.scene.collection.children.get(childCollection)
  485. parentCollection.children.link(childCollection)
  486. # if child collection is in scene collection unlink it
  487. if bpy.context.scene.collection.children.get(childCollection.name):
  488. bpy.context.scene.collection.children.unlink(childCollection)
  489. # bpy.context.scene.collection.children.unlink(childCollection)
  490. # link collection to collection
  491. def link_collection_to_collection_old(parentCollectionName, childCollection):
  492. if bpy.context.scene.collection.children:
  493. parentCollection = bpy.context.scene.collection.children.get(
  494. parentCollectionName
  495. )
  496. # Add it to the main collection
  497. try:
  498. childCollection = bpy.data.collections[childCollection]
  499. except:
  500. print("Collection not found.")
  501. return
  502. parentCollection.children.link(childCollection)
  503. bpy.context.scene.collection.children.unlink(childCollection)
  504. # function that checks if a collection exists
  505. def collection_exists(collection_name):
  506. return collection_name in bpy.data.collections
  507. # function that creates a new collection and adds it to the scene
  508. def create_collection(collection_name):
  509. new_collection = bpy.data.collections.new(collection_name)
  510. bpy.context.scene.collection.children.link(new_collection)
  511. def hide_collection(collection):
  512. collection.hide_render = True
  513. collection.hide_viewport = True
  514. def check_if_selected_objects_have_parent(self):
  515. for obj in bpy.context.selected_objects:
  516. if obj.parent is None:
  517. message = f"Object {obj.name} has no parent"
  518. self.report({"ERROR"}, message)
  519. return False
  520. return True
  521. def add_rig_controller_to_selection(self):
  522. # add "Rig_Controller_Main" to the selection
  523. has_controller_object = False
  524. for obj in bpy.data.objects:
  525. if obj.name == "Rig_Controller_Main":
  526. if obj.hide_viewport:
  527. self.report({"ERROR"}, "Rig_Controller_Main is hidden")
  528. return has_controller_object
  529. obj.select_set(True)
  530. # if object is not visible, make it visible
  531. has_controller_object = True
  532. return has_controller_object
  533. if not has_controller_object:
  534. message = f"Rig_Controller_Main not found"
  535. self.report({"ERROR"}, message)
  536. print("Rig_Controller_Main not found")
  537. return has_controller_object
  538. def select_objects_in_collection(collection):
  539. """Recursively select all objects in the given collection and its subcollections."""
  540. for obj in collection.objects:
  541. obj.select_set(True)
  542. for subcollection in collection.children:
  543. select_objects_in_collection(subcollection)
  544. def check_if_object_has_principled_material(obj):
  545. for slot in obj.material_slots:
  546. # if more than one slot is principled, return
  547. for node in slot.material.node_tree.nodes:
  548. # print(node.type)
  549. if node.type == "BSDF_PRINCIPLED":
  550. return True
  551. else:
  552. print("Object has no principled material", obj.name)
  553. return False
  554. return True
  555. def unhide_all_objects():
  556. # show all objects using operator
  557. bpy.ops.object.hide_view_clear()
  558. # -------------------------------------------------------------------
  559. # Scene optimization
  560. # -------------------------------------------------------------------
  561. def export_non_configurable_to_fbx():
  562. # Get the current .blend file path
  563. blend_filepath = bpy.data.filepath
  564. if not blend_filepath:
  565. print("Save the .blend file first.")
  566. return
  567. # Get the parent directory of the .blend file
  568. blend_dir = os.path.dirname(blend_filepath)
  569. parent_dir = os.path.dirname(blend_dir)
  570. blend_filename = os.path.splitext(os.path.basename(blend_filepath))[0]
  571. # Create the FBX export path
  572. fbx_export_path = os.path.join(parent_dir, "FBX", f"{blend_filename}.fbx")
  573. # Ensure the FBX directory exists
  574. os.makedirs(os.path.dirname(fbx_export_path), exist_ok=True)
  575. # Deselect all objects
  576. bpy.ops.object.select_all(action="DESELECT")
  577. # Select all objects in the NonConfigurable collection
  578. collection_name = "NonConfigurable"
  579. if collection_name in bpy.data.collections:
  580. collection = bpy.data.collections[collection_name]
  581. for obj in collection.objects:
  582. obj.select_set(True)
  583. else:
  584. print(f"Collection '{collection_name}' not found.")
  585. return
  586. def export_scene_to_fbx(self):
  587. # Ensure the .blend file is saved
  588. if not bpy.data.is_saved:
  589. print("Save the .blend file first.")
  590. self.report({"ERROR"}, "Save the .blend file first.")
  591. return
  592. # Get the current .blend file path
  593. blend_filepath = bpy.data.filepath
  594. if not blend_filepath:
  595. print("Unable to get the .blend file path.")
  596. return
  597. # Get the parent directory of the .blend file
  598. blend_dir = os.path.dirname(blend_filepath)
  599. parent_dir = os.path.dirname(blend_dir)
  600. blend_filename = os.path.splitext(os.path.basename(blend_filepath))[0]
  601. # Create the FBX export path
  602. fbx_export_path = os.path.join(parent_dir, "FBX", f"{blend_filename}.fbx")
  603. # Ensure the FBX directory exists
  604. os.makedirs(os.path.dirname(fbx_export_path), exist_ok=True)
  605. # unhide all objects
  606. unhide_all_objects()
  607. # Deselect all objects
  608. bpy.ops.object.select_all(action="DESELECT")
  609. # Select all objects in the NonConfigurable collection and its subcollections
  610. collection_name = "NonConfigurable"
  611. if collection_name in bpy.data.collections:
  612. collection = bpy.data.collections[collection_name]
  613. select_objects_in_collection(collection)
  614. else:
  615. print(f"Collection '{collection_name}' not found.")
  616. return
  617. # check if all objects selected have a parent, if not return
  618. if not check_if_selected_objects_have_parent(self):
  619. return
  620. if not add_rig_controller_to_selection(self):
  621. return
  622. # Export selected objects to FBX
  623. bpy.ops.export_scene.fbx(
  624. filepath=fbx_export_path,
  625. use_selection=True,
  626. global_scale=1.0,
  627. apply_unit_scale=True,
  628. bake_space_transform=False,
  629. object_types={"MESH", "ARMATURE", "EMPTY"},
  630. use_mesh_modifiers=True,
  631. mesh_smooth_type="FACE",
  632. use_custom_props=True,
  633. # bake_anim=False,
  634. )
  635. print(f"Exported to {fbx_export_path}")
  636. def export_scene_to_glb(self):
  637. # Ensure the .blend file is saved
  638. if not bpy.data.is_saved:
  639. print("Save the .blend file first.")
  640. self.report({"ERROR"}, "Save the .blend file first.")
  641. return
  642. # Get the current .blend file path
  643. blend_filepath = bpy.data.filepath
  644. if not blend_filepath:
  645. print("Unable to get the .blend file path.")
  646. return
  647. # Get the parent directory of the .blend file
  648. blend_dir = os.path.dirname(blend_filepath)
  649. parent_dir = os.path.dirname(blend_dir)
  650. blend_filename = os.path.splitext(os.path.basename(blend_filepath))[0]
  651. # Create the GLB export path
  652. glb_export_path = os.path.join(parent_dir, "WEB", f"{blend_filename}.glb")
  653. # Ensure the GLB directory exists
  654. os.makedirs(os.path.dirname(glb_export_path), exist_ok=True)
  655. # unhide all objects
  656. unhide_all_objects()
  657. # Deselect all objects
  658. bpy.ops.object.select_all(action="DESELECT")
  659. # Select all objects in the NonConfigurable collection and its subcollections
  660. collection_name = "WebGL"
  661. if collection_name in bpy.data.collections:
  662. collection = bpy.data.collections[collection_name]
  663. select_objects_in_collection(collection)
  664. else:
  665. print(f"Collection '{collection_name}' not found.")
  666. return
  667. if not check_if_selected_objects_have_parent(self):
  668. return
  669. # check if all objects selected have a parent, if not return
  670. if not add_rig_controller_to_selection(self):
  671. return
  672. # # for each selected objects, check if the the material is principled, if not return
  673. # for obj in bpy.context.selected_objects:
  674. # if not check_if_object_has_principled_material(obj):
  675. # message = f"Object {obj.name} has no principled material"
  676. # self.report({"ERROR"}, message)
  677. # return
  678. # Export selected objects to GLB
  679. bpy.ops.export_scene.gltf(
  680. filepath=glb_export_path,
  681. export_format="GLB",
  682. use_selection=True,
  683. export_apply=True,
  684. export_animations=False,
  685. export_yup=True,
  686. export_cameras=False,
  687. export_lights=False,
  688. export_materials="EXPORT",
  689. export_normals=True,
  690. export_tangents=True,
  691. export_morph=False,
  692. export_skins=False,
  693. export_draco_mesh_compression_enable=False,
  694. export_draco_mesh_compression_level=6,
  695. export_draco_position_quantization=14,
  696. export_draco_normal_quantization=10,
  697. export_draco_texcoord_quantization=12,
  698. export_draco_color_quantization=10,
  699. export_draco_generic_quantization=12,
  700. export_keep_originals=False,
  701. export_texture_dir="",
  702. )
  703. print(f"Exported to {glb_export_path}")
  704. # -------------------------------------------------------------------
  705. # Operators
  706. # -------------------------------------------------------------------
  707. # load scene operator
  708. class ZSSD_OT_LoadScene(bpy.types.Operator):
  709. bl_idname = "zs_sd_loader.load_scene"
  710. bl_label = "Load Scene"
  711. bl_description = "Load Scene"
  712. def execute(self, context):
  713. load_scene()
  714. return {"FINISHED"}
  715. # canvas exporter operator
  716. class ZSSD_OT_ExportAssets(bpy.types.Operator):
  717. bl_idname = "zs_canvas.export_assets"
  718. bl_label = "Export Assets"
  719. bl_description = "Export Scene Assets to FBX and GLB"
  720. def execute(self, context):
  721. export_scene_to_fbx(self)
  722. export_scene_to_glb(self)
  723. return {"FINISHED"}
  724. # parent class for panels
  725. class ZSSDPanel:
  726. bl_space_type = "VIEW_3D"
  727. bl_region_type = "UI"
  728. bl_category = "ZS SD Loader"
  729. # -------------------------------------------------------------------
  730. # Draw
  731. # -------------------------------------------------------------------
  732. # Panels
  733. class ZSSD_PT_Main(ZSSDPanel, bpy.types.Panel):
  734. bl_label = "SD Loader"
  735. def draw(self, context):
  736. layout = self.layout
  737. scene = context.scene
  738. col = layout.column()
  739. self.is_connected = False
  740. col.label(text="Stable Diffusion Connection")
  741. col.prop(context.scene, "load_local_DB")
  742. col.prop(context.scene, "config_string")
  743. # load scene button
  744. col.operator("zs_sd_loader.load_scene", text="Load Scene")
  745. col.separator()
  746. # export assets button
  747. col.operator("zs_canvas.export_assets", text="Export Assets")
  748. # modify after making products
  749. blender_classes = [
  750. ZSSD_PT_Main,
  751. ZSSD_OT_LoadScene,
  752. ZSSD_OT_ExportAssets,
  753. ]
  754. def register():
  755. # register classes
  756. for blender_class in blender_classes:
  757. bpy.utils.register_class(blender_class)
  758. bpy.types.Scene.shot_info_ai = bpy.props.StringProperty(
  759. name="Shot Info",
  760. )
  761. bpy.types.Scene.config_string = bpy.props.StringProperty( # type: ignore
  762. name="Configuration String",
  763. )
  764. bpy.types.Scene.load_local_DB = bpy.props.BoolProperty( # type: ignore
  765. name="Load Local DB",
  766. )
  767. # Has to be afqter class registering to correctly register property
  768. # register global properties
  769. # register list
  770. # list data
  771. # bpy.types.Scene.zs_product_list = bpy.props.CollectionProperty(
  772. # type=ZS_Product_ListItem)
  773. # current item in list
  774. def unregister():
  775. # unregister classes
  776. for blender_class in blender_classes:
  777. bpy.utils.unregister_class(blender_class)
  778. # unregister global properties
  779. del bpy.types.Scene.shot_info_ai
  780. del bpy.types.Scene.config_string
  781. del bpy.types.Scene.load_local_DB
  782. # unregister list items
  783. # del bpy.types.Scene.my_list
  784. # del bpy.types.Scene.product_product_list_index