rj-action-library/Runtime/Procedural/Baking/TextureMerger.cs

using System.Collections;
using System.Collections.Generic;
using Godot;
using System;



namespace Rokojori
{
  [Tool]
  [GlobalClass]
  public partial class TextureMerger:Node
  {    
    public enum SourceMode
    {
      MultiBaker,
      Viewports,
      Textures
    }

    [Export]
    public bool initialize;

    [Export]
    public bool createLayout;

    [Export]
    public int dilationRadius = 3;


    [ExportGroup("Source")]
    [Export]
    public SourceMode sourceMode;

    [ExportGroup("Source/MultiBaker")]
    [Export]
    public MultiBaker multiBaker = null;

    
    [ExportGroup("Source/Viewports")]
    [Export]
    public Node sourceViewportsContainer = null;

    [Export]
    public SubViewport[] sourceViewports = new SubViewport[ 0 ];

    [ExportGroup("Source/Textures")]
    [Export]
    public Texture2D[] sourceTextures = new Texture2D[ 0 ];

    [ExportGroup("Output")]
    [Export]
    public Node outputTarget;

    public enum LayoutMode
    {
      Grid,
      Custom
    }

    [Export]
    public LayoutMode layoutMode = LayoutMode.Grid;

    [ExportGroup("Output/Custom")]
    [Export]
    public Vector2[] customPositions = new Vector2[ 0 ];

    [Export]
    public Vector2[] customSizes = new Vector2[ 0 ];

    [ExportGroup("Viewport")]
    [Export]
    public Vector2 textureSize = new Vector2( 2048, 2048 );
    [Export]
    public float alphaThreshold = 0.1f;

    [Export]
    public BaseMaterial3D.TransparencyEnum transparencyMode = BaseMaterial3D.TransparencyEnum.AlphaScissor;

    [ExportGroup("Debugging")]
    [Export]
    public SubViewport X_textureMergerViewport;

    [Export]
    public Camera3D X_mergerCamera;

    List<Texture2D> _textures = new List<Texture2D>();

    public override void _Process( double delta )
    {
      if ( initialize )
      {
        initialize = false;
        Initialize();
      }

      if ( createLayout )
      {
        createLayout = false;
        CreateLayout();
      }
    }

    public void Initialize()
    {
      if ( outputTarget == null )
      {
        outputTarget = this;
      }

      if ( X_textureMergerViewport == null )
      {
        X_textureMergerViewport = outputTarget.CreateChild<SubViewport>( "Texture Merger Viewport" );
      }

      Nodes.RemoveAndDeleteChildren( X_textureMergerViewport );

      X_mergerCamera = X_textureMergerViewport.CreateChild<Camera3D>( "Texture Merger Camera" );
      X_mergerCamera.Projection = Camera3D.ProjectionType.Orthogonal;   
      X_mergerCamera.Size = 2;  

      X_textureMergerViewport.Size = (Vector2I) textureSize;
      X_textureMergerViewport.OwnWorld3D = true;
      X_textureMergerViewport.TransparentBg = true;
      X_textureMergerViewport.UseHdr2D = true;

    }

    void GrabTextures()
    {
      _textures = new List<Texture2D>();

      if ( SourceMode.Textures == sourceMode )
      {
        _textures.AddRange( sourceTextures );
        return;
      }

      if ( SourceMode.MultiBaker == sourceMode )
      {        
        var viewports = multiBaker.GetAllViewports().ToArray();
        
        var vpTextures = Arrays.Map( viewports, 
          s => 
          {
            var vt = new ViewportTexture();
            vt.ViewportPath = s.GetPath();
            return vt;          
          }
        );

        _textures.AddRange( vpTextures );
      }


      if ( SourceMode.Viewports == sourceMode )
      {
        if ( sourceViewportsContainer != null )
        {
          sourceViewports = Nodes.GetAll<SubViewport>( sourceViewportsContainer, null, false ).ToArray();
        }

        var vpTextures = Arrays.Map( sourceViewports, 
          s => 
          {
            var vt = new ViewportTexture();
            vt.ViewportPath = s.GetPath();
            return vt;          
          }
        );

        _textures.AddRange( vpTextures );
      }
    }

    public void CreateLayout()
    {
      GrabTextures();

      if ( LayoutMode.Grid == layoutMode )
      {
        CreateGridLayout();
      }
      else if ( LayoutMode.Custom == layoutMode )
      {
        CreateCustomLayout();
      }
    }

    Vector2 ConvertUVtoCameraSpace( Vector2 uv )
    {
      var scale = Vector2.One;
      var w = 1;
      var h = 1;

      if ( X_textureMergerViewport.Size.X != X_textureMergerViewport.Size.Y )
      {
        if ( X_textureMergerViewport.Size.X > X_textureMergerViewport.Size.Y )
        {
          w = X_textureMergerViewport.Size.X / X_textureMergerViewport.Size.Y;
        }
        else
        {
          h = X_textureMergerViewport.Size.Y / X_textureMergerViewport.Size.X;
        }
      }

      var x = Mathf.Remap( uv.X, 0, 1, -w, w );
      var y = Mathf.Remap( uv.Y, 0, 1, -h, h );

      return new Vector2( x, y );
    }

    public Vector2 GetGridTextureSize()
    {
      var alignment = ComputeTextureAlignment( _textures.Count );

      return new Vector2( textureSize.X / alignment.X, textureSize.Y / alignment.Y );
    }

    public int GetMaxTextureDimension()
    {
      var size = GetGridTextureSize();

      return Mathf.RoundToInt( Mathf.Max( size.X, size.Y ) );
    }

    int dilationAlphaTreshold = 30;
    int dilationAssignedMinimum = 0;

    void CreateGridLayout()
    {
      var alignment = ComputeTextureAlignment( _textures.Count );

      // RJLog.Log( "Grid Alignment:", alignment );

      Nodes.RemoveAndDeleteChildrenOfType<CsgMesh3D>( X_textureMergerViewport );
      
      for ( int i = 0; i < _textures.Count; i++ )
      {
        var meshInstance = X_textureMergerViewport.CreateChild<MeshInstance3D>( "Texture " + ( i + 1 ) );

        var uvRectangle = GetUVRectangle( alignment, i );

        // RJLog.Log( "Set Texture" + ( i + 1 ), uvRectangle.min, uvRectangle.max );
        SetMeshCoordinates( meshInstance, uvRectangle, i );      

        var dilatedMaterial = new DilationDrawerMaterial();
        dilatedMaterial.textureAlbedo.Set( _textures[ i ] );
        dilatedMaterial.resolution.Set( _textures[ i ].GetSize() );
        dilatedMaterial.maxRadius.Set( dilationRadius );
        dilatedMaterial.alphaTreshold.Set( dilationAlphaTreshold );
        dilatedMaterial.assignedColorAlphaMinimum.Set( dilationAssignedMinimum );
        dilatedMaterial.genericAlphaOffset.Set( 0 );
        dilatedMaterial.amount.Set( 1 );

        var material = new StandardMaterial3D();
        material.Transparency = transparencyMode;
        material.AlphaScissorThreshold = alphaThreshold;
        material.ShadingMode = BaseMaterial3D.ShadingModeEnum.Unshaded;
        material.AlbedoTexture = _textures[ i ];

        Materials.Set( meshInstance, dilatedMaterial );
      }
    }

    void CreateCustomLayout()
    {
      for ( int i = 0; i < _textures.Count; i++ )
      {
        var meshInstance = outputTarget.CreateChild<MeshInstance3D>( "Texture " + ( i + 1 ) );

        SetMeshCoordinates( meshInstance, customPositions[ i ], customSizes[ i ], i );
        
        var material = new StandardMaterial3D();
        material.Transparency = transparencyMode;
        material.ShadingMode = BaseMaterial3D.ShadingModeEnum.Unshaded;
        material.AlphaScissorThreshold = alphaThreshold;

        var dilatedMaterial = new DilationDrawerMaterial();
        dilatedMaterial.textureAlbedo.Set( _textures[ i ] );
        dilatedMaterial.resolution.Set( _textures[ i ].GetSize() );
        dilatedMaterial.maxRadius.Set( dilationRadius );
        dilatedMaterial.alphaTreshold.Set( dilationAlphaTreshold );
        dilatedMaterial.assignedColorAlphaMinimum.Set( dilationAssignedMinimum );
        dilatedMaterial.genericAlphaOffset.Set( 0 );
        dilatedMaterial.amount.Set( 1 );


        material.AlbedoTexture = _textures[ i ];
        Materials.Set( meshInstance, dilatedMaterial );
      }
    }

    void SetMeshCoordinates( MeshInstance3D mesh, Box2 rectangle, int index )
    {
      SetMeshCoordinates( mesh, rectangle.min, rectangle.max, index ); 
    }

    void SetMeshCoordinates( MeshInstance3D mesh, Vector2 min, Vector2 max, int index )
    {
      var start =  ConvertUVtoCameraSpace( min );
      var end   =  ConvertUVtoCameraSpace( max );

      var size = end - start;

      var quadMesh = new QuadMesh();
      quadMesh.Size = size;
      mesh.Mesh = quadMesh;

      mesh.GlobalPosition = new Vector3( start.X + size.X/2, start.Y + size.Y/2, -1 );

      // RJLog.Log( "Set Mesh", index, "Size:", size, "Min:", min, ">>", start, "Max:", max, ">>", end );
    }

    public static Box2 GetUVAlignmentBoxFor( Vector2I alignment, int index )
    {
      var x = index % alignment.X;
      var y = index / alignment.X;
      
      var sizeX = 1f / alignment.X;
      var sizeY = 1f / alignment.Y;

      return Box2.With( sizeX, sizeY, x * sizeX, y * sizeY );
    }

    public static Vector2I ComputeTextureAlignment( int numElements )
    {
      var root    = Mathf.Sqrt( numElements );     

      var ceiled  = Mathf.CeilToInt( root );
      var floored = ceiled - 1;

      if ( ceiled * floored >= numElements )
      {
        return new Vector2I( ceiled, floored );
      }

      return new Vector2I( ceiled, ceiled );
      
    }

    public static Box2 GetUVRectangle( Vector2I textureAlignment, int index, bool flipY )
    {
      var b = GetUVRectangle( textureAlignment, index );
      
      if ( ! flipY )
      {
        return b;
      }
      
      var uvA = b.min;
      var uvB = b.max;
      uvA.Y = 1f - uvA.Y;
      uvB.Y = 1f - uvB.Y;

      b.min = uvA;
      b.max = uvB;
      
      return b;
    } 

    public static Box2 GetUVRectangle( Vector2I textureAlignment, int index )
    {
      var x = index % textureAlignment.X;
      var y = index / textureAlignment.X;

      var xs = 1f / textureAlignment.X;
      var ys = 1f / textureAlignment.Y;

      var size = new Vector2( xs, ys );

      var min = new Vector2( x * xs, y * ys );
      var max = min + size;

      return new Box2( min, max );
    }

  
  }
}