From 338522b259e4b06d756fabf1cf938c13937f4928 Mon Sep 17 00:00:00 2001
From: AlvaroHG <alvaroherrasti@gmail.com>
Date: Fri, 22 Nov 2024 14:27:41 -0800
Subject: [PATCH] Adds back distortion map

---
 test_distortion.py                       |  2 +-
 unity/Assets/Scripts/CapturePass.cs      | 43 ++++++++++++++++++++----
 unity/Assets/Scripts/RenderingManager.cs |  5 +++
 3 files changed, 43 insertions(+), 7 deletions(-)

diff --git a/test_distortion.py b/test_distortion.py
index 74af454ef1..92b4732237 100644
--- a/test_distortion.py
+++ b/test_distortion.py
@@ -26,7 +26,7 @@ def load_scene(scene_name, house_path=None, run_in_editor=False, platform=None,
             server_class=ai2thor.wsgi_server.WsgiServer,
         )
 
-    enableDistortionMap = False
+    enableDistortionMap = True
 
     all_args = dict(
         # local_executable_path="unity/builds/thor-OSXIntel64-local/thor-OSXIntel64-local.app/Contents/MacOS/AI2-THOR",
diff --git a/unity/Assets/Scripts/CapturePass.cs b/unity/Assets/Scripts/CapturePass.cs
index 97cc483a28..80ee117773 100644
--- a/unity/Assets/Scripts/CapturePass.cs
+++ b/unity/Assets/Scripts/CapturePass.cs
@@ -58,6 +58,7 @@ public class CaptureConfig {
 
         public bool cloudRendering;
         
+        public RenderTextureFormat renderTextureFormat = RenderTextureFormat.ARGB32;
         
     }
     
@@ -118,7 +119,7 @@ public static void SetupCameraWithPostShader(
         
         private TextureFormat readTextureFormat;
 
-        
+        private RenderTextureFormat renderTextureFormat;
 
         // private Texture2D readTexture;
 
@@ -133,6 +134,7 @@ public RenderToTexture(CaptureConfig config, Camera camera) {
             this.name = config.name;
             this.cloudRendering = config.cloudRendering;
             this.toDisplayId = config.toDisplay;
+            this.renderTextureFormat = config.renderTextureFormat;
 
             // TODO. if config.toDisplay is present then render to display buffer and copy to render texture 
             // for debugging purposes
@@ -214,10 +216,26 @@ private RenderTexture CreateRenderTexture(int width, int height) {
        
          
         if (cloudRendering) {
+            GraphicsFormat cloudRenderingRTFormat;
+            if (this.renderTextureFormat == RenderTextureFormat.RGFloat) {
+                readTextureFormat = TextureFormat.RGFloat;
+                cloudRenderingRTFormat = GraphicsFormat.R32G32_SFloat;
+            } 
+            else if (this.renderTextureFormat == RenderTextureFormat.RFloat) {
+                readTextureFormat = TextureFormat.RFloat;
+                cloudRenderingRTFormat = GraphicsFormat.R32_SFloat;
+            } 
+            else {
+                readTextureFormat = TextureFormat.RGBA32;
+                cloudRenderingRTFormat = GraphicsFormat.R8G8B8A8_UNorm;
+            }
+
             // Why 0 for depth here ?
-            rt = new RenderTexture(Screen.width, Screen.height, 0, GraphicsFormat.R8G8B8A8_UNorm);
-            // TODO: if 0 then RGB24? if not RGB32?
-            readTextureFormat = TextureFormat.RGBA32;
+            rt = new RenderTexture(Screen.width, Screen.height, 0, cloudRenderingRTFormat);
+            // // Why 0 for depth here ?
+            // rt = new RenderTexture(Screen.width, Screen.height, 0, GraphicsFormat.R8G8B8A8_UNorm);
+            // // TODO: if 0 then RGB24? if not RGB32?
+            // readTextureFormat = TextureFormat.RGBA32;
             
         //     RenderTexture rt = new RenderTexture(
         //     width: width,
@@ -227,8 +245,21 @@ private RenderTexture CreateRenderTexture(int width, int height) {
         // );
         }
         else {
-            rt = new RenderTexture(Screen.width, Screen.height, 24, RenderTextureFormat.ARGB32, RenderTextureReadWrite.Default);
-            readTextureFormat = TextureFormat.RGBA32; 
+            if (this.renderTextureFormat == RenderTextureFormat.RGFloat) {
+                readTextureFormat = TextureFormat.RGFloat;
+            } 
+            else if (this.renderTextureFormat == RenderTextureFormat.RFloat) {
+                readTextureFormat = TextureFormat.RFloat;
+            } 
+            else {
+                readTextureFormat = TextureFormat.RGBA32;
+                this.renderTextureFormat = RenderTextureFormat.ARGB32;
+            }
+
+            rt = new RenderTexture(Screen.width, Screen.height, 24, renderTextureFormat, RenderTextureReadWrite.Default);
+
+            // rt = new RenderTexture(Screen.width, Screen.height, 24, RenderTextureFormat.ARGB32, RenderTextureReadWrite.Default);
+            // readTextureFormat = TextureFormat.RGBA32; 
         }
 
          if (this.tex == null) {
diff --git a/unity/Assets/Scripts/RenderingManager.cs b/unity/Assets/Scripts/RenderingManager.cs
index 8c4be8f381..438af2cbd7 100644
--- a/unity/Assets/Scripts/RenderingManager.cs
+++ b/unity/Assets/Scripts/RenderingManager.cs
@@ -189,6 +189,10 @@ void Awake() {
             new CaptureConfig() { name = "_normals", antiAliasLevel = antiAliasLevel, shaderName = "Hidden/UberReplacement", replacementMode = ReplacelementMode.Normals },
             cameraParent: camera.transform
         );
+
+        this.distortionMap = new OnDemandCapture(
+            new CaptureConfig() { name = "_distortion_map", antiAliasLevel = antiAliasLevel, shaderName = "Custom/BarrelDistortionMap" , cloudRendering = cloudRenderingCapture, toDisplay = 7, renderTextureFormat = RenderTextureFormat.RGFloat }
+        );
         
         // make first _img capture created render to Display
         int? toDisplay = null;
@@ -203,6 +207,7 @@ void Awake() {
             this.mainPass,
             depthPass,
             distPass,
+            this.distortionMap,
             idPass,
             classPass
         }.ToDictionary(x => x.GetName(), x => x);