From f973275fa7f0bdcf553cba8b0a09d4f987b75ea8 Mon Sep 17 00:00:00 2001 From: AlvaroHG Date: Tue, 15 Oct 2024 15:25:13 -0700 Subject: [PATCH] Third party image distortion --- ai2thor/interact.py | 17 ++ ai2thor/server.py | 1 + test_distortion.py | 65 +++++- unity/Assets/Scripts/AgentManager.cs | 28 +-- .../Scripts/ImageSynthesis/ImageSynthesis.cs | 17 ++ .../Shaders/BarrelDistortion.shader | 19 +- .../Shaders/BarrelDistortionImageBlend.shader | 194 ++++++++++++++++++ 7 files changed, 314 insertions(+), 27 deletions(-) create mode 100644 unity/Assets/Scripts/ImageSynthesis/Shaders/BarrelDistortionImageBlend.shader diff --git a/ai2thor/interact.py b/ai2thor/interact.py index b0bc10ef48..2adde955e5 100644 --- a/ai2thor/interact.py +++ b/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, diff --git a/ai2thor/server.py b/ai2thor/server.py index 7f03bc6bb1..5796d19451 100644 --- a/ai2thor/server.py +++ b/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 = [] diff --git a/test_distortion.py b/test_distortion.py index f278ab70b8..3d56819d61 100644 --- a/test_distortion.py +++ b/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() diff --git a/unity/Assets/Scripts/AgentManager.cs b/unity/Assets/Scripts/AgentManager.cs index 111c8ece7f..7e968f26b0 100644 --- a/unity/Assets/Scripts/AgentManager.cs +++ b/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 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 imageSynths = mainCamera ? new List() {this.primaryAgent.imageSynthesis} : new List(); + imageSynths = thidPartyCameraIndices != null ? imageSynths.Concat(this.thirdPartyCameras.Where((cam, i) => thidPartyCameraIndices.Contains(i)).Select(cam => cam.GetComponent())) : 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; } diff --git a/unity/Assets/Scripts/ImageSynthesis/ImageSynthesis.cs b/unity/Assets/Scripts/ImageSynthesis/ImageSynthesis.cs index 27a5a48a5d..b2eb8ff72b 100644 --- a/unity/Assets/Scripts/ImageSynthesis/ImageSynthesis.cs +++ b/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); diff --git a/unity/Assets/Scripts/ImageSynthesis/Shaders/BarrelDistortion.shader b/unity/Assets/Scripts/ImageSynthesis/Shaders/BarrelDistortion.shader index d1285438a6..580fc37eb6 100644 --- a/unity/Assets/Scripts/ImageSynthesis/Shaders/BarrelDistortion.shader +++ b/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 diff --git a/unity/Assets/Scripts/ImageSynthesis/Shaders/BarrelDistortionImageBlend.shader b/unity/Assets/Scripts/ImageSynthesis/Shaders/BarrelDistortionImageBlend.shader new file mode 100644 index 0000000000..6ab756a7b1 --- /dev/null +++ b/unity/Assets/Scripts/ImageSynthesis/Shaders/BarrelDistortionImageBlend.shader @@ -0,0 +1,194 @@ +// latest + +Shader "Custom/BarrelDistortion" { + Properties + { + [MainTexture] _MainTex ("Base (RGB)", 2D) = "white" {} + + _LensDistortionStrength ("Lens Distortion Strength", Range (-20.0, 20.0)) = 1.0 + // _LensDistortionTightness ("Lens Distortion Power", Range (-20.0, 20.0)) = 7.0 + _ZoomPercent ("Zoom Percent", Range (0.0, 5.0)) = 1.0 + + + _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 (-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 + + + _OutOfBoundColour ("Outline Color", Color) = (0.0, 0.0, 0.0, 1.0) + + _fov_y ("Vertical Fov", Range (-360.0, 360.0)) = 0.0 + + _Virtual_To_Real_Blend ("Blend From Virtual To Real", Range (0.0, 1.0)) = 0.0 + + _RealImage ("Real Image To Bleand", 2D) = "white" {} + } + SubShader + { + Pass + { + CGPROGRAM + + #pragma vertex vert + #pragma fragment frag + #pragma target 3.0 + #include "UnityCG.cginc" + + uniform sampler2D _MainTex; + uniform sampler2D _RealImage; + uniform sampler2D _CameraDepthTexture; + uniform half4 _MainTex_TexelSize; + + uniform float _LensDistortionStrength; + uniform float4 _OutOfBoundColour; + // uniform float _LensDistortionTightness; + + uniform float _ZoomPercent; + + uniform float _DistortionIntensityX; + uniform float _DistortionIntensityY; + + uniform float _k1; + uniform float _k2; + uniform float _k3; + uniform float _k4; + + uniform float _fov_y; + + uniform float _Virtual_To_Real_Blend; + + // uniform float4 _ScreenParams; + + + // uniform float4 _ScreenParams; + // uniform float4 _ProjectionParams; + + struct input + { + float4 pos : POSITION; + half2 uv : TEXCOORD0; + }; + + struct output + { + 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); + // o.uv.y = 1 - o.uv.y; + // // why do we need this? cause sometimes the image I get is flipped. see: http://docs.unity3d.com/Manual/SL-PlatformDifferences.html + // #if UNITY_UV_STARTS_AT_TOP + // if (_MainTex_TexelSize.y < 0) + // o.uv.y = 1 - o.uv.y; + // #endif + // if (_ProjectionParams.x >= 0) + // o.uv.y = 1 - o.uv.y; + + return o; + } + + float3x3 cam_intrinsics(float fov_y, float frame_height, float frame_width) { + float focal_length = 0.5 * frame_height / tan((UNITY_PI / 180.0) *(fov_y / 2)); + float f_x = focal_length; + float f_y = f_x; + + float c_x = frame_width / 2; + float c_y = frame_height / 2; + return float3x3(float3(f_x, 0, c_x), float3(0, f_y, c_y), float3(0, 0, 1)); + } + + 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); + + zoom_offset = (1.0 - _ZoomPercent); + centered_uv = (o.uv*2.0 * _ZoomPercent - float2(1.0, 1.0)) + float2(zoom_offset, zoom_offset); + + + + float uv_dot = dot(centered_uv, centered_uv); + float r = sqrt(uv_dot); + // // For Atan based distortion + float z = sqrt(1.0 - uv_dot * effect); + // float z = sqrt(1.0 - r * r); + float atan_distort = atan2(r, z) / UNITY_PI; + + float atan_distortX = atan2(r, sqrt(1.0 - uv_dot * _DistortionIntensityX)) / UNITY_PI; + float atan_distortY = atan2(r, sqrt(1.0 - uv_dot * _DistortionIntensityY)) / UNITY_PI; + float2 distort_uv = (centered_uv / r) * atan_distort;//+ float2(0.5, 0.5); + + distort_uv = (centered_uv / r) * float2(atan_distortX, atan_distortY); + + //return t; + + const float distortionMagnitude=abs(centered_uv.x*centered_uv.y); + + float smoothDistortionMagnitude = 1.0 + _k1 * pow(r, 2.0) + _k2 * pow(r, 4.0) + _k3 * pow(r,6.0) + _k4 * pow(r,8.0); + + float zoom_percent = 0.2; + float translate = (1.0 - zoom_percent) / 2.0; + float2 centered_uv_norm = normalize(centered_uv); + + + + // Works with zoom for poly distortion + float2 uvDistorted = centered_uv * smoothDistortionMagnitude * distortionStrengthXY / 2.0 + float2(0.5, 0.5) ; + + // working for atan distortion + // float2 uvDistorted = ( centered_uv / r ) * atan_distort + float2(0.5, 0.5); + + // float2 uvDistorted = distort_uv + float2(0.5, 0.5); + + // 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 lerp(col, tex2D(_RealImage, o.uv), _Virtual_To_Real_Blend); + } + + // FOR atan radius cutoff + // float radius = 1.0; + // if (uvDistorted.x > radius) { + // return _OutOfBoundColour; + // } + // else { + // return lerp(col1, col2, _BlendAmount); + // } + + + } + + ENDCG + } + } + } + \ No newline at end of file