diff --git a/ai2thor/fifo_server.py b/ai2thor/fifo_server.py index f3afcd44bd..52725fc023 100644 --- a/ai2thor/fifo_server.py +++ b/ai2thor/fifo_server.py @@ -42,6 +42,8 @@ class FieldType(IntEnum): THIRD_PARTY_IMAGE_IDS = 14 THIRD_PARTY_CLASSES = 15 THIRD_PARTY_FLOW = 16 + DISTORTION_IMAGE = 17, + THIRD_PARTY_DISTORTION=18, END_OF_MESSAGE = 255 @@ -83,6 +85,7 @@ def __init__( # allows us to map the enum to form field names # for backwards compatibility # this can be removed when the wsgi server is removed + # TODO make more generic cod ewise, one new value you have to modify 3 different arrays self.form_field_map = { FieldType.RGB_IMAGE: "image", FieldType.DEPTH_IMAGE: "image_depth", @@ -90,12 +93,14 @@ def __init__( FieldType.IDS_IMAGE: "image_ids", FieldType.NORMALS_IMAGE: "image_normals", FieldType.FLOWS_IMAGE: "image_flow", + FieldType.DISTORTION_IMAGE: "image_distortion", FieldType.THIRD_PARTY_IMAGE: "image-thirdParty-camera", FieldType.THIRD_PARTY_DEPTH: "image_thirdParty_depth", FieldType.THIRD_PARTY_NORMALS: "image_thirdParty_normals", FieldType.THIRD_PARTY_IMAGE_IDS: "image_thirdParty_image_ids", FieldType.THIRD_PARTY_CLASSES: "image_thirdParty_classes", FieldType.THIRD_PARTY_FLOW: "image_thirdParty_flow", + FieldType.THIRD_PARTY_DISTORTION: "image_thirdParty_distortion", } self.image_fields = { @@ -111,6 +116,8 @@ def __init__( FieldType.THIRD_PARTY_IMAGE_IDS, FieldType.THIRD_PARTY_CLASSES, FieldType.THIRD_PARTY_FLOW, + FieldType.DISTORTION_IMAGE, + FieldType.THIRD_PARTY_DISTORTION } self.eom_header = self._create_header(FieldType.END_OF_MESSAGE, b"") diff --git a/ai2thor/interact.py b/ai2thor/interact.py index 341b39e894..1a17e8eb90 100644 --- a/ai2thor/interact.py +++ b/ai2thor/interact.py @@ -224,10 +224,11 @@ def write_image( image_dir, suffix, image_per_frame=False, + color_frame=False, semantic_segmentation_frame=False, instance_segmentation_frame=False, depth_frame=False, - color_frame=False, + distortion_frame=False, metadata=False, ): def save_image(name, image, flip_br=False): @@ -271,6 +272,13 @@ def json_write(name, obj): array_to_image, save_image, ), + ( + "distortion", + semantic_segmentation_frame, + lambda event: event.distortion_frame, + array_to_image, + lambda x, y: save_image(x, y, flip_br=True), + ), ( "depth", depth_frame, diff --git a/ai2thor/server.py b/ai2thor/server.py index 047e5cecdd..7f03bc6bb1 100644 --- a/ai2thor/server.py +++ b/ai2thor/server.py @@ -646,6 +646,17 @@ def add_image_flows(self, image_flows_data): image_flows_data, self.screen_width, self.screen_height ) + # TODO make more generic, above methids minus depth are basically the same + def add_image_distortion(self, image_distortion_data): + self.distortion_frame = read_buffer_image( + image_distortion_data, self.screen_width, self.screen_height + ) + + def add_third_party_image_distortion(self, distortion_data): + self.third_party_distortion_frames.append( + read_buffer_image(distortion_data, self.screen_width, self.screen_height) + ) + def add_third_party_camera_image(self, third_party_image_data): self.third_party_camera_frames.append( read_buffer_image(third_party_image_data, self.screen_width, self.screen_height) @@ -775,6 +786,7 @@ def create_event(self, metadata, files): image_classes=e.add_image_classes, image_normals=e.add_image_normals, image_flow=e.add_image_flows, + image_distortion=e.add_image_distortion ) for key in image_mapping.keys(): @@ -796,6 +808,7 @@ def create_event(self, metadata, files): "image_thirdParty_classes": e.add_third_party_image_classes, "image_thirdParty_normals": e.add_third_party_image_normals, "image_thirdParty_flows": e.add_third_party_image_flows, + "image_thirdParty_distortion": e.add_third_party_image_distortion } if a["thirdPartyCameras"] is not None: diff --git a/test_distortion.py b/test_distortion.py new file mode 100644 index 0000000000..e5539d777a --- /dev/null +++ b/test_distortion.py @@ -0,0 +1,116 @@ +import json +import ai2thor.controller +import argparse +import sys +from ai2thor.interact import InteractiveControllerPrompt + +def load_scene(scene_name, house_path=None, run_in_editor=False, platform=None, local_build=False, commit_id=None): + if not run_in_editor: + args = dict( + # commit_id="13ef2aba9d0228c30d775cbae0b674f0826a97f2", + # commit_id="b9a854e9b7fdde450460fe054fb75ba6572e022e", + commit_id=commit_id, + server_class=ai2thor.fifo_server.FifoServer, + ) + if local_build: + del args["commit_id"] + args["local_build"] = True + else: + args = dict( + start_unity=False, + port=8200, + server_class=ai2thor.wsgi_server.WsgiServer, + ) + + all_args = dict( + # local_executable_path="unity/builds/thor-OSXIntel64-local/thor-OSXIntel64-local.app/Contents/MacOS/AI2-THOR", + # local_build=True, + # agentMode="stretch", + platform=platform, + scene=scene_name, + gridSize=0.25, + width=300, + height=300, + visibilityScheme="Distance", + renderDepthImage=True, + renderDistortionImage=True, + renderSemanticSegmentation=True, + renderInstanceSegmentation=True, + **args, + ) + + print("Controller args: ") + print(all_args) + + controller = ai2thor.controller.Controller(**all_args) + if house_path != None: + with open(house_path, "r") as f: + house = json.load(f) + + evt = controller.step(action="CreateHouse", house=house) + + print(f"Action {controller.last_action['action']} success: {evt.metadata['lastActionSuccess']}") + print(f'Error: {evt.metadata["errorMessage"]}') + + agent = house["metadata"]["agent"] + evt = controller.step( + action="TeleportFull", + x=agent["position"]["x"], + y=agent["position"]["y"], + z=agent["position"]["z"], + rotation=agent["rotation"], + horizon=agent["horizon"], + standing=agent["standing"], + ) + print(f"Action {controller.last_action['action']} success: {evt.metadata['lastActionSuccess']}") + print(f'Error: {evt.metadata["errorMessage"]}') + + 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 + ) + InteractiveControllerPrompt.write_image(controller.last_event, "distortion", "", semantic_segmentation_frame=True, depth_frame=True, color_frame=True, distortion_frame=True) + # input() + + +if __name__ == "__main__": + parser = argparse.ArgumentParser() + + parser.add_argument("--house_path", type=str, default=None) + parser.add_argument( + "--platform", + type=str, + default=None, + help='Platform "CloudRendering", "OSXIntel64"', + ) + parser.add_argument( + "--commit_id", + type=str, + default=None, + ) + + parser.add_argument( + "--scene", + type=str, + default="FloorPlan227_physics", + ) + + parser.add_argument("--local_build", action="store_true", help="Uses the local build.") + + parser.add_argument("--editor", action="store_true", help="Runs in editor.") + + args = parser.parse_args(sys.argv[1:]) + load_scene( + scene_name=args.scene, + house_path=args.house_path, + run_in_editor=args.editor, + local_build=args.local_build, + commit_id=args.commit_id, + platform=args.platform, + ) # platform="CloudRendering") diff --git a/unity/Assets/Editor/Build.cs b/unity/Assets/Editor/Build.cs index 388b3c377f..20d9832675 100644 --- a/unity/Assets/Editor/Build.cs +++ b/unity/Assets/Editor/Build.cs @@ -181,8 +181,8 @@ private static List GetAllScenePaths() } } - return scenes; - // return scenes.Where(x => x.Contains("FloorPlan1_") || x.Contains("Procedural") || x.Contains("ProceduralAB")).ToList(); + // return scenes; + return scenes.Where(x => x.Contains("FloorPlan1_") || x.Contains("Procedural") || x.Contains("ProceduralAB")).ToList(); // uncomment for faster builds for testing return scenes; //.Where(x => x.Contains("Procedural.unity") || x.Contains("Procedural.unity")).ToList(); //.Where(x => x.Contains("FloorPlan1_") || x.Contains("Procedural")).ToList(); } diff --git a/unity/Assets/Scripts/AgentManager.cs b/unity/Assets/Scripts/AgentManager.cs index 04e4c9f90e..be6b1a2b10 100644 --- a/unity/Assets/Scripts/AgentManager.cs +++ b/unity/Assets/Scripts/AgentManager.cs @@ -1374,13 +1374,6 @@ bool shouldRenderImageSynthesis "_flow", "image_thirdParty_flow" ); // XXX fix this in a bit - addImageSynthesisImage( - renderPayload, - imageSynthesis, - this.renderSemanticSegmentation, - "_distortion", - "image_thirdParty_distortion" - ); addImageSynthesisImage( renderPayload, imageSynthesis, @@ -1863,7 +1856,7 @@ public void SetCriticalErrorState() { this.agentManagerState = AgentState.Error; } - public ActionFinished SetDistortionShaderParams(float zoomPercent, float k1, float k2, float k3, float k4, float strength = 1.0f) { + 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) { return new ActionFinished(success: false, errorMessage: "No imageSynthesis, make sure you pass 'renderDistortionImage = true' to the agent constructor."); @@ -1874,6 +1867,8 @@ public ActionFinished SetDistortionShaderParams(float zoomPercent, float k1, flo 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/FifoServer.cs b/unity/Assets/Scripts/FifoServer.cs index 86dd645258..15a6c65afe 100644 --- a/unity/Assets/Scripts/FifoServer.cs +++ b/unity/Assets/Scripts/FifoServer.cs @@ -23,12 +23,14 @@ public class Client { { "image_classes", FieldType.ClassesImage }, { "image_flow", FieldType.FlowsImage }, { "image_ids", FieldType.IDsImage }, + { "image_distortion", FieldType.DistortionImage }, { "image-thirdParty-camera", FieldType.ThirdPartyCameraImage }, { "image_thirdParty_depth", FieldType.ThirdPartyDepth }, { "image_thirdParty_normals", FieldType.ThirdPartyNormals }, { "image_thirdParty_classes", FieldType.ThirdPartyClasses }, { "image_thirdParty_image_ids", FieldType.ThirdPartyImageIds }, - { "image_thirdParty_flow", FieldType.ThirdPartyFlow } + { "image_thirdParty_flow", FieldType.ThirdPartyFlow }, + { "image_thirdParty_distortion", FieldType.ThirdPartyDistortion }, }; public static int UnpackNetworkBytes(byte[] data, int offset = 0) { @@ -142,6 +144,10 @@ public enum FieldType : byte { ThirdPartyImageIds = 0x0e, ThirdPartyClasses = 0x0f, ThirdPartyFlow = 0x10, + + DistortionImage = 0x11, + ThirdPartyDistortion = 0x12, + EndOfMessage = 0xff } } diff --git a/unity/Assets/Scripts/ImageSynthesis/ImageSynthesis.cs b/unity/Assets/Scripts/ImageSynthesis/ImageSynthesis.cs index 9037a4c81e..9597833355 100644 --- a/unity/Assets/Scripts/ImageSynthesis/ImageSynthesis.cs +++ b/unity/Assets/Scripts/ImageSynthesis/ImageSynthesis.cs @@ -300,7 +300,7 @@ public void OnEnable() { #endif if (!distortionShader) { - distortionShader = Shader.Find("Hidden/BarrelDistortion"); + distortionShader = Shader.Find("Custom/BarrelDistortion"); } // if (!positionShader) diff --git a/unity/Assets/Scripts/ImageSynthesis/Shaders/BarrelDistortion.shader b/unity/Assets/Scripts/ImageSynthesis/Shaders/BarrelDistortion.shader index adb57acecf..d1285438a6 100644 --- a/unity/Assets/Scripts/ImageSynthesis/Shaders/BarrelDistortion.shader +++ b/unity/Assets/Scripts/ImageSynthesis/Shaders/BarrelDistortion.shader @@ -1,6 +1,6 @@ // latest -Shader "Hidden/BarrelDistortion" { +Shader "Custom/BarrelDistortion" { Properties { [MainTexture] _MainTex ("Base (RGB)", 2D) = "white" {} @@ -15,7 +15,7 @@ Shader "Hidden/BarrelDistortion" { _k3 ("K3 polynomial dist coeff", Range (-18.0, 18.0)) = 0.0 _k4 ("K4 polynomial dist coeff", Range (-18.0, 18.0)) = 0.0 - _DistortionIntensityX ("Distort Strength X", Range (0.0, 6.0)) = 1.0 + _DistortionIntensityX ("Distort Strength X", Range (0.0, 6.0)) = 1.0 _DistortionIntensityY ("Distort Strength Y", Range (0.0, 6.0)) = 1.0 diff --git a/unity/ProjectSettings/GraphicsSettings.asset b/unity/ProjectSettings/GraphicsSettings.asset index 0b74deae17..95cf203587 100644 --- a/unity/ProjectSettings/GraphicsSettings.asset +++ b/unity/ProjectSettings/GraphicsSettings.asset @@ -28,6 +28,7 @@ GraphicsSettings: m_LensFlare: m_Mode: 1 m_Shader: {fileID: 102, guid: 0000000000000000f000000000000000, type: 0} + m_VideoShadersIncludeMode: 2 m_AlwaysIncludedShaders: - {fileID: 7, guid: 0000000000000000f000000000000000, type: 0} - {fileID: 15104, guid: 0000000000000000f000000000000000, type: 0} @@ -46,6 +47,7 @@ GraphicsSettings: - {fileID: 4800000, guid: 4cf408482db174fc4b862238d304ef73, type: 3} - {fileID: 16001, guid: 0000000000000000f000000000000000, type: 0} - {fileID: 16002, guid: 0000000000000000f000000000000000, type: 0} + - {fileID: 4800000, guid: 488c87645652f4b30a1a5f7d3d6a30b1, type: 3} m_PreloadedShaders: [] m_SpritesDefaultMaterial: {fileID: 10754, guid: 0000000000000000f000000000000000, type: 0} @@ -172,5 +174,5 @@ GraphicsSettings: m_AlbedoSwatchInfos: [] m_LightsUseLinearIntensity: 0 m_LightsUseColorTemperature: 0 + m_DefaultRenderingLayerMask: 1 m_LogWhenShaderIsCompiled: 0 - m_AllowEnlightenSupportForUpgradedProject: 1 diff --git a/unity/ProjectSettings/ProjectSettings.asset b/unity/ProjectSettings/ProjectSettings.asset index 29c345dbd9..6a9ed65b4e 100644 --- a/unity/ProjectSettings/ProjectSettings.asset +++ b/unity/ProjectSettings/ProjectSettings.asset @@ -644,7 +644,7 @@ PlayerSettings: webGLThreadsSupport: 0 webGLDecompressionFallback: 0 scriptingDefineSymbols: - 1: UNITY_POST_PROCESSING_STACK_V2 + 1: 4: CROSS_PLATFORM_INPUT;MOBILE_INPUT;UNITY_POST_PROCESSING_STACK_V2 7: CROSS_PLATFORM_INPUT;MOBILE_INPUT;UNITY_POST_PROCESSING_STACK_V2 13: UNITY_POST_PROCESSING_STACK_V2