atlas maker + rgb packer

atlas.py

import os 
import argparse
import cv2
import numpy as np
from tqdm import tqdm 
from math import ceil, sqrt

def atlasFromFolder( name, src, dst, tile, packChannels, quality ):

    if src[:-1] != "/":
        src += "/"

    if dst[:-1] != "/":
        dst += "/"

    os.makedirs(dst, exist_ok=True)
    files = os.listdir(src)
    count = 0
    for f in files:
        n,e = os.path.splitext( f )
        if n == ".DS_Store" or e == ".DS_Store": #pesky macs
            continue
        count += 1

    channels = 1
    if packChannels:
       channels = 3
    
    print( "files", count )
    perBoard = ceil(count / channels)
    size = tile * ceil( sqrt(perBoard ) )
    columns = int( size / tile )
    
    print( "image size", size, ", per board", perBoard, "tile", tile ) 
    print( "grid width", columns )
    print( len( files ), columns ** 2 * channels )
    print( '----------' )

    # creates separate frames for the R,G,B channels
    frames = []
    for i in range(channels):
        if channels == 1:
            img = np.zeros( (size, size, 3 ) )
        else:
            img = np.zeros( (size, size, 1 ) )
        frames.append(img)
    frameId = len( frames )- 1

    # meta    
    widths = []
    index = 0
    iterator = 0
    for f in tqdm( files ):
        n,e = os.path.splitext( f )
        if n == ".DS_Store" or e == ".DS_Store": #pesky macs...
            continue

        # texture overflow: too many images, not enough space 
        if iterator >= columns * columns * channels:
            print( 'texture overflow')
            break
        iterator += 1

        # open and resize image
        tmp = cv2.imread( src + f, cv2.IMREAD_UNCHANGED )
        
        # store the width (to rescale instances)
        iw = tmp.shape[1] / tmp.shape[0] 
        widths.append(str( iw ))

        # and resize image 
        tmp = cv2.resize(tmp, (tile,tile) )

        # convert to greyscale if packing into R,G,B channels
        if channels == 3:
            if len( tmp.shape ) != 2:
                tmp = cv2.cvtColor(tmp, cv2.COLOR_BGR2GRAY ) 
            tmp = tmp.reshape((tile,tile, 1))
            
        # change frame on overflow
        if index >= columns * columns:
            print( index, "change frame")
            frameId -= 1
            index = 0

        # find where to paste image
        x = ( index % columns ) * tile
        y = int(index / columns ) * tile

        # # pick frame and paste
        frame = frames[frameId]
        frame[ y:(y+tile), x:(x+tile), : ] = tmp

        index += 1 

    # save metadata 
    # pack frames and save
    if channels == 3:
        img = cv2.merge(frames, img )

    # save the atlas
    cv2.imwrite( dst + name + ".jpg", img, [int(cv2.IMWRITE_JPEG_QUALITY), quality] )

    # save images widths
    f = open( dst + name + ".txt", "w" )
    f.writelines(','.join( widths ))
    f.close()

parser = argparse.ArgumentParser(description='build a texture atlas from a folder of images.')
parser.add_argument('-name', '-n',  type=str, default="atlas", help='atlas name')
parser.add_argument('-source', '-src', type=str, required=True, help='source folder')
parser.add_argument('-destination', '-dst',  type=str, default="atlas/", help='destination folder')
parser.add_argument('-size', '-s',  type=int, default=64, help='tile size')
parser.add_argument('-pack', '-p',  type=bool, default=False, help='pack rgb channels')
parser.add_argument('-quality', '-q',  type=int, default=80, help='JPEG quality')

if __name__ == "__main__":

    args = parser.parse_args()
    name = args.name
    srcFolder = args.source
    dstFolder = args.destination
    size = args.size
    pack = args.pack
    quality = args.quality

    atlasFromFolder( name, srcFolder, dstFolder, size, pack, quality )

call.md

script's arguments

-name, -n: string
    the atlas' name default "atlas" 
    
-source, -src: string
    the source folder *Required*
    
-destination, -dst: string
    the destination folder default "atlas" 
    
-size, -s: int
    the tile size default 64

-pack, -p : boolean
    whether or not to pack rgb channels default: false

-quality, -q: int
    the JPEG quality [1-100] default 80

sample usage to create a rgb atlas from the images of the source folder

python atlas.py -src source -dst atlas -n color -s 64 -q 60

to pack geryscale values into the RGB channels

python atlas.py -src source -dst atlas -n packed -s 64 -p true -q 90