OpenDroneMap-ODM/opendm/gltf.py

310 wiersze
10 KiB
Python

import os
import rasterio
from rasterio.io import MemoryFile
import warnings
import numpy as np
import pygltflib
from opendm import system
from opendm import io
warnings.filterwarnings("ignore", category=rasterio.errors.NotGeoreferencedWarning)
def load_obj(obj_path, _info=print):
if not os.path.isfile(obj_path):
raise IOError("Cannot open %s" % obj_path)
obj_base_path = os.path.dirname(os.path.abspath(obj_path))
obj = {
'materials': {},
}
vertices = []
uvs = []
normals = []
faces = {}
current_material = "_"
with open(obj_path) as f:
_info("Loading %s" % obj_path)
for line in f:
if line.startswith("mtllib "):
# Materials
mtl_file = "".join(line.split()[1:]).strip()
obj['materials'].update(load_mtl(mtl_file, obj_base_path, _info=_info))
elif line.startswith("v "):
# Vertices
vertices.append(list(map(float, line.split()[1:4])))
elif line.startswith("vt "):
# UVs
uvs.append(list(map(float, line.split()[1:3])))
elif line.startswith("vn "):
normals.append(list(map(float, line.split()[1:4])))
elif line.startswith("usemtl "):
mtl_name = "".join(line.split()[1:]).strip()
if not mtl_name in obj['materials']:
raise Exception("%s material is missing" % mtl_name)
current_material = mtl_name
elif line.startswith("f "):
if current_material not in faces:
faces[current_material] = []
a,b,c = line.split()[1:]
if a.count("/") == 2:
av, at, an = map(int, a.split("/")[0:3])
bv, bt, bn = map(int, b.split("/")[0:3])
cv, ct, cn = map(int, c.split("/")[0:3])
faces[current_material].append((av - 1, bv - 1, cv - 1, at - 1, bt - 1, ct - 1, an - 1, bn - 1, cn - 1))
else:
av, at = map(int, a.split("/")[0:2])
bv, bt = map(int, b.split("/")[0:2])
cv, ct = map(int, c.split("/")[0:2])
faces[current_material].append((av - 1, bv - 1, cv - 1, at - 1, bt - 1, ct - 1))
obj['vertices'] = np.array(vertices, dtype=np.float32)
obj['uvs'] = np.array(uvs, dtype=np.float32)
obj['normals'] = np.array(normals, dtype=np.float32)
obj['faces'] = faces
obj['materials'] = convert_materials_to_jpeg(obj['materials'])
return obj
def convert_materials_to_jpeg(materials):
min_value = 0
value_range = 0
skip_conversion = False
for mat in materials:
image = materials[mat]
# Stop here, assuming all other materials are also uint8
if image.dtype == np.uint8:
skip_conversion = True
break
# Find common min/range values
try:
data_range = np.iinfo(image.dtype)
min_value = min(min_value, 0)
value_range = max(value_range, float(data_range.max) - float(data_range.min))
except ValueError:
# For floats use the actual range of the image values
min_value = min(min_value, float(image.min()))
value_range = max(value_range, float(image.max()) - min_value)
if value_range == 0:
value_range = 255 # Should never happen
for mat in materials:
image = materials[mat]
if not skip_conversion:
image = image.astype(np.float32)
image -= min_value
image *= 255.0 / value_range
np.around(image, out=image)
image[image > 255] = 255
image[image < 0] = 0
image = image.astype(np.uint8)
with MemoryFile() as memfile:
bands, h, w = image.shape
bands = min(3, bands)
with memfile.open(driver='JPEG', jpeg_quality=90, count=bands, width=w, height=h, dtype=rasterio.dtypes.uint8) as dst:
for b in range(1, min(3, bands) + 1):
dst.write(image[b - 1], b)
memfile.seek(0)
materials[mat] = memfile.read()
return materials
def load_mtl(mtl_file, obj_base_path, _info=print):
mtl_file = os.path.join(obj_base_path, mtl_file)
if not os.path.isfile(mtl_file):
raise IOError("Cannot open %s" % mtl_file)
mats = {}
current_mtl = ""
with open(mtl_file) as f:
for line in f:
if line.startswith("newmtl "):
current_mtl = "".join(line.split()[1:]).strip()
elif line.startswith("map_Kd ") and current_mtl:
map_kd_filename = "".join(line.split()[1:]).strip()
map_kd = os.path.join(obj_base_path, map_kd_filename)
if not os.path.isfile(map_kd):
raise IOError("Cannot open %s" % map_kd)
_info("Loading %s" % map_kd_filename)
with rasterio.open(map_kd, 'r') as src:
mats[current_mtl] = src.read()
return mats
def paddedBuffer(buf, boundary):
r = len(buf) % boundary
if r == 0:
return buf
pad = boundary - r
return buf + b'\x00' * pad
def obj2glb(input_obj, output_glb, rtc=(None, None), draco_compression=True, _info=print):
_info("Converting %s --> %s" % (input_obj, output_glb))
obj = load_obj(input_obj, _info=_info)
vertices = obj['vertices']
uvs = obj['uvs']
# Flip Y
uvs = (([0, 1] - (uvs * [0, 1])) + uvs * [1, 0]).astype(np.float32)
normals = obj['normals']
binary = b''
accessors = []
bufferViews = []
primitives = []
materials = []
textures = []
images = []
bufOffset = 0
def addBufferView(buf, target=None):
nonlocal bufferViews, bufOffset
bufferViews += [pygltflib.BufferView(
buffer=0,
byteOffset=bufOffset,
byteLength=len(buf),
target=target,
)]
bufOffset += len(buf)
return len(bufferViews) - 1
for material in obj['faces'].keys():
faces = obj['faces'][material]
faces = np.array(faces, dtype=np.uint32)
prim_vertices = vertices[faces[:,0:3].flatten()]
prim_uvs = uvs[faces[:,3:6].flatten()]
if faces.shape[1] == 9:
prim_normals = normals[faces[:,6:9].flatten()]
normals_blob = prim_normals.tobytes()
else:
prim_normals = None
normals_blob = None
vertices_blob = prim_vertices.tobytes()
uvs_blob = prim_uvs.tobytes()
binary += vertices_blob + uvs_blob
if normals_blob is not None:
binary += normals_blob
verticesBufferView = addBufferView(vertices_blob, pygltflib.ARRAY_BUFFER)
uvsBufferView = addBufferView(uvs_blob, pygltflib.ARRAY_BUFFER)
normalsBufferView = None
if normals_blob is not None:
normalsBufferView = addBufferView(normals_blob, pygltflib.ARRAY_BUFFER)
accessors += [
pygltflib.Accessor(
bufferView=verticesBufferView,
componentType=pygltflib.FLOAT,
count=len(prim_vertices),
type=pygltflib.VEC3,
max=prim_vertices.max(axis=0).tolist(),
min=prim_vertices.min(axis=0).tolist(),
),
pygltflib.Accessor(
bufferView=uvsBufferView,
componentType=pygltflib.FLOAT,
count=len(prim_uvs),
type=pygltflib.VEC2,
max=prim_uvs.max(axis=0).tolist(),
min=prim_uvs.min(axis=0).tolist(),
),
]
if prim_normals is not None:
accessors += [
pygltflib.Accessor(
bufferView=normalsBufferView,
componentType=pygltflib.FLOAT,
count=len(prim_normals),
type=pygltflib.VEC3,
max=prim_normals.max(axis=0).tolist(),
min=prim_normals.min(axis=0).tolist(),
)
]
primitives += [pygltflib.Primitive(
attributes=pygltflib.Attributes(POSITION=verticesBufferView, TEXCOORD_0=uvsBufferView, NORMAL=normalsBufferView), material=len(primitives)
)]
for material in obj['faces'].keys():
texture_blob = paddedBuffer(obj['materials'][material], 4)
binary += texture_blob
textureBufferView = addBufferView(texture_blob)
images += [pygltflib.Image(bufferView=textureBufferView, mimeType="image/jpeg")]
textures += [pygltflib.Texture(source=len(images) - 1, sampler=0)]
mat = pygltflib.Material(pbrMetallicRoughness=pygltflib.PbrMetallicRoughness(baseColorTexture=pygltflib.TextureInfo(index=len(textures) - 1), metallicFactor=0, roughnessFactor=1),
alphaMode=pygltflib.OPAQUE)
mat.extensions = {
'KHR_materials_unlit': {}
}
materials += [mat]
gltf = pygltflib.GLTF2(
scene=0,
scenes=[pygltflib.Scene(nodes=[0])],
nodes=[pygltflib.Node(mesh=0)],
meshes=[pygltflib.Mesh(
primitives=primitives
)],
materials=materials,
textures=textures,
samplers=[pygltflib.Sampler(magFilter=pygltflib.LINEAR, minFilter=pygltflib.LINEAR)],
images=images,
accessors=accessors,
bufferViews=bufferViews,
buffers=[pygltflib.Buffer(byteLength=len(binary))],
)
gltf.extensionsRequired = ['KHR_materials_unlit']
gltf.extensionsUsed = ['KHR_materials_unlit']
if rtc != (None, None) and len(rtc) >= 2:
gltf.extensionsUsed.append('CESIUM_RTC')
gltf.extensions = {
'CESIUM_RTC': {
'center': [float(rtc[0]), float(rtc[1]), 0.0]
}
}
gltf.set_binary_blob(binary)
_info("Writing...")
gltf.save(output_glb)
_info("Wrote %s" % output_glb)
if draco_compression:
_info("Compressing with draco")
try:
compressed_glb = io.related_file_path(output_glb, postfix="_compressed")
system.run('draco_transcoder -i "{}" -o "{}" -qt 16 -qp 16'.format(output_glb, compressed_glb))
if os.path.isfile(compressed_glb) and os.path.isfile(output_glb):
os.remove(output_glb)
os.rename(compressed_glb, output_glb)
except Exception as e:
log.ODM_WARNING("Cannot compress GLB with draco: %s" % str(e))