rj-action-library/Runtime/Procedural/Baking/MultiBaker/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
  {
    [ExportToolButton( "Bake")]
    public Callable BakeButton => Callable.From( () => StartBaking() );

    [Export]
    public bool cleanUpAfterBaking = true;

    [Export]
    public _XX_MultiBakeMode bakeMode;

    public enum MaterialMode
    {
      Full_Seperated, 
      Simple_Prebaked
    }

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

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


    [ExportGroup( "Object")]
    
    [Export]
    public Node3D target;
    
    [Export]
    public bool autoTargetPivot = false;

    [Export]
    public Vector3 targetPivot;

    
  
    [ExportGroup( "Camera")]

    [Export]
    public BakingViewSettings viewSettings;

    public float cameraZoom
    {
      get 
      {
        if ( viewSettings.fovDistance is AutoDistance_BakingFDSettings ad )
        {
          return ad.zoom;
        }
        
        if ( viewSettings.fovDistance is AutoDistance_BakingFDSettings afd )
        {
          return afd.zoom;
        }


        return 1;
      }
    }

    /*
    [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 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;


    [Export]
    public bool X_baking = false;

    public void StartBaking()
    {
      Initialize();
      Bake();
    }

    public void Initialize()
    {     
      this.LogInfo( "Initializing" );
      
      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_textureMerger = this.CreateChild<TextureMerger>( "Texture Merger" );
      X_textureMerger.multiBaker = this;
      X_textureMerger.dilationRadius = dilationRadius;
      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 );                  

    }


    public async Task Bake()
    {
      if ( X_baking )
      {
        return;
      }

      X_baking = true;


      try
      {
        bakeMode.multiBaker = this;

        this.LogInfo( "Started baking" );

        Nodes.RemoveAndDeleteChildren( X_bakingTargetContainer );
        target.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;

        bakeMode.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 fovDistance = viewSettings.fovDistance.ComputeFOVDistance( 
        var objectDistance = GetCameraDistance();

        this.LogInfo( "Set Camera Distance", objectDistance );
        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 ] );

          _bakers.ForEach( b => b.Bake() );

          await this.RequestNextFrame();

          Texture2D texture = X_textureMerger.X_textureMergerViewport.GetTexture();

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

          texture = Textures.Copy( texture );
          await this.RequestNextFrame();
          await this.RequestNextFrame();

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

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

          await this.RequestNextFrame();

        }    

        await this.RequestNextFrame();

        X_outputMesh.GlobalPosition = target.GlobalPosition - targetPivot;

        X_outputMesh.Scale = Vector3.One * cameraZoom;

        if ( cleanUpAfterBaking )
        {
          this.LogInfo( "Cleaning Up" );
          CleanUp();
        }

        this.LogInfo( "All Baking done" );

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

      X_baking = false;
      return;

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

    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 CleanUp()
    {
      var mesh = X_outputMesh;
      mesh.Reparent( GetParent(), true );

      X_bakingViewport = null;

      X_bakingTargetContainer = null;

      X_views = null;

      X_worldEnvironment = null;    
    
      X_outputMesh = null;

      X_textureMerger = null;

      X_setBakingMaterials = null;

      X_texturePreview = null;

      X_depthMapper = null;

      X_bakedTextureAlbedo = null;

      X_bakedTextureNormal = null;

      X_bakedTextureORM = null;

      X_bakedTextureDepth = null;

      Nodes.RemoveAndDeleteChildren( this );

      mesh.Reparent( this, true );
    }

    

    public int GetNumViews()
    {
      return bakeMode.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;


        

        _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;
      _targetBoundingBox = null;

      ComputeBoundingSphere();

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

      ComputeCameraViewSettings();

      bakeMode.CreateBakers();

      

    }

    Sphere _targetBoundingSphere;
    Box3 _targetBoundingBox;

    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;
      }

      _targetBoundingBox = X_bakingTargetContainer.GetWorldBounds();

      if ( autoTargetPivot )
      { 
        Box3 box = target.GetWorldBounds();
        targetPivot = new Vector3( 0, -box.center.Y, 0 );
      }
      
      _targetBoundingSphere = Sphere.ContainingBox( _targetBoundingBox );
    }

    float _cameraFOV = 0;
    float _cameraDistance = 0;
    float _outputScale = 1;
    
    void ComputeCameraViewSettings()
    {
      var fovDistance = viewSettings.fovDistance;
      var size = targetBoundingSphere.radius;

      var computeScale = false;

      if ( fovDistance is AutoDistance_BakingFDSettings ad )
      {
        size = ad.sizeEstimation == _XX_BakingFDSettings.SizeEstimationType.Bounding_Sphere ? 
                targetBoundingSphere.radius : ( _targetBoundingBox.size.Y / 2f );
        computeScale = true;
      } 

      if ( fovDistance is AutoFOVDistance_BakingFDSettings afd )
      {
        size = afd.sizeEstimation == _XX_BakingFDSettings.SizeEstimationType.Bounding_Sphere ? 
                targetBoundingSphere.radius : ( _targetBoundingBox.size.Y / 2f );
        computeScale = true;
      } 
      
      var fd = fovDistance.ComputeFOVDistance( size / 2f );

      _cameraFOV = fd.X;
      _cameraDistance = fd.Y;
      _outputScale = computeScale ? Cameras.ComputeCameraFittingScale( _cameraFOV, _cameraDistance ) : 1;
    }
    
    public float GetCameraFOV()
    {
      return _cameraFOV;
    }

    public float GetCameraDistance()
    {
      return _cameraDistance;
    }

    public float GetOutputScale()
    {
      return _outputScale;
    }
  }
}