Adds back distortion map

allenai · Nov 22, 2024 · 338522b · 338522b
1 parent 837f7a2
commit 338522b
Show file tree

Hide file tree

Showing 3 changed files with 43 additions and 7 deletions.
diff --git 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
@@ -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
@@ -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);