Third party image distortion

ai2thor/interact.py

@@ -229,6 +229,7 @@ def write_image(
         instance_segmentation_frame=False,
         depth_frame=False,
         distortion_frame=False,
+        third_party_camera_frames=False,
         metadata=False,
     ):
         def save_image(name, image, flip_br=False):
@@ -298,6 +299,22 @@ def json_write(name, obj):
                     x.astype(np.float32),
                 ),
             ),
+            (
+                "third_party_camera_frames",
+                third_party_camera_frames,
+                lambda event: event.third_party_camera_frames,
+                lambda x: x,
+                lambda name, images: [save_image(name, image, flip_br=True) for  (i, image) in zip(range(len(images)), images)]
+
+            ),
+            (
+                "third_party_distortion_frames",
+                third_party_camera_frames and distortion_frame,
+                lambda event: event.third_party_distortion_frames,
+                lambda x: x,
+                lambda name, images: [save_image(name, image, flip_br=True) for  (i, image) in zip(range(len(images)), images)]
+
+            ),
             (
                 "metadata",
                 metadata,

ai2thor/server.py

@@ -459,6 +459,7 @@ def __init__(self, metadata):
         self.third_party_instance_masks = []
         self.third_party_class_masks = []
         self.third_party_camera_frames = []
+        self.third_party_distortion_frames = []
         self.third_party_semantic_segmentation_frames = []
         self.third_party_instance_segmentation_frames = []
         self.third_party_depth_frames = []

test_distortion.py

@@ -68,16 +68,38 @@ def load_scene(scene_name, house_path=None, run_in_editor=False, platform=None,
         )
         print(f"Action {controller.last_action['action']} success: {evt.metadata['lastActionSuccess']}")
         print(f'Error: {evt.metadata["errorMessage"]}')
+
+
+    addThirdPartyCam = {'action': 'AddThirdPartyCamera', 'agentPositionRelativeCoordinates': True, 'fieldOfView': 139, 'parent': 'agent', 'position': {'x': 0.04, 'y': 0.5560812, 'z': 0.0}, 'rotation': {'x': 30.0, 'y': 120.0, 'z': 0.0}}
+
+    evt = controller.step(**addThirdPartyCam)
+    print(f"Action {controller.last_action['action']} success: {evt.metadata['lastActionSuccess']}")
+    print(f'Error: {evt.metadata["errorMessage"]}')
+
     if distortion:
+        # Good parameters too, for previous frame
+        # controller.step(
+        #     action="SetDistortionShaderParams",
+        #     zoomPercent=0.46,
+        #     k1=1.09,
+        #     k2=1.92,
+        #     k3=3.1,
+        #     k4=1.8,
+        #     intensityX=0.9,
+        #     intensityY=0.97
+        # )
+
         controller.step(
             action="SetDistortionShaderParams",
-            zoomPercent=0.46,
-            k1=1.09,
-            k2=1.92,
-            k3=3.1,
-            k4=1.8,
-            intensityX=0.9,
-            intensityY=0.97
+            mainCamera=True,
+            thidPartyCameraIndices=[0],
+            zoomPercent=0.54,
+            k1=0.5,
+            k2=4.5,
+            k3=-13.1,
+            k4=14.1,
+            intensityX=0.91,
+            intensityY=0.93
         )
 
     xpos = dict(x=0.0, y=0.900992214679718, z=0.0786)
@@ -93,19 +115,40 @@ def load_scene(scene_name, house_path=None, run_in_editor=False, platform=None,
     # controller.step(
     #     {"action": "RotateCameraMount", "degrees": 13, "secondary": False}
     # )
+    cam_param = {
+        "position": {"x": -0.1211464, "y": 0.561659, "z": 0.03892733},
+        "rotation": {"x": 20.0, "y": 0.0, "z": 0.0},
+        "fov": 120,
+        "index": "main",
+    }
     cam_pos = {"x": -0.1211464, "y": 0.561659, "z": 0.03892733}
     cam_rot = {"x": 13.0, "y": 0.0, "z": 0.0}
     event = controller.step(
         action="UpdateMainCamera",
-        position=cam_pos,
-        rotation=cam_rot,
-        fieldOfView=120,
+        position=cam_param["position"],
+        rotation=cam_param["rotation"],
+        fieldOfView=cam_param["fov"],
         agentId=0,
     )
 
-    InteractiveControllerPrompt.write_image(controller.last_event, image_dir, "", semantic_segmentation_frame=True, depth_frame=True, color_frame=True, distortion_frame=distortion)
+    print(f"Action {controller.last_action['action']} success: {evt.metadata['lastActionSuccess']}")
+    print(f'Error: {evt.metadata["errorMessage"]}')
+
+    InteractiveControllerPrompt.write_image(controller.last_event, image_dir, "", semantic_segmentation_frame=True, depth_frame=True, color_frame=True, third_party_camera_frames=True, distortion_frame=distortion)
     # input()
 
+# def reset_agent_embodiment(self, controller, main_camera_params, vertical_fov):
+#     position = main_camera_params["position"].copy()
+#     rotation = main_camera_params["rotation"].copy()
+#     fov = vertical_fov
+
+#     event = controller.step(
+#         action="UpdateMainCamera",
+#         position=position,
+#         rotation=rotation,
+#         fieldOfView=fov,
+#         agentId=0,
+#     )
 
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()

unity/Assets/Scripts/AgentManager.cs

@@ -1856,20 +1856,24 @@ public void SetCriticalErrorState() {
         this.agentManagerState = AgentState.Error;
     }
 
-    public ActionFinished SetDistortionShaderParams(float zoomPercent, float k1, float k2, float k3, float k4, float strength = 1.0f, float intensityX = 1.0f, float intensityY = 1.0f) {
-
-        if (this.primaryAgent.imageSynthesis == null) {
+    public ActionFinished SetDistortionShaderParams(bool mainCamera = true, IEnumerable<int> thidPartyCameraIndices = null, float zoomPercent = 1.0f, float k1 = 0.0f, float k2 = 0.0f, float k3 = 0.0f, float k4 = 0.0f, float strength = 1.0f, float intensityX = 1.0f, float intensityY = 1.0f) {
+
+        IEnumerable<ImageSynthesis> imageSynths = mainCamera ? new List<ImageSynthesis>() {this.primaryAgent.imageSynthesis} : new List<ImageSynthesis>();
+        imageSynths = thidPartyCameraIndices != null ? imageSynths.Concat(this.thirdPartyCameras.Where((cam, i) => thidPartyCameraIndices.Contains(i)).Select(cam => cam.GetComponent<ImageSynthesis>())) : imageSynths;
+        if (imageSynths.Any(x => x == null)) {
             return new ActionFinished(success: false, errorMessage: "No imageSynthesis, make sure you pass 'renderDistortionImage = true' to the agent constructor.");
         }
-        var material = this.primaryAgent.imageSynthesis.distortionMaterial;
-        material.SetFloat("_ZoomPercent", zoomPercent);
-        material.SetFloat("_k1", k1);
-        material.SetFloat("_k2", k2);
-        material.SetFloat("_k3", k3);
-        material.SetFloat("_k4", k4);
-        material.SetFloat("_DistortionIntensityX", intensityX);
-        material.SetFloat("_DistortionIntensityY", intensityY);
-        material.SetFloat("_LensDistortionStrength", strength);
+        foreach (var imageSynthesis in imageSynths) {
+            var material = imageSynthesis.distortionMaterial;
+            material.SetFloat("_ZoomPercent", zoomPercent);
+            material.SetFloat("_k1", k1);
+            material.SetFloat("_k2", k2);
+            material.SetFloat("_k3", k3);
+            material.SetFloat("_k4", k4);
+            material.SetFloat("_DistortionIntensityX", intensityX);
+            material.SetFloat("_DistortionIntensityY", intensityY);
+            material.SetFloat("_LensDistortionStrength", strength);
+        }
         return ActionFinished.Success;
 
     }

unity/Assets/Scripts/ImageSynthesis/ImageSynthesis.cs

@@ -3,6 +3,7 @@
 using System.Collections.Generic;
 using System.Collections.Specialized;
 using System.IO;
+using System.Linq;
 using UnityEngine;
 using UnityEngine.Rendering;
 
@@ -489,6 +490,22 @@ public void OnCameraChange() {
         if (!distortionMaterial || distortionMaterial.shader != distortionShader) {
             distortionMaterial = new Material(distortionShader);
         }
+        var distortionPass = capturePasses.First(x => x.name == "_distortion");
+
+        distortionMaterial.SetFloat("_fov_y", distortionPass.camera.fieldOfView);
+
+        Texture2D realTex = null;
+        byte[] fileData;
+        var filePath =  Application.dataPath + "/real_camera/" + "frame_1.png";
+
+        if (File.Exists(filePath)) 	{
+            fileData = File.ReadAllBytes(filePath);
+            realTex = new Texture2D(2, 2);
+            realTex.LoadImage(fileData); //..this will auto-resize the texture dimensions.
+
+             distortionMaterial.SetTexture("_RealImage", realTex);
+        }
+
 
         // capturePasses [1].camera.farClipPlane = 100;
         // SetupCameraWithReplacementShader(capturePasses[1].camera, uberReplacementShader, ReplacelementModes.DepthMultichannel);

unity/Assets/Scripts/ImageSynthesis/Shaders/BarrelDistortion.shader

@@ -12,8 +12,8 @@ Shader "Custom/BarrelDistortion" {
 
           _k1 ("K1 polynomial dist coeff", Range (-8.0, 8.0))  = -0.126
           _k2 ("K2 polynomial dist coeff", Range (-8.0, 8.0))  = 0.004
-          _k3 ("K3 polynomial dist coeff", Range (-18.0, 18.0))  = 0.0
-          _k4 ("K4 polynomial dist coeff", Range (-18.0, 18.0))  = 0.0
+          _k3 ("K3 polynomial dist coeff", Range (-25.0, 25.0))  = 0.0
+          _k4 ("K4 polynomial dist coeff", Range (-25.0, 25.0))  = 0.0
 
          _DistortionIntensityX ("Distort Strength X", Range (0.0, 6.0)) = 1.0
          _DistortionIntensityY ("Distort Strength Y", Range (0.0, 6.0)) = 1.0
@@ -29,9 +29,10 @@ Shader "Custom/BarrelDistortion" {
 
               #pragma vertex vert
              #pragma fragment frag
+             #pragma target 3.0 
              #include "UnityCG.cginc"
 
-              uniform sampler2D _MainTex;
+            uniform sampler2D _MainTex;
              uniform sampler2D _CameraDepthTexture;
              uniform half4 _MainTex_TexelSize;
 
@@ -49,6 +50,8 @@ Shader "Custom/BarrelDistortion" {
              uniform float _k3;
              uniform float _k4;
 
+            //  uniform float4 _ScreenParams;
+
 
             //  uniform float4 _ScreenParams;
             //  uniform float4 _ProjectionParams;
@@ -63,12 +66,14 @@ Shader "Custom/BarrelDistortion" {
              {
                  float4 pos : SV_POSITION;
                  half2 uv : TEXCOORD0;
+                //  float2 clipPos : UNITY_VPOS_TYPE;
              };
 
 
               output vert(input i)
              {
                  output o;
+                //  o.screenPos = i.pos;
                  o.pos = UnityObjectToClipPos(i.pos);
                 //  o.uv = i.uv;
                  o.uv = MultiplyUV(UNITY_MATRIX_TEXTURE0, i.uv);
@@ -86,10 +91,14 @@ Shader "Custom/BarrelDistortion" {
 
               fixed4 frag(output o) : COLOR
              {
+
+                // return float4(o.screenPos.x, 0.0, o.screenPos.y, 1.0);
                 float effect = _LensDistortionStrength;
                 float2 distortionStrengthXY = float2(_DistortionIntensityX, _DistortionIntensityY);
                 float zoom_offset = (1.0 - _ZoomPercent) / 2.0;
 
+                float3x3 k = cam_intrinsics(_fov_y, _ScreenParams.y, _ScreenParams.x);
+
                 float2 centered_uv = o.uv - float2(0.5, 0.5);
                 centered_uv = o.uv*2.0 -  float2(1.0, 1.0);
                 //  centered_uv = o.uv -  float2(0.5, 0.5);
@@ -135,12 +144,14 @@ Shader "Custom/BarrelDistortion" {
                 // for xydistortion
                 // float2 uvDistorted = distort_uv*2.0 + float2(0.5, 0.5);
 
+
+
             fixed4 col = tex2D(_MainTex, uvDistorted);
             // Handle out of bound uv
             if (uvDistorted.x < 0 || uvDistorted.x > 1 || uvDistorted.y < 0 || uvDistorted.y > 1) {
                 return _OutOfBoundColour;//uv out of bound so display out of bound color
             } else {
-                return col;
+                return col
             }
 
             // FOR atan radius cutoff