working distrotion get in python

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,19 +85,22 @@ 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",
             FieldType.CLASSES_IMAGE: "image_classes",
             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"")

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,

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:

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")

unity/Assets/Editor/Build.cs

@@ -181,8 +181,8 @@ private static List<string> 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();
     }

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;
 

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
     }
 }

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)

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
 
 

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

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