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

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

using System.Threading.Tasks;

namespace Rokojori
{
  [Tool]
  [GlobalClass]
  public partial class MultiBaker:Node
  {
    [Export]
    public bool initialize;

    public bool bake;

    public bool saveTexture;

    [ExportGroup("Preview")]
        
    [Export]
    public bool preview_UpdateAlways = true;

    [Export]
    public bool preview_DilateTextures = true;  

    public enum MaterialMode
    {
      Full_Seperated, 
      Simple_Prebaked
    }

    [ExportGroup("Material")]
    [Export]
    public MaterialMode materialMode;

    [ExportGroup("Material/Full Seperated")]
    [Export]
    public bool mmfs_Normals = true; 
    [Export]
    public bool mmfs_Depth = true;
    [Export]
    public bool mmfs_ORM = true;


    public enum BakeMode
    {
      Cylinder,
      Cross_Braces,
      Octahedral,
      Spherical,
      Cloud
    }

    [ExportGroup("BakeMode")]
    
    [Export]
    public BakeMode bakeMode;



    [ExportGroup("BakeMode/Cylinder")]
    [Export]
    public int cylinderSides = 4;
    [Export]
    public bool createBackFacesForSides = false;
    [Export]
    public bool cylinderTop = false;
    [Export]
    public bool createBackFacesForTop = false;
    [Export]
    public bool cylinderBottom = false;
    [Export]
    public bool createBackFacesForBottom = false;
    [Export( PropertyHint.Range, "-1,1" )]
    public float cylinderSideOffset = 0;
    [Export( PropertyHint.Range, "-1,1" )]
    public float cylinderTopOffset = 0f;
    [Export( PropertyHint.Range, "-1,1" )]
    public float cylinderBottomOffset = 0f;    

    [ExportGroup("BakeMode/Cross Braces")]
    [Export]
    public int crossAngles = 2;
    [Export]
    public int crossBraces = 3;
    [Export]
    public float crossSpreadDistance = 1;
    [Export]
    public float crossAngleOffset = 0;

    [Export]
    public bool crossTop = false;

    [Export]
    public bool crossBottom = false;


    [ExportGroup("BakeMode/Octahedral")]
    [Export]
    public int octahedralSides = 4;
    [Export]
    public bool octahedralFullSphere = false;

    [ExportGroup("BakeMode/Spherical")]
    [Export]
    public int sphericalSides = 4;
    [Export]
    public int sphericalRows = 3;
    [Export]
    public float minPitch = -30;    
    [Export]
    public float maxPitch = 30;

    [ExportGroup( "Object")]
    
    [Export]
    public Node3D sourceTarget;
    
    [Export]
    public bool autoCenter = false;

    [Export]
    public Vector3 sourceOffset;


  
    [ExportGroup( "Camera")]
    [Export]
    public Baker.CameraDistanceDetectionType distanceDetectionType = Baker.CameraDistanceDetectionType.Automatic_Distance_Detection;

    [Export]
    public float customDistance = 50;

    [Export]
    public float cameraZoom = 1;

    [Export]
    public Baker.CameraFOVMode fovMode = Baker.CameraFOVMode.Compute_Fov_With_Distance;
    [Export]
    public float originalFOV = 75;
    [Export]
    public float fovPlacingDistance = 200;
    [Export]
    public float customFOV = 75;



    [ExportGroup( "Output")]


    [Export]
    public string outputDirectory;

    [Export]
    public string outputFileName;

    [Export]
    public float outputQuality = 1f;

    [Export]
    public Vector2 outputTextureSize = new Vector2( 2048, 2048 );    

    [Export]
    public bool showOutputTexture = false;


    [ExportGroup("Read Only")]
    [Export]
    public SubViewport X_bakingViewport;

    [Export]
    public Node3D X_bakingTargetContainer;

    [Export]
    public Node X_views;

    [Export]
    public WorldEnvironment X_worldEnvironment;    

    [Export]
    public DilateTexture X_dilateTexture;
  
    [Export]
    public MeshInstance3D X_outputMesh;

    [Export]
    public TextureMerger X_textureMerger;

    [Export]
    public SetBakingMaterials X_setBakingMaterials;

    [Export]
    public MeshInstance3D X_texturePreview;

    [Export]
    public CsgMesh3D X_depthMapper;

    [Export]
    public Texture2D X_bakedTextureAlbedo;

    [Export]
    public Texture2D X_bakedTextureNormal;

    [Export]
    public Texture2D X_bakedTextureORM;

    [Export]
    public Texture2D X_bakedTextureDepth;




    bool _initialized = false;
    bool _baking = false;


    SerializedGodotObject _cached;

    
    public override void _Ready()
    {
      Initialize();
    }
    
    public override void _Process( double delta )
    { 
      X_texturePreview.Visible = showOutputTexture;

      if ( _baking )
      {
        return;
      }
     
      if ( initialize || ! _initialized )
      {
        initialize = false;
        Initialize();
      } 


      var current = SerializedGodotObject.Create( this ); 

      var changed = _cached != null && ! _cached.Equals( current );

      _cached = current;

      if ( bake || changed && preview_UpdateAlways )
      {
        bake = false;
        Bake();
      }
     
    } 

    
    public async Task Bake()
    {
      _baking = true;

      try
      {

        this.LogInfo( "Started baking" );

        Nodes.RemoveAndDeleteChildren( X_bakingTargetContainer );
        sourceTarget.DeepCopyTo( X_bakingTargetContainer );

        if ( _bakers == null || _bakers.Count != GetNumViews() )
        {
          CreateViews();
          await this.RequestNextFrame();
        }


        SetupViews();

        this.LogInfo( "Views set up" );
        await this.RequestNextFrame();

        X_setBakingMaterials.SetTarget( X_bakingTargetContainer );


        var bakingMaterialModes = new List<BakingMaterialMode>();

        var preview_QuickMaterial = MaterialMode.Simple_Prebaked == materialMode;

        GetBakeModeImplementation().CreateMaterial( preview_QuickMaterial );

        if ( preview_QuickMaterial )
        { 
          bakingMaterialModes.Add( BakingMaterialMode.Preview );
        }
        else
        {
          bakingMaterialModes.Add( BakingMaterialMode.Albedo );

          if ( mmfs_Normals )
          {
            bakingMaterialModes.Add( BakingMaterialMode.Normals );
          }

          if ( mmfs_Depth )
          {
            bakingMaterialModes.Add( BakingMaterialMode.Depth );
          }

          if ( mmfs_ORM )
          {
            bakingMaterialModes.Add( BakingMaterialMode.ORM );
          }
        }

        this.LogInfo( "Prepared baking modes" );

        X_textureMerger.textureSize = outputTextureSize;
        X_textureMerger.Initialize();
        X_textureMerger.CreateLayout();

        this.LogInfo( "Prepared texture merger" );

        var objectDistance = GetCameraDistance();

        for ( int i = 0; i < bakingMaterialModes.Count; i++ )
        {
          this.LogInfo( "Baking mode:", bakingMaterialModes[ i ] );   

          X_setBakingMaterials.mode = bakingMaterialModes[ i ];
          X_setBakingMaterials.ApplyBakingMaterials( objectDistance, _targetBoundingSphere.radius );

          this.LogInfo( "Materials changed:", bakingMaterialModes[ i ] );

          await this.RequestNextFrame();

          Texture2D texture = X_textureMerger.X_textureMergerViewport.GetTexture();

          this.LogInfo( "Texture created:", bakingMaterialModes[ i ] );

          if ( preview_DilateTextures )
          {
            this.LogInfo( "Dilating:", bakingMaterialModes[ i ] );
            texture = await CreateDilatedTexture();

            this.LogInfo( "Dilating done:", bakingMaterialModes[ i ] );
          }        
          else
          {
            texture = Textures.Copy( texture );
            await this.RequestNextFrame();
            await this.RequestNextFrame();
          }


          this.LogInfo( "Assigning Texture", bakingMaterialModes[ i ] );
          GetBakeModeImplementation().AssignMaterial( bakingMaterialModes[ i ], texture ); 

          this.LogInfo( "Baking done:", bakingMaterialModes[ i ] );

          await this.RequestNextFrame();

        }    

        await this.RequestNextFrame();

        // this.LogInfo( "Baking done" );

      }
      catch ( System.Exception e )
      {
        this.LogError( "Baking failed" );
        this.LogError( e );
      }

      _baking = false;
      return;

    }
    
    
    public async Task<Texture2D> CreateDilatedTexture()
    { 
      var viewports = GetAllViewports();
      var textures = Lists.Map( viewports, v => v.GetTexture() as Texture2D );
      var dilatedTextures = new List<Texture2D>();

      var index = 0;

      foreach ( var t in textures )
      {
        index ++;

        var dilatedTexture = await X_dilateTexture.Create( t );
        dilatedTextures.Add( dilatedTexture );
      }

      X_textureMerger.sourceTextures = dilatedTextures.ToArray();
      X_textureMerger.sourceMode = TextureMerger.SourceMode.Textures;      

      X_textureMerger.Initialize();

      await this.RequestNextFrame();

      X_textureMerger.CreateLayout();

      await this.RequestNextFrame();

      var finalTexture = await X_dilateTexture.Create( X_textureMerger.X_textureMergerViewport.GetTexture() );

      return finalTexture;
    }

    public List<SubViewport> GetAllViewports()
    {
      return Nodes.AllIn<SubViewport>( X_views, null, false );
    }

    List<MultiBakeModeImplementation> _bakeModeImplementations = new List<MultiBakeModeImplementation>()
    {
      new MultiBakeModeCylinder(),
      new MultiBakeModeCrossBraces(),
      new MultiBakeModeOctahedral()
    };

    public MultiBakeModeImplementation GetBakeModeImplementation()
    {
      var bm = _bakeModeImplementations.Find( bm => bm.GetBakeMode() == bakeMode );
      bm.multiBaker = this;
      return bm;
    }

    public void CacheTexture( BakingMaterialMode mode, Texture2D texture )
    {
      if ( BakingMaterialMode.Albedo == mode )
      {
        X_bakedTextureAlbedo = texture;
      }
      else if ( BakingMaterialMode.Normals == mode )
      {
        X_bakedTextureNormal = texture;
      }
      else if ( BakingMaterialMode.ORM == mode )
      {
        X_bakedTextureORM = texture;
      }
      else if ( BakingMaterialMode.Depth == mode )
      {
        X_bakedTextureDepth = texture;
      }

    }

    public void Initialize()
    {     
      Nodes.RemoveAndDeleteChildren( this );

      _bakers = null;

      X_bakingViewport = this.CreateChild<SubViewport>( "Multi Baker Viewport" );

      X_bakingViewport.Size = (Vector2I) outputTextureSize;
      X_bakingViewport.OwnWorld3D = true;
      X_bakingViewport.TransparentBg = true;

      X_worldEnvironment = X_bakingViewport.CreateChild<WorldEnvironment>( "Multi Baker Environment" );
      X_worldEnvironment.Environment = new Godot.Environment();
      X_worldEnvironment.Environment.AmbientLightSource = Godot.Environment.AmbientSource.Color;
      X_worldEnvironment.Environment.AmbientLightColor = HSLColor.white;

      X_bakingTargetContainer = X_bakingViewport.CreateChild<Node3D>( "Target Container" );

      X_views = X_bakingViewport.CreateChild<Node>( "Views" );

      X_dilateTexture = this.CreateChild<DilateTexture>( "Dilate Texture" ); 

      X_textureMerger = this.CreateChild<TextureMerger>( "Texture Merger" );
      X_textureMerger.multiBaker = this;
      X_textureMerger.Initialize();

      X_outputMesh = this.CreateChild<MeshInstance3D>( "Output Mesh" );

      X_texturePreview = this.CreateChild<MeshInstance3D>( "Texture Preview" );
      X_texturePreview.Mesh = new QuadMesh();
      X_texturePreview.Scale = Vector3.One * 100;

      var pm = new StandardMaterial3D();
      pm.Transparency = BaseMaterial3D.TransparencyEnum.AlphaScissor;
      pm.ResourceLocalToScene = true;

      var vt = new ViewportTexture();
      vt.ViewportPath = X_textureMerger.X_textureMergerViewport.GetPath();
      pm.AlbedoTexture = vt;

      X_setBakingMaterials = this.CreateChild<SetBakingMaterials>( "Set Baking Materials" );

      Materials.Set( X_texturePreview, pm );            
      
      _initialized = true;      

    }

    public int GetNumViews()
    {
      return GetBakeModeImplementation().GetNumViews();
    }

    List<Baker> _bakers;

    public List<Baker> bakers => _bakers;

    public void CreateViews()
    {    

      Nodes.RemoveAndDeleteChildren( X_views );

      var numViews = GetNumViews();

      _bakers = new List<Baker>();

      var minViewsPerAxis = TextureMerger.ComputeTextureAlignment( numViews ).Y;

      for ( int i = 0; i < numViews; i++ )
      {
        var userIndex = ( i + 1 );
        var bakingView = X_views.CreateChild<SubViewport>( "Baking View " + userIndex );
        bakingView.TransparentBg = true;
        bakingView.Size = (Vector2I) ( outputTextureSize / minViewsPerAxis );
        
        var bakingCamera = bakingView.CreateChild<Camera3D>( "Camera View " + userIndex );
        var baker = bakingView.CreateChild<Baker>( "Baker " + userIndex );

        baker.camera   = bakingCamera;
        baker.target   = X_bakingTargetContainer;
        baker.viewport = bakingView;


        baker.update = true;

        _bakers.Add( baker );
      }
    }

    public void SetupViews()
    {
      var numViews = GetNumViews();
      var minViewsPerAxis = TextureMerger.ComputeTextureAlignment( numViews ).Y;

      X_bakingViewport.Size = (Vector2I) outputTextureSize;

      for ( int i = 0; i < numViews; i++ )
      {
        var baker = _bakers[ i ];
        var bakingView = baker.viewport as SubViewport;
        bakingView.Size = (Vector2I) ( outputTextureSize / minViewsPerAxis );
      }

      _targetBoundingSphere = null;
      ComputeBoundingSphere();

      if ( _targetBoundingSphere == null )
      {
        this.LogError( "No bounding sphere created, ensure there are visible targets" );
        return;
      }

      var bmi = GetBakeModeImplementation();
      
      if ( bmi == null )
      {
        return;
      }

      bmi.CreateBakes();

      /*if ( X_bakedTextureAlbedo != null )
      {
        bmi.AssignMaterial( BakingMaterialMode.Albedo, X_bakedTextureAlbedo );
      }*/
    }

    Sphere _targetBoundingSphere;

    Sphere targetBoundingSphere 
    {
      get 
      {
        if ( _targetBoundingSphere == null )
        {
          ComputeBoundingSphere();
        }

        return _targetBoundingSphere;
      }
    }

    void ComputeBoundingSphere()
    {
      if ( X_bakingTargetContainer.GetChildCount() == 0 )
      {
        _targetBoundingSphere = null;
        return;
      }

      var firstChild = X_bakingTargetContainer.GetChild( 0 ) as Node3D;

      if ( firstChild == null )
      {
        _targetBoundingSphere = null;
        return;
      }

      if ( ! firstChild.Visible )
      {
        firstChild.Visible = true;
      }

      var worldBounds = X_bakingTargetContainer.GetWorldBounds();

      
      _targetBoundingSphere = Sphere.ContainingBox( worldBounds );
    }

    public float GetCameraFOV()
    {
      if ( Baker.CameraFOVMode.Custom_Fov == fovMode )
      {
        return customFOV;
      }

      if ( Baker.CameraFOVMode.Keep_Fov == fovMode )
      {
        return originalFOV;
      }

      return Cameras.ComputeFOVForBillboard( originalFOV, targetBoundingSphere.radius, fovPlacingDistance );
    }

    public float GetCameraDistance()
    {
      if ( Baker.CameraDistanceDetectionType.Custom_Distance == distanceDetectionType )
      {
        return customDistance;
      }

      var fov = GetCameraFOV();

      return Cameras.ComputeCameraFrameFittingDistance( fov, targetBoundingSphere.radius / cameraZoom );
    }

    public float GetOutputScale()
    {
      var fov = GetCameraFOV();
      var distance = GetCameraDistance();

      return Cameras.ComputeCameraFittingScale( fov, distance );
    }
  }
}