src/manhole_model.py

import io
import numpy as np
from PIL import Image as PIL_Image # gltflibのImageと被るので別名にする。
from PIL import ImageDraw as PIL_ImageDraw # 上記と同じ命名規則にする。
from scipy.spatial import Delaunay
import cv2
import os
import struct
import matplotlib.pyplot as plt
import triangle
import uuid

from gltflib import (
    GLTF, GLTFModel, Asset, Scene, Node, Mesh, Primitive, Attributes, Buffer, BufferView, Image, Texture, TextureInfo, Material, Sampler, Accessor, AccessorType,
    BufferTarget, ComponentType, GLBResource, FileResource, PBRMetallicRoughness)

# Common configuration information
# Parts that are independent of the binary section's offset and size can be predefined.

# Asset
asset=Asset()

# Image
images=[
    Image(mimeType='image/jpeg', bufferView=4),
]

# Sampler
samplers = [Sampler(magFilter=9728, minFilter=9984)] # magFilter:最近傍フィルタリング、minFilter:ミップマップ+最近傍フィルタリング

# Texture
textures = [
    Texture(name='Main',sampler=0,source=0),
]

# Material
materials = [
    Material(
        pbrMetallicRoughness=PBRMetallicRoughness(
            baseColorTexture=TextureInfo(index=0),
            metallicFactor=0,
            roughnessFactor=1
        ),
        name='Material0',
        alphaMode='OPAQUE',
        doubleSided=True
    ),
]

# Mesh
meshes = [
    Mesh(name='Main', primitives=[Primitive(attributes=Attributes(POSITION=0, NORMAL=1,TEXCOORD_0=2),
                                            indices=3, material=0, mode=4)]),
]

# Scene
scene = 0
scenes = [Scene(name='Scene', nodes=[0])]

def create_manhole_model(original_img_bytearray):
    
    # Image
    img = PIL_Image.open(original_img_bytearray).convert('RGB')
    
    # Add edge color space.
    # (edge color is decided at 'Decide edge color' section)
    temp_img = PIL_Image.new('RGB', (img.size[0] + 3, img.size[1] + 3), 'black')
    temp_img.paste(img, (3, 3))
    img = temp_img.copy()

    # Get coordinates of manhole
    ret, mask = cv2.threshold(cv2.cvtColor(np.array(img), cv2.COLOR_RGB2GRAY), 0, 255, cv2.THRESH_OTSU+cv2.THRESH_BINARY_INV)
    edges = cv2.Canny(mask , 50, 150, apertureSize=3)
    circle = cv2.HoughCircles(edges, 
                            cv2.HOUGH_GRADIENT, 1, int(max(img.size)*2), 
                            param1=50, param2=30, minRadius=0, maxRadius=0)[0][0]
    
    coordinate_list = []
    num_polygons = 360
    for i in range(num_polygons):
        coordinate_list.append((
            int(circle[0] + np.cos((i / num_polygons) * (2 * np.pi)) * circle[2]),
            int(circle[1] + np.sin((i / num_polygons) * (2 * np.pi)) * circle[2])
        ))
    coordinate_list = [coordinate_list]

    # Decide edge color
    coordinate_colors = list(map(lambda x: img.getpixel(x), coordinate_list[0]))
    edge_color = tuple(np.mean(np.array(coordinate_colors), axis=0).astype(np.uint8))

    draw = PIL_ImageDraw.Draw(img)
    draw.point(tuple((i, j) for i in range(3) for j in range(3)), fill=edge_color)
    
    # image binary data
    img_ext = "JPEG"
    img_bytearray = io.BytesIO()
    img.save(img_bytearray, format=img_ext, quality=95)
    img_bytearray = img_bytearray.getvalue()
    img_bytelen = len(img_bytearray)

    # calculate scale of 3d model
    scale_factor = np.power(img.size[0] * img.size[1], 0.5)
    scale = (img.size[0] / scale_factor, img.size[1] / scale_factor, 0.4)
    
    # Thickness of manhole
    thickness = 0.2 # manhole

    # Vertex data(POSITION)
    front_vertex_list = []
    back_vertex_list = []
    for coordinates in coordinate_list:
        front_vertices = []
        back_vertices = []
        # Be aware that the Y-axis direction is inverted between images and GLB files
        for coordinate in coordinates:
            front_vertices.append((coordinate[0] * 2 / img.size[0] - 1.0, -(coordinate[1] * 2 / img.size[1] - 1.0), thickness))
            back_vertices.append((coordinate[0] * 2 / img.size[0] - 1.0, -(coordinate[1] * 2 / img.size[1] - 1.0), -thickness))

        front_vertex_list.append(front_vertices)
        back_vertex_list.append(back_vertices)

    vertices = front_vertex_list[0] + back_vertex_list[0]
    vertex_bytearray = bytearray()

    # Set vertices with 'Front+Back' and 'Edge'
    for i in range(2): # first for 'Front+Back', second for 'Edge'
        for vertex in vertices:
            for value in vertex:
                vertex_bytearray.extend(struct.pack('f', value))

    vertex_bytelen = len(vertex_bytearray)
    mins = [min([vertex[i] for vertex in vertices]) for i in range(3)]
    maxs = [max([vertex[i] for vertex in vertices]) for i in range(3)]
    
     # Normal data(NORMAL)
    normals = [( 0.0,  0.0, 1.0)] * len(front_vertex_list[0]) + [( 0.0,  0.0, -1.0)] * len(back_vertex_list[0])

    normal_bytearray = bytearray()

    # Consider not only 'Front+Back' but also 'Edge'
    for i in range(2): # first for 'Front+Back', second for 'Edge'
        for normal in normals:
            for value in normal:
                normal_bytearray.extend(struct.pack('f', value))

    normal_bytelen = len(normal_bytearray)

    # Texture coordinates(TEXCOORD_0)
    texcoord_0s = [
        ((vertex[0] + 1.0) / 2.0, 1.0 - ((vertex[1] + 1.0) / 2.0) ) for vertex in vertices
    ]
    texcoord_0_bytearray = bytearray()

    # 'Front+Back'
    for texcoord_0 in texcoord_0s:
        for value in texcoord_0:
            texcoord_0_bytearray.extend(struct.pack('f', value))
    # 'Edge'
    for texcoord_0 in texcoord_0s:
        for value in texcoord_0:
            texcoord_0_bytearray.extend(struct.pack('f', 0.0))
            
    texcoord_0_bytelen = len(texcoord_0_bytearray)

    # Vertex indices
    polygon = {
        'vertices': np.array(front_vertex_list[0])[:, :2],
        'segments': np.array([( i, (i + 1) % (len(front_vertex_list[0])) ) for i in range(len(front_vertex_list[0]))])
    }
    triangulate_result = triangle.triangulate(polygon, 'p')

    vertex_indices = []
    
    # Front
    vertex_indices.extend(list(np.array(triangulate_result['triangles']).flatten())) 

    # Back
    temp_list = list((np.array(triangulate_result['triangles'])+len(front_vertex_list[0])).flatten())
    # Swap vertex indices to get Back vertex indices from Front vertex indices 
    def swap_elements(lst):
        if len(lst) < 2:
            return lst
        for i in range(len(lst) - 1):
            if i % 3 == 1:
                lst[i], lst[i + 1] = lst[i + 1], lst[i]

        return lst
    
    temp_list = swap_elements(temp_list)
    vertex_indices.extend(temp_list) # Back
    
    # Edge 
    vertex_indices.extend(list(np.array([[len(vertices) + i, 
                                        len(vertices) + (i + 1) % len(front_vertex_list[0]), 
                                        len(vertices) + len(front_vertex_list[0]) + i] 
                                        for i in range(len(front_vertex_list[0]))]).flatten())) 
    vertex_indices.extend(list(np.array([[len(vertices) + len(front_vertex_list[0]) + (i + 1) % len(front_vertex_list[0]), 
                                        len(vertices) + len(front_vertex_list[0]) + i, 
                                        len(vertices) + (i + 1) % len(front_vertex_list[0])] 
                                        for i in range(len(front_vertex_list[0]))]).flatten()))

    vertex_index_bytearray = bytearray()
    for value in vertex_indices:
        vertex_index_bytearray.extend(struct.pack('H', value))
    vertex_index_bytelen = len(vertex_index_bytearray)
    
    # Concatenate binar data
    bytearray_list = [
        vertex_bytearray,
        normal_bytearray,
        texcoord_0_bytearray,
        vertex_index_bytearray,
        img_bytearray,
    ]
    bytelen_list = [
        vertex_bytelen,
        normal_bytelen,
        texcoord_0_bytelen,
        vertex_index_bytelen,
        img_bytelen,
    ]
    bytelen_cumsum_list = list(np.cumsum(bytelen_list))
    bytelen_cumsum_list = list(map(lambda x: int(x), bytelen_cumsum_list))

    all_bytearray = bytearray()
    for temp_bytearray in bytearray_list:
        all_bytearray.extend(temp_bytearray)
    offset_list = [0] + bytelen_cumsum_list # First offset is 0
    offset_list.pop() # Delete the last element
    
    # Resource
    resources = [GLBResource(data=all_bytearray)]
    
    # Buffer
    buffers = [Buffer(byteLength=len(all_bytearray))]

    # BufferView
    bufferViews = [
        BufferView(buffer=0, byteOffset=offset_list[0], byteLength=bytelen_list[0], target=BufferTarget.ARRAY_BUFFER.value),
        BufferView(buffer=0, byteOffset=offset_list[1], byteLength=bytelen_list[1], target=BufferTarget.ARRAY_BUFFER.value),
        BufferView(buffer=0, byteOffset=offset_list[2], byteLength=bytelen_list[2], target=BufferTarget.ARRAY_BUFFER.value),
        BufferView(buffer=0, byteOffset=offset_list[3], byteLength=bytelen_list[3], target=BufferTarget.ELEMENT_ARRAY_BUFFER.value),
        BufferView(buffer=0, byteOffset=offset_list[4], byteLength=bytelen_list[4], target=None),
    ]

    # Accessor
    accessors = [
        Accessor(bufferView=0, componentType=ComponentType.FLOAT.value, count=len(vertices*2), type=AccessorType.VEC3.value, max=maxs, min=mins),
        Accessor(bufferView=1, componentType=ComponentType.FLOAT.value, count=len(normals*2), type=AccessorType.VEC3.value, max=None, min=None),
        Accessor(bufferView=2, componentType=ComponentType.FLOAT.value, count=len(texcoord_0s*2), type=AccessorType.VEC2.value, max=None, min=None),
        Accessor(bufferView=3, componentType=ComponentType.UNSIGNED_SHORT.value, count=len(vertex_indices), type=AccessorType.SCALAR.value, max=None, min=None)
    ]
    
    # Node
    nodes = [
        Node(mesh=0,rotation=None, scale=scale),
    ]
    
    model = GLTFModel(
        asset=asset,
        buffers=buffers,
        bufferViews=bufferViews,
        accessors=accessors,
        images=images,
        samplers=samplers,
        textures=textures,
        materials=materials,
        meshes=meshes,
        nodes=nodes,
        scene=scene,
        scenes=scenes,
    )
    
    gltf = GLTF(model=model, resources=resources)
    
    tmp_filename = uuid.uuid4().hex
    model_path = f'../tmp/{tmp_filename}.glb'

    gltf.export(model_path)
    
    return model_path