src/extracted_objects_model.py

import io
import numpy as np
from PIL import Image as PIL_Image # Renaming to avoid conflict with Image from gltflib
import cv2
import struct
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, PBRMetallicRoughness)

# Create vertex lists for both the front and back surfaces
def make_front_and_back_vertex_list(coordinate_list, img):

    # Front surface vertices
    front_vertex_list = []
    # Back surface vertices
    back_vertex_list = []
    for coordinates in coordinate_list:
        front_vertices = []
        back_vertices = []
        # Note that the Y-axis direction is inverted between the image and GLB, so be careful
        for coordinate in coordinates:
            front_vertices.append((coordinate[0] * 2 / img.size[0] - 1.0, -(coordinate[1] * 2 / img.size[1] - 1.0), 0.2))
            back_vertices.append((coordinate[0] * 2 / img.size[0] - 1.0, -(coordinate[1] * 2 / img.size[1] - 1.0), -0.2))
    
        front_vertex_list.append(front_vertices)
        back_vertex_list.append(back_vertices)

    return front_vertex_list, back_vertex_list
    
# Creation of various information for the mesh
def make_mesh_data(coordinate_list, img):
    front_vertex_list, back_vertex_list = make_front_and_back_vertex_list(coordinate_list, img)

    # Vertex data(POSITION)
    vertices = []
    # List of offset values used when determining vertex indices
    front_offset = 0
    front_offset_list = []
    back_offset_list = []
    for front_vertices, back_vertices in zip(front_vertex_list, back_vertex_list):
        vertices.extend(front_vertices)
        vertices.extend(back_vertices)

        back_offset = front_offset + len(front_vertices)
        front_offset_list.append(front_offset)
        back_offset_list.append(back_offset)
        front_offset += len(front_vertices) + len(back_vertices)        

    # Normal data(NORMAL)
    normals = []
    for front_vertices, back_vertices in zip(front_vertex_list, back_vertex_list):
        normals.extend([( 0.0,  0.0,  1.0)] * len(front_vertices))
        normals.extend([( 0.0,  0.0, -1.0)] * len(back_vertices))

    # Texture coordinates (TEXCOORD_0)
    # The image origin is at the top-left, requiring an inversion of the Y-axis.
    texcoord_0s = [((vertex[0] + 1.0) / 2.0, 1.0 - ((vertex[1] + 1.0) / 2.0) ) for vertex in vertices]

    # Vertex indices
    vertex_indices = []
    for front_vertices, back_vertices, front_offset, back_offset \
        in zip(front_vertex_list, back_vertex_list, front_offset_list, back_offset_list):
        polygon = {
            'vertices': np.array(front_vertices)[:, :2],
            'segments': np.array([( i, (i + 1) % (len(front_vertices)) ) for i in range(len(front_vertices))]) # Define each edge
        }
        triangulate_result = triangle.triangulate(polygon, 'p')
        vertex_indices.extend(list(np.array(triangulate_result['triangles']+front_offset).flatten())) # Front surface
        vertex_indices.extend(list((np.array(triangulate_result['triangles'])+back_offset).flatten())) # Back surface
        vertex_indices.extend(list(np.array([[front_offset + i, 
                                              front_offset + (i + 1) % len(front_vertices), 
                                              back_offset + i] 
                                               for i in range(len(front_vertices))]).flatten())) # Side surface 1
        vertex_indices.extend(list(np.array([[back_offset + i,
                                              back_offset + (i + 1) % len(back_vertices), 
                                              front_offset+ (i + 1) % len(front_vertices)] for i in range(len(front_vertices))]).flatten())) # Side surface 2
            
    return vertices, normals, texcoord_0s, vertex_indices

def create_extracted_objects_model(img_bytearray):

    # Retrieve the image
    img = PIL_Image.open(img_bytearray).convert('RGB')
    img_bytearray = io.BytesIO()
    img.save(img_bytearray, format="JPEG", quality=95)
    img_bytearray = img_bytearray.getvalue()
    img_bytelen = len(img_bytearray)

    # Calculate the scale of the 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)

    # Retrieve vertices of the main part of the image
    base_color = img.getpixel((0, 0))
    mask = PIL_Image.new('RGB', img.size)
    for i in range(img.size[0]):
        for j in range(img.size[1]):
            if base_color == img.getpixel((i, j)):
                mask.putpixel((i, j), (0, 0, 0))
            else:
                mask.putpixel((i, j), (255, 255, 255))

    opening = cv2.morphologyEx(np.array(mask), cv2.MORPH_OPEN, kernel=np.ones((15, 15),np.uint8))
    contours, _ = cv2.findContours(cv2.cvtColor(np.array(opening), cv2.COLOR_RGB2GRAY), cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    coordinate_list = []
    for contour in contours:
        coordinates = []
        for [[x, y]] in contour:
            coordinates.append((x, y))
        coordinate_list.append(coordinates)
    
    # Creation of associated data for the mesh
    vertices, normals, texcoord_0s, vertex_indices = make_mesh_data(coordinate_list, img)
    
    # Vertex data(POSITION)
    vertex_bytearray = bytearray()
    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)
    normal_bytearray = bytearray()
    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()
    for texcoord_0 in texcoord_0s:
        for value in texcoord_0:
            texcoord_0_bytearray.extend(struct.pack('f', value))
    texcoord_0_bytelen = len(texcoord_0_bytearray)

    # Vertex indices
    vertex_index_bytearray = bytearray()
    for value in vertex_indices:
        vertex_index_bytearray.extend(struct.pack('H', value))
    vertex_index_bytelen = len(vertex_index_bytearray)

    # Concatenation of the binary data section
    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 # The first offset is 0
    offset_list.pop() # 末尾を削除

    # GLBResource
    resources = [GLBResource(data=all_bytearray)]

    # Asset
    asset=Asset()

    # 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), type=AccessorType.VEC3.value, max=maxs, min=mins),
        Accessor(bufferView=1, componentType=ComponentType.FLOAT.value, count=len(normals), type=AccessorType.VEC3.value, max=None, min=None),
        Accessor(bufferView=2, componentType=ComponentType.FLOAT.value, count=len(texcoord_0s), 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)
    ]

    # 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)]),
    ]

    # Node
    nodes = [
        Node(mesh=0,rotation=None, scale=scale),
    ]

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

    # GLTFModel
    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