new space

G_checkpoint.pkl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:423997d54f967d4ac949a2fdc2008a25a70e1c005954eaf18d7151728b974023
+size 121792171

app.py

@@ -0,0 +1,122 @@
+import gradio as gr
+import glob
+import utils.utiles as ut
+import numpy as np
+from pywavefront import Wavefront
+from pathlib import Path
+
+# https://huggingface.co/spaces/radames/dpt-depth-estimation-3d-voxels/blob/main/app.py
+
+
+def create_cube_around_point(point, cube_size, idx=0):
+    x, y, z = point
+    half_size = cube_size / 2.0
+
+    vertices = [
+        # Vertices delanteros inferiores
+        (x - half_size, y - half_size, z - half_size),
+        (x + half_size, y - half_size, z - half_size),
+        (x + half_size, y + half_size, z - half_size),
+        (x - half_size, y + half_size, z - half_size),
+        # Vertices traseros inferiores
+        (x - half_size, y - half_size, z + half_size),
+        (x + half_size, y - half_size, z + half_size),
+        (x + half_size, y + half_size, z + half_size),
+        (x - half_size, y + half_size, z + half_size),
+    ]
+
+    faces = [
+        (1, 2, 3, 4),  # Cara delantera
+        (5, 6, 7, 8),  # Cara trasera
+        (1, 5, 8, 4),  # Lado izquierdo
+        (2, 6, 7, 3),  # Lado derecho
+        (1, 2, 6, 5),  # Lado inferior
+        (4, 3, 7, 8),  # Lado superior
+    ]
+
+    return vertices, list(np.array(faces) + (idx*8))
+
+
+def vox_to_obj(voxel_array, name_obj):
+
+    vertices = []
+    faces = []
+    vertices_reales = []
+    faces_reales = []
+
+    idx = 0
+
+    for x in range(voxel_array.shape[0]):
+        for y in range(voxel_array.shape[1]):
+            for z in range(voxel_array.shape[2]):
+                if voxel_array[x, y, z] > 0:
+                    v, f = create_cube_around_point((x, y, z), 0.7, idx)
+                    vertices_reales += v
+                    faces_reales += f
+                    idx += 1
+                    vertices.append((x, y, z))
+                    #print(np.mean(v, 0), (x, y, z))
+
+    # Crear un archivo .obj
+    obj_file = open(name_obj, "w")
+    for vertex in vertices_reales:
+        obj_file.write(f"v {vertex[0]} {vertex[1]} {vertex[2]}\n")
+
+    for face in faces_reales:
+        obj_file.write(f"f {' '.join(map(str, face))}\n")
+
+    obj_file.close()
+
+
+def obj_to_vox(obj_file):
+
+    # Leer el archivo .obj
+    obj_mesh = Wavefront(obj_file)
+
+    # Crear un nuevo arreglo voxel
+    new_voxel_array = np.zeros((64, 64, 64), dtype=np.uint8)
+    for vertex in np.array(obj_mesh.vertices).reshape(-1, 8, 3):
+        x, y, z = np.mean(vertex, 0)
+        new_voxel_array[round(x), round(y), round(z)] = 1
+    return new_voxel_array
+
+
+def process_image(v):
+    vox_frag = obj_to_vox(v)
+    image_path = Path(v)
+
+    PATH = './G_checkpoint.pkl'
+    G_encode_decode = ut.load_generator(PATH)
+    fake = ut.generate(G_encode_decode, vox_frag)
+    _, fake_posprocess = ut.posprocessing(fake, vox_frag)
+    gltf_path = f'./{image_path.stem}_result.obj'
+    vox_to_obj(fake_posprocess, gltf_path)
+    print(gltf_path)
+    return gltf_path
+
+
+title = "IberianVoxel: Automatic Completion of Iberian Ceramics for Cultural Heritage Studies"
+description = "Accurate completion of archaeological artifacts is a critical aspect in several archaeological studies, including documentation of variations in style, inference of chronological and ethnic groups, and trading routes trends, among many others. However, most available pottery is fragmented, leading to missing textural and morphological cues. Currently, the reassembly and completion of fragmented ceramics is a daunting and time-consuming task, done almost exclusively by hand, which requires the physical manipulation of the fragments. To overcome the challenges of manual reconstruction, reduce the materials' exposure and deterioration, and improve the quality of reconstructed samples, we present IberianVoxel, a novel 3D Autoencoder Generative Adversarial Network (3D AE-GAN) framework tested on an extensive database with complete and fragmented references. We generated a collection of  3D voxelized samples and their fragmented references from Iberian wheel-made pottery profiles. The fragments generated are stratified into different size groups and across multiple pottery classes. Lastly, we provide quantitative and qualitative assessments to measure the quality of the reconstructed voxelized samples by our proposed method and archaeologists' evaluation."
+examples = [img for img in glob.glob("examples/*.obj")]
+
+# background-color: black;
+#css = svelte-wn75i6
+
+
+iface = gr.Interface(fn=process_image,
+                     inputs=[
+                         gr.Model3D(label="Fragment Input",
+                                    elem_id="model-in", clear_color=['black']),
+
+                     ],
+                     outputs=[
+                         gr.Model3D(label="3d mesh reconstruction",
+                                    clear_color=['black'])
+
+                     ],
+                     title=title,
+                     description=description,
+                     examples=examples,
+
+                     )
+iface.launch(share=False)

pyvox/__init__.py

@@ -0,0 +1 @@
+__version__ = '0.1'

pyvox/defaultpalette.py

@@ -0,0 +1,19 @@
+
+default_palette = [
+    0x00000000, 0xffffffff, 0xffccffff, 0xff99ffff, 0xff66ffff, 0xff33ffff, 0xff00ffff, 0xffffccff, 0xffccccff, 0xff99ccff, 0xff66ccff, 0xff33ccff, 0xff00ccff, 0xffff99ff, 0xffcc99ff, 0xff9999ff,
+    0xff6699ff, 0xff3399ff, 0xff0099ff, 0xffff66ff, 0xffcc66ff, 0xff9966ff, 0xff6666ff, 0xff3366ff, 0xff0066ff, 0xffff33ff, 0xffcc33ff, 0xff9933ff, 0xff6633ff, 0xff3333ff, 0xff0033ff, 0xffff00ff,
+    0xffcc00ff, 0xff9900ff, 0xff6600ff, 0xff3300ff, 0xff0000ff, 0xffffffcc, 0xffccffcc, 0xff99ffcc, 0xff66ffcc, 0xff33ffcc, 0xff00ffcc, 0xffffcccc, 0xffcccccc, 0xff99cccc, 0xff66cccc, 0xff33cccc,
+    0xff00cccc, 0xffff99cc, 0xffcc99cc, 0xff9999cc, 0xff6699cc, 0xff3399cc, 0xff0099cc, 0xffff66cc, 0xffcc66cc, 0xff9966cc, 0xff6666cc, 0xff3366cc, 0xff0066cc, 0xffff33cc, 0xffcc33cc, 0xff9933cc,
+    0xff6633cc, 0xff3333cc, 0xff0033cc, 0xffff00cc, 0xffcc00cc, 0xff9900cc, 0xff6600cc, 0xff3300cc, 0xff0000cc, 0xffffff99, 0xffccff99, 0xff99ff99, 0xff66ff99, 0xff33ff99, 0xff00ff99, 0xffffcc99,
+    0xffcccc99, 0xff99cc99, 0xff66cc99, 0xff33cc99, 0xff00cc99, 0xffff9999, 0xffcc9999, 0xff999999, 0xff669999, 0xff339999, 0xff009999, 0xffff6699, 0xffcc6699, 0xff996699, 0xff666699, 0xff336699,
+    0xff006699, 0xffff3399, 0xffcc3399, 0xff993399, 0xff663399, 0xff333399, 0xff003399, 0xffff0099, 0xffcc0099, 0xff990099, 0xff660099, 0xff330099, 0xff000099, 0xffffff66, 0xffccff66, 0xff99ff66,
+    0xff66ff66, 0xff33ff66, 0xff00ff66, 0xffffcc66, 0xffcccc66, 0xff99cc66, 0xff66cc66, 0xff33cc66, 0xff00cc66, 0xffff9966, 0xffcc9966, 0xff999966, 0xff669966, 0xff339966, 0xff009966, 0xffff6666,
+    0xffcc6666, 0xff996666, 0xff666666, 0xff336666, 0xff006666, 0xffff3366, 0xffcc3366, 0xff993366, 0xff663366, 0xff333366, 0xff003366, 0xffff0066, 0xffcc0066, 0xff990066, 0xff660066, 0xff330066,
+    0xff000066, 0xffffff33, 0xffccff33, 0xff99ff33, 0xff66ff33, 0xff33ff33, 0xff00ff33, 0xffffcc33, 0xffcccc33, 0xff99cc33, 0xff66cc33, 0xff33cc33, 0xff00cc33, 0xffff9933, 0xffcc9933, 0xff999933,
+    0xff669933, 0xff339933, 0xff009933, 0xffff6633, 0xffcc6633, 0xff996633, 0xff666633, 0xff336633, 0xff006633, 0xffff3333, 0xffcc3333, 0xff993333, 0xff663333, 0xff333333, 0xff003333, 0xffff0033,
+    0xffcc0033, 0xff990033, 0xff660033, 0xff330033, 0xff000033, 0xffffff00, 0xffccff00, 0xff99ff00, 0xff66ff00, 0xff33ff00, 0xff00ff00, 0xffffcc00, 0xffcccc00, 0xff99cc00, 0xff66cc00, 0xff33cc00,
+    0xff00cc00, 0xffff9900, 0xffcc9900, 0xff999900, 0xff669900, 0xff339900, 0xff009900, 0xffff6600, 0xffcc6600, 0xff996600, 0xff666600, 0xff336600, 0xff006600, 0xffff3300, 0xffcc3300, 0xff993300,
+    0xff663300, 0xff333300, 0xff003300, 0xffff0000, 0xffcc0000, 0xff990000, 0xff660000, 0xff330000, 0xff0000ee, 0xff0000dd, 0xff0000bb, 0xff0000aa, 0xff000088, 0xff000077, 0xff000055, 0xff000044,
+    0xff000022, 0xff000011, 0xff00ee00, 0xff00dd00, 0xff00bb00, 0xff00aa00, 0xff008800, 0xff007700, 0xff005500, 0xff004400, 0xff002200, 0xff001100, 0xffee0000, 0xffdd0000, 0xffbb0000, 0xffaa0000,
+    0xff880000, 0xff770000, 0xff550000, 0xff440000, 0xff220000, 0xff110000, 0xffeeeeee, 0xffdddddd, 0xffbbbbbb, 0xffaaaaaa, 0xff888888, 0xff777777, 0xff555555, 0xff444444, 0xff222222, 0xff111111
+]

pyvox/models.py

@@ -0,0 +1,87 @@
+from collections import namedtuple
+
+from .defaultpalette import default_palette
+from .utils import chunks
+
+Size = namedtuple('Size', 'x y z')
+Color = namedtuple('Color', 'r g b a')
+Voxel = namedtuple('Voxel', 'x y z c')
+Model = namedtuple('Model', 'size voxels')
+Material = namedtuple('Material', 'id type weight props')
+
+def get_default_palette():
+    return [ Color( *tuple(i.to_bytes(4,'little')) ) for i in default_palette ]
+
+
+class Vox(object):
+
+    def __init__(self, models, palette=None, materials=None):
+        self.models = models
+        self.default_palette = not palette
+        self._palette = palette or get_default_palette()
+        self.materials = materials or []
+
+    @property
+    def palette(self):
+        return self._palette
+
+    @palette.setter
+    def palette(self, val):
+        self._palette = val
+        self.default_palette = False
+
+    def to_dense_rgba(self, model_idx=0):
+
+        import numpy as np
+        m = self.models[model_idx]
+        res = np.zeros(( m.size.y, m.size.z, m.size.x, 4 ), dtype='B')
+
+        for v in m.voxels:
+            res[v.y, m.size.z-v.z-1, v.x] = self.palette[v.c]
+
+        return res
+
+    def to_dense(self, model_idx=0):
+
+        import numpy as np
+        m = self.models[model_idx]
+        res = np.zeros(( m.size.y, m.size.z, m.size.x ), dtype='B')
+
+        for v in m.voxels:
+            res[v.y, m.size.z-v.z-1, v.x] = v.c
+
+        return res
+
+    def __str__(self):
+        return 'Vox(%s)'%(self.models)
+
+    @staticmethod
+    def from_dense(a, black=[0,0,0]):
+
+        palette = None
+
+        if len(a.shape) == 4:
+            from PIL import Image
+            import numpy as np
+
+            mask = np.all(a == np.array([[black]]), axis=3)
+
+            x,y,z,_ = a.shape
+
+            # color index 0 is reserved for empty, so we get 255 colors
+            img = Image.fromarray(a.reshape(x,y*z,3)).quantize(255)
+            palette = img.getpalette()
+            palette = [ Color(0,0,0,0) ] + [ Color(*c, 255) for c in chunks(palette, 3) ]
+            a = np.asarray(img, dtype='B').reshape(x,y,z).copy() + 1
+            a[mask] = 0
+
+
+        if len(a.shape) != 3: raise Exception("I expect a 4 or 3 dimensional matrix")
+
+        y,z,x = a.shape
+
+        nz = a.nonzero()
+
+        voxels = [ Voxel( nz[2][i], nz[0][i], z-nz[1][i]-1, a[nz[0][i], nz[1][i], nz[2][i]] ) for i in range(nz[0].shape[0]) ]
+
+        return Vox([ Model(Size(x,y,z), voxels)], palette)

pyvox/parser.py

@@ -0,0 +1,136 @@
+from struct import unpack_from as unpack, calcsize
+import logging
+
+from .models import Vox, Size, Voxel, Color, Model, Material
+
+log = logging.getLogger(__name__)
+
+class ParsingException(Exception): pass
+
+def bit(val, offset):
+    mask = 1 << offset
+    return(val & mask)
+
+class Chunk(object):
+    def __init__(self, id, content=None, chunks=None):
+        self.id = id
+        self.content = content or b''
+        self.chunks = chunks or []
+
+        if id == b'MAIN':
+            if len(self.content): raise ParsingException('Non-empty content for main chunk')
+        elif id == b'PACK':
+            self.models = unpack('i', content)[0]
+        elif id == b'SIZE':
+            self.size = Size(*unpack('iii', content))
+        elif id == b'XYZI':
+            n = unpack('i', content)[0]
+            log.debug('xyzi block with %d voxels (len %d)', n, len(content))
+            self.voxels = []
+            self.voxels = [ Voxel(*unpack('BBBB', content, 4+4*i)) for i in range(n) ]
+        elif id == b'RGBA':
+            self.palette = [ Color(*unpack('BBBB', content, 4*i)) for i in range(255) ]
+            # Docs say:  color [0-254] are mapped to palette index [1-255]
+            # hmm
+            # self.palette = [ Color(0,0,0,0) ] + [ Color(*unpack('BBBB', content, 4*i)) for i in range(255) ]
+        elif id == b'MATT':
+            _id, _type, weight, flags = unpack('iifi', content)
+            props = {}
+            offset = 16
+            for b,field in [ (0, 'plastic'),
+                             (1, 'roughness'),
+                             (2, 'specular'),
+                             (3, 'IOR'),
+                             (4, 'attenuation'),
+                             (5, 'power'),
+                             (6, 'glow'),
+                             (7, 'isTotalPower') ]:
+                if bit(flags, b) and b<7: # no value for 7 / isTotalPower
+                    props[field] = unpack('f', content, offset)
+                    offset += 4
+
+            self.material = Material(_id, _type, weight, props)
+
+        else:
+            raise ParsingException('Unknown chunk type: %s'%self.id)
+
+class VoxParser(object):
+
+    def __init__(self, filename):
+        with open(filename, 'rb') as f:
+            self.content = f.read()
+
+        self.offset = 0
+
+    def unpack(self, fmt):
+
+        r = unpack(fmt, self.content, self.offset)
+        self.offset += calcsize(fmt)
+
+        return r
+
+    def _parseChunk(self):
+
+        _id, N, M = self.unpack('4sii')
+
+        log.debug("Found chunk id %s / len %s / children %s", _id, N, M)
+
+        content = self.unpack('%ds'%N)[0]
+
+        start = self.offset
+        chunks = [ ]
+        while self.offset<start+M:
+            chunks.append(self._parseChunk())
+
+        return Chunk(_id, content, chunks)
+
+    def parse(self):
+
+        header, version = self.unpack('4si')
+        if header != b'VOX ': raise ParsingException("This doesn't look like a vox file to me")
+
+        if version != 150: raise ParsingException("Unknown vox version: %s expected 150"%version)
+
+        main = self._parseChunk()
+
+        if main.id != b'MAIN': raise ParsingException("Missing MAIN Chunk")
+
+        chunks = list(reversed(main.chunks))
+        if chunks[-1].id == b'PACK':
+            models = chunks.pop().models
+        else:
+            models = 1
+
+        log.debug("file has %d models", models)
+
+        models = [ self._parseModel(chunks.pop(), chunks.pop()) for _ in range(models) ]
+
+        if chunks and chunks[0].id == b'RGBA':
+            palette = chunks.pop().palette
+        else:
+            palette = None
+
+        materials = [ c.material for c in chunks ]
+
+        return Vox(models, palette, materials)
+
+
+
+    def _parseModel(self, size, xyzi):
+        if size.id != b'SIZE': raise ParsingException('Expected SIZE chunk, got %s', size.id)
+        if xyzi.id != b'XYZI': raise ParsingException('Expected XYZI chunk, got %s', xyzi.id)
+
+        return Model(size.size, xyzi.voxels)
+
+
+
+
+if __name__ == '__main__':
+
+    import sys
+    import coloredlogs
+
+    coloredlogs.install(level=logging.DEBUG)
+
+
+    VoxParser(sys.argv[1]).parse()

pyvox/utils.py

@@ -0,0 +1,4 @@
+def chunks(l, n):
+    """Yield successive n-sized chunks from l."""
+    for i in range(0, len(l), n):
+        yield l[i:i + n]

pyvox/writer.py

@@ -0,0 +1,61 @@
+from struct import pack
+
+
+class VoxWriter(object):
+
+    def __init__(self, filename, vox):
+        self.filename = filename
+        self.vox = vox
+
+    def _chunk(self, id, content, chunks=[]):
+
+        res = b''
+        for c in chunks:
+            res += self._chunk(*c)
+
+        return pack('4sii', id, len(content), len(res)) + content + res
+
+    def _matflags(self, props):
+        flags = 0
+        res = b''
+        for b,field in [ (0, 'plastic'),
+                         (1, 'roughness'),
+                         (2, 'specular'),
+                         (3, 'IOR'),
+                         (4, 'attenuation'),
+                         (5, 'power'),
+                         (6, 'glow'),
+                         (7, 'isTotalPower') ]:
+            if field in props:
+                flags |= 1<<b
+                res += pack('f', props[field])
+
+        return pack('i', flags)+res
+
+
+
+    def write(self):
+
+        res = pack('4si', b'VOX ', 150)
+
+        chunks = []
+
+        if len(self.vox.models):
+            chunks.append((b'PACK', pack('i', len(self.vox.models))))
+
+        for m in self.vox.models:
+            chunks.append((b'SIZE', pack('iii', *m.size)))
+            chunks.append((b'XYZI', pack('i', len(m.voxels)) + b''.join(pack('BBBB', *v) for v in m.voxels)))
+
+        if not self.vox.default_palette:
+            chunks.append((b'RGBA', b''.join(pack('BBBB', *c) for c in self.vox.palette)))
+
+        for m in self.vox.materials:
+            chunks.append((b'MATT', pack('iif', m.id, m.type, m.weight) + self._matflags(m.props)))
+
+        # TODO materials
+
+        res += self._chunk(b'MAIN', b'', chunks)
+
+        with open(self.filename, 'wb') as f:
+            f.write(res)

utils/FragmentDataset.py

@@ -0,0 +1,83 @@
+import glob
+from torch.utils.data import Dataset
+import numpy as np
+import pyvox.parser
+
+
+class FragmentDataset(Dataset):
+    def __init__(self, vox_path, vox_type, dim_size=64, transform=None):
+        self.vox_type = vox_type
+        self.vox_path = vox_path
+        self.transform = transform
+        self.dim_size = dim_size
+        self.vox_files = sorted(
+            glob.glob('{}/{}/*/*.vox'.format(self.vox_path, self.vox_type)))
+
+    def __len__(self):
+        return len(self.vox_files)
+
+    def __read_vox__(self, path):
+        vox = pyvox.parser.VoxParser(path).parse()
+        a = vox.to_dense()
+        caja = np.zeros((64, 64, 64))
+        caja[0:a.shape[0], 0:a.shape[1], 0:a.shape[2]] = a
+        return caja
+
+    def __select_fragment__(self, v):
+        frag_id = np.unique(v)[1:]
+        #select_frag = np.random.choice(frag_id, np.random.choice(np.arange(1, len(frag_id)), 1)[0], replace=False)
+        select_frag = np.random.choice(frag_id, 1, replace=False)
+        for f in frag_id:
+            if not(f in select_frag):
+                v[v == f] = 0
+            else:
+                v[v == f] = 1
+        return v, select_frag
+
+    def __non_select_fragment__(self, v, select_frag):
+        frag_id = np.unique(v)[1:]
+        for f in frag_id:
+            if not(f in select_frag):
+                v[v == f] = 1
+            else:
+                v[v == f] = 0
+        return v
+
+    def __select_fragment_specific__(self, v, select_frag):
+        frag_id = np.unique(v)[1:]
+        for f in frag_id:
+            if not(f in select_frag):
+                v[v == f] = 0
+            else:
+                v[v == f] = 1
+        return v, select_frag
+
+    def __getitem__(self, idx):
+        img_path = self.vox_files[idx]
+        vox = self.__read_vox__(img_path)
+        label = img_path.replace(self.vox_path, '').split('/')[2]
+        frag, select_frag = self.__select_fragment__(vox.copy())
+
+        if self.transform:
+            vox = self.transform(vox)
+            frag = self.transform(frag)
+
+        return frag, vox,  # select_frag, int(label)-1#, img_path
+
+    def __getitem_specific_frag__(self, idx, select_frag):
+        img_path = self.vox_files[idx]
+        vox = self.__read_vox__(img_path)
+        label = img_path.replace(self.vox_path, '').split('/')[2]
+        frag, select_frag = self.__select_fragment_specific__(
+            vox.copy(), select_frag)
+
+        if self.transform:
+            vox = self.transform(vox)
+            frag = self.transform(frag)
+
+        return frag, vox,  # select_frag, int(label)-1, img_path
+
+    def __getfractures__(self, idx):
+        img_path = self.vox_files[idx]
+        vox = self.__read_vox__(img_path)
+        return np.unique(vox)  # select_frag, int(label)-1, img_path

utils/network_vox.py

@@ -0,0 +1,155 @@
+import torch
+
+class _D(torch.nn.Module):
+    def __init__(self, cube_len=64):
+        super(_D, self).__init__()
+        self.leak = 0.2
+        self.cube_len = cube_len
+
+        padd = (0,0,0)
+        if self.cube_len == 32:
+            padd = (1,1,1)
+
+        self.layer1 = torch.nn.Sequential(
+
+            torch.nn.Conv3d(1, self.cube_len, kernel_size=4, stride=2, bias=False, padding=(1, 1, 1)),
+            torch.nn.BatchNorm3d(self.cube_len),
+            torch.nn.LeakyReLU(self.leak)
+        )
+        self.layer2 = torch.nn.Sequential(
+            torch.nn.Conv3d(self.cube_len, self.cube_len*2, kernel_size=4, stride=2, bias=False, padding=(1, 1, 1)),
+            torch.nn.BatchNorm3d(self.cube_len*2),
+            torch.nn.LeakyReLU(self.leak)
+        )
+        self.layer3 = torch.nn.Sequential(
+            torch.nn.Conv3d(self.cube_len*2, self.cube_len*4, kernel_size=4, stride=2, bias=False, padding=(1, 1, 1)),
+            torch.nn.BatchNorm3d(self.cube_len*4),
+            torch.nn.LeakyReLU(self.leak)
+        )
+        self.layer4 = torch.nn.Sequential(
+            torch.nn.Conv3d(self.cube_len*4, self.cube_len*8, kernel_size=4, stride=2, bias=False, padding=(1, 1, 1)),
+            torch.nn.BatchNorm3d(self.cube_len*8),
+            torch.nn.LeakyReLU(self.leak)
+        )
+        self.layer5 = torch.nn.Sequential(
+            torch.nn.Conv3d(self.cube_len*8, 1, kernel_size=4, stride=2, bias=False, padding=padd),
+        )
+        self.layer6 = torch.nn.Sequential(
+            torch.nn.Sigmoid(),
+            #torch.nn.Linear(64, 1),
+            #torch.nn.LogSoftmax(dim=1)
+        )
+
+    def forward(self, x):
+        out = x.view((-1, 1, self.cube_len, self.cube_len, self.cube_len))
+        out = self.layer1(out)
+        out = self.layer2(out)
+        out = self.layer3(out)
+        out = self.layer4(out)
+        #print(out.shape)
+        out = self.layer5(out).view(-1, 1)
+        #print(out.shape)
+        out = self.layer6(out)
+        #print(out.shape)
+        return out
+
+class _G_encode_decode(torch.nn.Module):
+    def __init__(self, cube_len=64, z_latent_space=64, z_intern_space=64):
+        super(_G_encode_decode, self).__init__()
+        self.leak = 0.01
+        self.cube_len = cube_len
+        self.z_latent_space = z_latent_space
+        self.z_intern_space = z_intern_space
+
+        padd = (0,0,0)
+        if self.cube_len == 32:
+            padd = (1,1,1)
+
+        self.layer1 = torch.nn.Sequential(
+            torch.nn.Conv3d(1, self.cube_len, kernel_size=4, stride=2, bias=False, padding=(1, 1, 1)),
+            torch.nn.BatchNorm3d(self.cube_len),
+            torch.nn.LeakyReLU(self.leak)
+        )
+        self.layer2 = torch.nn.Sequential(
+            torch.nn.Conv3d(self.cube_len, self.cube_len*2, kernel_size=4, stride=2, bias=False, padding=(1, 1, 1)),
+            torch.nn.BatchNorm3d(self.cube_len*2),
+            torch.nn.LeakyReLU(self.leak)
+        )
+        self.layer3 = torch.nn.Sequential(
+            torch.nn.Conv3d(self.cube_len*2, self.cube_len*4, kernel_size=4, stride=2, bias=False, padding=(1, 1, 1)),
+            torch.nn.BatchNorm3d(self.cube_len*4),
+            torch.nn.LeakyReLU(self.leak)
+        )
+        self.layer4 = torch.nn.Sequential(
+            torch.nn.Conv3d(self.cube_len*4, self.cube_len*8, kernel_size=4, stride=2, bias=False, padding=(1, 1, 1)),
+            torch.nn.BatchNorm3d(self.cube_len*8),
+            torch.nn.LeakyReLU(self.leak)
+        )
+        self.layer5 = torch.nn.Sequential(
+            torch.nn.Conv3d(self.cube_len*8, self.cube_len*2, kernel_size=4, stride=2, bias=False, padding=padd),
+            #torch.nn.BatchNorm3d(self.cube_len*8),
+
+        )
+
+        self.layer6 = torch.nn.Sequential(
+            torch.nn.ConvTranspose3d(self.z_latent_space, self.cube_len*8, kernel_size=4, stride=2, bias=False, padding=padd),
+            torch.nn.BatchNorm3d(self.cube_len*8),
+            torch.nn.ReLU()
+        )
+        self.layer7 = torch.nn.Sequential(
+            torch.nn.ConvTranspose3d(self.cube_len*8, self.cube_len*4, kernel_size=4, stride=2, bias=False, padding=(1, 1, 1)),
+            torch.nn.BatchNorm3d(self.cube_len*4),
+            torch.nn.ReLU()
+        )
+        self.layer8 = torch.nn.Sequential(
+            torch.nn.ConvTranspose3d(self.cube_len*4, self.cube_len*2, kernel_size=4, stride=2, bias=False, padding=(1, 1, 1)),
+            torch.nn.BatchNorm3d(self.cube_len*2),
+            torch.nn.ReLU()
+        )
+        self.layer9 = torch.nn.Sequential(
+            torch.nn.ConvTranspose3d(self.cube_len*2, self.cube_len, kernel_size=4, stride=2, bias=False, padding=(1, 1, 1)),
+            torch.nn.BatchNorm3d(self.cube_len),
+            torch.nn.ReLU()
+        )
+        self.layer10 = torch.nn.Sequential(
+            torch.nn.ConvTranspose3d(self.cube_len, 1, kernel_size=4, stride=2, bias=False, padding=(1, 1, 1)),
+
+
+       )
+        self.linear_layer = torch.nn.Sequential(
+            torch.nn.Linear(self.z_intern_space, self.z_latent_space),
+            torch.nn.BatchNorm1d(self.z_latent_space),
+            torch.nn.ReLU()
+        )
+        self.normalized_layer = torch.nn.Sequential(
+            torch.nn.BatchNorm2d(self.z_latent_space),
+        )
+
+    def normalized_vector(self, z0, z1):
+        x_encode = torch.concat((z0, z1), 1)
+        return self.linear_layer(x_encode)
+
+    def forward(self, x):
+        x_encode = forward_encode(x)
+        x_decode = forward_decode(x)
+
+        return x_decode
+
+    def forward_encode(self, x):
+        out_x = x.view((-1, 1, self.cube_len, self.cube_len, self.cube_len))
+        out = self.layer1(out_x)
+        out = self.layer2(out)
+        out = self.layer3(out)
+        out = self.layer4(out)
+        out = self.layer5(out)
+
+        return out.view(out_x.shape[0], -1)
+
+    def forward_decode(self, x):
+        out = x.view(-1, self.z_latent_space, 1, 1, 1)
+        out = self.layer6(out)
+        out = self.layer7(out)
+        out = self.layer8(out)
+        out = self.layer9(out)
+        out = self.layer10(out)
+        return out

utils/utiles.py

@@ -0,0 +1,141 @@
+import plotly.graph_objects as go
+import numpy as np
+import pyvox.parser
+import torch
+import utils.network_vox as nv
+from scipy import ndimage as ndi
+
+available_device = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+def __read_vox_frag__(path, fragment_idx):
+    vox_pottery = __read_vox__(path)
+    try:
+        assert(fragment_idx in np.unique(vox_pottery))
+        vox_frag = vox_pottery.copy()
+        vox_frag[vox_pottery != fragment_idx] = 0
+        vox_frag[vox_pottery == fragment_idx] = 1
+        return vox_frag
+    except AssertionError:
+        print('fragment_idx not found. Possible fragment_idx {}'.format(
+            np.unique(vox_pottery)[1:]))
+
+
+def __read_vox__(path):
+    vox = pyvox.parser.VoxParser(path).parse()
+    a = vox.to_dense()
+    caja = np.zeros((64, 64, 64))
+    caja[0:a.shape[0], 0:a.shape[1], 0:a.shape[2]] = a
+    return caja
+
+
+def plot(voxel_matrix):
+    voxels = np.array(np.where(voxel_matrix)).T
+    x, y, z = voxels[:, 0], voxels[:, 1], voxels[:, 2]
+    fig = go.Figure(data=go.Scatter3d(x=x, y=y, z=z, mode='markers', marker=dict(size=5, symbol='square', color='#ceabb2', line=dict(width=2,
+                                                                                                                                     color='DarkSlateGrey',))))
+    fig.update_layout()
+
+    fig.show()
+
+
+def plot_frag(vox_pottery):
+    stts = []
+    colors = ['#ceabb2', '#d05d86', '#7e1b2f', '#c1375b', '#cdc1c3',
+              '#ceabb2', '#d05d86', '#7e1b2f', '#c1375b', '#cdc1c3']
+    for i, frag in enumerate(np.unique(vox_pottery)[1:][::-1]):
+        vox_frag = vox_pottery.copy()
+        vox_frag[vox_pottery != frag] = 0
+        voxels = np.array(np.where(vox_frag)).T
+        x, y, z = voxels[:, 0], voxels[:, 1], voxels[:, 2]
+        # ut.plot(vox_frag)
+        scatter = go.Scatter3d(x=x, y=y, z=z,
+                               mode='markers',
+                               name='Fragment {} ({})'.format(i+1, int(frag)),
+                               marker=dict(size=5, symbol='square', color=colors[i],
+                                           line=dict(width=2, color='DarkSlateGrey',)))
+        stts.append(scatter)
+
+    fig = go.Figure(data=stts)
+    fig.update_layout()
+
+    fig.show()
+
+
+def plot_join(vox_1, vox_2):
+    stts = []
+    colors = ['#ceabb2', '#d05d86', '#7e1b2f', '#c1375b', '#cdc1c3',
+              '#ceabb2', '#d05d86', '#7e1b2f', '#c1375b', '#cdc1c3']
+    voxels = np.array(np.where(vox_1)).T
+    x, y, z = voxels[:, 0], voxels[:, 1], voxels[:, 2]
+    # ut.plot(vox_frag)
+    scatter = go.Scatter3d(x=x, y=y, z=z,
+                           mode='markers',
+                           name='Fragment 1',
+                           marker=dict(size=5, symbol='square', color=colors[0],
+                                       line=dict(width=2, color='DarkSlateGrey',)))
+    stts.append(scatter)
+
+    voxels = np.array(np.where(vox_2)).T
+    x, y, z = voxels[:, 0], voxels[:, 1], voxels[:, 2]
+    # ut.plot(vox_frag)
+    scatter = go.Scatter3d(x=x, y=y, z=z,
+                           mode='markers',
+                           name='Fragment 2',
+                           marker=dict(size=5, symbol='square', color=colors[2],
+                                       line=dict(width=2, color='DarkSlateGrey',)))
+    stts.append(scatter)
+
+    fig = go.Figure(data=stts)
+    fig.update_layout()
+
+    fig.show()
+
+
+def posprocessing(fake, mesh_frag):
+    a_p = (mesh_frag > 0.5)
+    a_fake = (fake[0] > np.mean(fake[0]))
+    #a_fake = (fake[0] > 0.1)
+    a_fake = np.array(a_fake, dtype=np.int32).reshape(1, -1)
+
+    diamond = ndi.generate_binary_structure(rank=3, connectivity=1)
+    a_fake = ndi.binary_erosion(a_fake.reshape(
+        64, 64, 64), diamond, iterations=1)
+    _a_p = ndi.binary_erosion(a_p.reshape(64, 64, 64), diamond, iterations=1)
+
+    a_fake = ndi.binary_dilation(
+        a_fake.reshape(64, 64, 64), diamond, iterations=1)
+
+    a_p = ndi.binary_dilation(a_p.reshape(64, 64, 64), diamond, iterations=1)
+    a_fake = a_fake + _a_p
+    #a_fake = (a_fake > 0.5)
+    # make a little 3D diamond:
+    diamond = ndi.generate_binary_structure(rank=3, connectivity=1)
+    dilated = ndi.binary_erosion(
+        a_fake.reshape(64, 64, 64), diamond, iterations=1)
+    dilated = ndi.binary_dilation(
+        a_fake.reshape(64, 64, 64), diamond, iterations=1)
+
+    return a_fake, dilated
+
+
+def load_generator(path_checkpoint):
+
+    G_encode_decode = nv._G_encode_decode(
+        cube_len=64, z_latent_space=128, z_intern_space=136).to(available_device)
+    checkpoint = torch.load(path_checkpoint, map_location=available_device)
+    G_encode_decode.load_state_dict(checkpoint)
+    G_encode_decode = G_encode_decode.eval()
+
+    return G_encode_decode
+
+
+def generate(model, vox_frag):
+    mesh_frag = torch.Tensor(vox_frag).unsqueeze(
+        0).float().to(available_device)
+    output_g_encode = model.forward_encode(mesh_frag)
+    fake = model.forward_decode(output_g_encode)
+    fake = fake + (mesh_frag.unsqueeze(1))
+    fake = fake.detach().cpu().numpy()
+    mesh_frag = mesh_frag.detach().cpu().numpy()
+    return fake, mesh_frag