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

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

using System.Threading.Tasks;
using System.Linq;

namespace Rokojori
{
  [Tool]
  [GlobalClass]
  public partial class MultiBaker:Node
  {
    [ExportToolButton( "Bake")]
    public Callable BakeButton => Callable.From( () => Bake() );

    [Export]
    public bool cleanUpAfterBaking = true;

    [Export]
    public _XX_MultiBakeMode bakeMode;

    [Export]
    public _XX_MultiTextureBaker textureBaker;


    [Export]
    public BakingViewSettings viewSettings;

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

    [Export]
    public Vector3 targetPivot;

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

    [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 TextureDilate X_textureDilate;

    [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 async Task Bake()
    {
      Initialize();
      await CreateBakes();

      return;
    }

    public void Initialize()
    {     
      this.LogInfo( "Initializing" );
      
      Nodes.RemoveAndDeleteChildren( this );

      _bakerCameras = 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_textureDilate = this.CreateChild<TextureDilate>( "Texture Dilate" );
      X_textureDilate.viewport = X_textureMerger.X_textureMergerViewport;

      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<Texture2D> GrabDilatedTexture( bool srgb, bool alpha, bool mipmaps )
    {
      X_textureDilate.viewport = X_textureMerger.X_textureMergerViewport;
      var maxRadius = X_textureMerger.GetMaxTextureDimension() / 2;
      X_textureDilate.SetDilationRadius( maxRadius );
      var texture = await X_textureDilate.Grab( srgb, alpha, mipmaps );

      return texture;
    }   

    bool _targetIsInInitialState = false;

    public void SetTargetDirty()
    {
      _targetIsInInitialState = false;
    }

    public void ResetOriginalTargetState()
    {
      if ( _targetIsInInitialState )
      {
        return;
      }

      // _bakerCameras.ForEach( b => b.viewport.en)

      Nodes.RemoveAndDeleteChildren( X_bakingTargetContainer );
      var root = (Node3D)target.DeepCopyTo( X_bakingTargetContainer );

      root.GlobalPosition = Vector3.Zero;
      root.SetGlobalQuaternion( Quaternion.Identity );
      root.Scale = Vector3.One;


      _targetIsInInitialState = true;

      return;
    }


    bool _compositorsClean = false;

    public void SetCompositorsDirty()
    {
      _compositorsClean = false;
    }


    public void ResetCompositorStates()
    {
      X_worldEnvironment.Compositor = null;
      _compositorsClean = true;
    }

    public void AddCompositorEffect( CompositorEffect compositorEffect )
    {
      if ( X_worldEnvironment.Compositor == null )
      {
        X_worldEnvironment.Compositor = new Compositor();
      }

      var compositorEffects = X_worldEnvironment.Compositor.CompositorEffects;
      compositorEffects.Add( compositorEffect );

      X_worldEnvironment.Compositor.CompositorEffects = compositorEffects;

      _compositorsClean = false; 
    } 
   

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

      X_baking = true;


      try
      {
        bakeMode.multiBaker = this;

        _targetIsInInitialState = false;

        this.LogInfo( "Started baking" );

        ResetOriginalTargetState();

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

        SetupBakerCameras();

        this.LogInfo( "Bake Cameras set up" );
        await this.RequestNextFrame();

        this.LogInfo( "Render passes" );
        var passes = textureBaker.GetPasses( this );



        bakeMode.CreateMaterial( passes );


        this.LogInfo( "Prepared baking modes" );
        X_textureMerger.textureSize = outputTextureSize;
        X_textureMerger.Initialize();
        X_textureMerger.CreateLayout();


        for ( int i = 0; i < passes.Count; i++ )
        { 
          ResetOriginalTargetState();
          
          _bakerCameras.ForEach( b => b.ApplyCameraSettings() );
          await passes[ i ].Bake();
          await this.RequestNextFrame();

        }

        bakeMode.AssignMaterial( passes );

        X_outputMesh.GlobalPosition = target.GlobalPosition;// - targetPivot;

        X_outputMesh.Scale = Vector3.One * cameraZoom;

       
        this.LogInfo( "All Baking done" );

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

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


      X_baking = false;
      return;

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

    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_textureDilate = 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> _bakerCameras;

    public List<Baker> bakerCameras => _bakerCameras;

    public void CreateBakerCameras()
    {    

      Nodes.RemoveAndDeleteChildren( X_views );

      var numViews = GetNumViews();

      _bakerCameras = new List<Baker>();

      var alginment = TextureMerger.ComputeTextureAlignment( numViews );
      var minViewsPerAxis = Mathf.Max( alginment.X, alginment.Y );

      this.LogInfo( "Size per Baker:", "nv", numViews, "al", alginment, "mvpa", minViewsPerAxis,"s", (Vector2I) ( outputTextureSize / minViewsPerAxis ) );
      
      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.viewport.Msaa3D = Viewport.Msaa.Msaa8X;
        


        

        _bakerCameras.Add( baker );
      }
    }

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

      X_bakingViewport.Size = (Vector2I) outputTextureSize;

      for ( int i = 0; i < numViews; i++ )
      {
        var baker = _bakerCameras[ 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 ( _targetBoundingBox == null )
      {
        return;
      }

      if ( autoTargetPivot )
      { 
        Box3 box = target.GetWorldBounds();

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

        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 = Mathf.Clamp( fd.X, 1, 179 );
      _cameraDistance = fd.Y;
      _outputScale = computeScale ? Cameras.ComputeCameraFittingScale( _cameraFOV, _cameraDistance ) : 1;

      RJLog.Log( "Computed FOV Distance", _cameraFOV, _cameraDistance );
    }
    
    public float GetCameraFOV()
    {
      return Mathf.Clamp( _cameraFOV, 1, 179 );
    }

    public float GetCameraDistance()
    {
      return _cameraDistance;
    }

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