working distortion

ai2thor/interact.py

@@ -274,7 +274,7 @@ def json_write(name, obj):
             ),
             (
                 "distortion",
-                semantic_segmentation_frame,
+                distortion_frame,
                 lambda event: event.distortion_frame,
                 array_to_image,
                 lambda x, y: save_image(x, y, flip_br=True),
@@ -308,19 +308,20 @@ def json_write(name, obj):
         ]
 
         for frame_filename, condition, frame_func, transform, save in frame_writes:
-            frame = frame_func(event)
-            if frame is not None and condition:
-                frame = transform(frame)
-                image_name = os.path.join(
-                    image_dir,
-                    "{}{}".format(frame_filename, "{}".format(suffix) if image_per_frame else ""),
-                )
-                print("Image {}, {}".format(image_name, image_dir))
-                save(image_name, frame)
-
-            elif condition:
-                print(
-                    "No frame '{}' present, call initialize with the right parameters".format(
-                        frame_filename
+            if condition:
+                frame = frame_func(event)
+                if frame is not None:
+                    frame = frame_func(event)
+                    frame = transform(frame)
+                    image_name = os.path.join(
+                        image_dir,
+                        "{}{}".format(frame_filename, "{}".format(suffix) if image_per_frame else ""),
+                    )
+                    print("Image {}, {}".format(image_name, image_dir))
+                    save(image_name, frame)
+                else:
+                    print(
+                        "No frame '{}' present, call initialize with the right parameters".format(
+                            frame_filename
+                        )
                     )
-                )

test_distortion.py

@@ -2,9 +2,13 @@
 import ai2thor.controller
 import argparse
 import sys
+import os
 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):
+
+def load_scene(scene_name, house_path=None, run_in_editor=False, platform=None, local_build=False, commit_id=None, fov=120, distortion=False, image_dir=None):
+    if image_dir is not None:
+        os.makedirs(image_dir, exist_ok=True)
     if not run_in_editor:
         args = dict(
             # commit_id="13ef2aba9d0228c30d775cbae0b674f0826a97f2",
@@ -33,7 +37,7 @@ def load_scene(scene_name, house_path=None, run_in_editor=False, platform=None,
         height=300,
         visibilityScheme="Distance",
         renderDepthImage=True,
-        renderDistortionImage=True,
+        renderDistortionImage=distortion,
         renderSemanticSegmentation=True,
         renderInstanceSegmentation=True,
         **args,
@@ -64,18 +68,42 @@ 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"]}')
+    if distortion:
+        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
+    xpos = dict(x=0.0, y=0.900992214679718, z=0.0786)
+    sr = controller.step(
+        action="Teleport", position=xpos, rotation=dict(x=0, y=0, z=0), forceAction=True
     )
-    InteractiveControllerPrompt.write_image(controller.last_event, "distortion", "", semantic_segmentation_frame=True, depth_frame=True, color_frame=True, distortion_frame=True)
+    if not sr.metadata["lastActionSuccess"]:
+        print(f"Error teleporting to {xpos}")
+        current_pos = controller.last_event.metadata["agent"]["position"]
+        print(f"Current position: {current_pos}")
+    # print("Teleported to calibration room")
+    # print(f"Current position: {thor_controller.last_event.metadata['agent']['position']}")
+    # controller.step(
+    #     {"action": "RotateCameraMount", "degrees": 13, "secondary": False}
+    # )
+    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,
+        agentId=0,
+    )
+
+    InteractiveControllerPrompt.write_image(controller.last_event, image_dir, "", semantic_segmentation_frame=True, depth_frame=True, color_frame=True, distortion_frame=distortion)
     # input()
 
 
@@ -95,6 +123,12 @@ def load_scene(scene_name, house_path=None, run_in_editor=False, platform=None,
         default=None,
     )
 
+    parser.add_argument(
+        "--output",
+        type=str,
+        default=None,
+    )
+
     parser.add_argument(
         "--scene",
         type=str,
@@ -105,12 +139,26 @@ def load_scene(scene_name, house_path=None, run_in_editor=False, platform=None,
 
     parser.add_argument("--editor", action="store_true", help="Runs in editor.")
 
+
+    parser.add_argument("--distortion", action="store_true", help="Distortion.")
+
+    parser.add_argument(
+        "--fov",
+        type=str,
+        default="120",
+        help='fov',
+    )
+
     args = parser.parse_args(sys.argv[1:])
+    print(args.distortion)
     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,
+        fov=float(args.fov),
+        distortion=args.distortion,
+        image_dir=args.output
     )  # 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") || x.Contains("Calibration_Room")).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

@@ -1222,6 +1222,7 @@ private void addObjectImage(
         BaseFPSAgentController agent,
         ref MetadataWrapper metadata
     ) {
+        Debug.Log("Add object image enter");
         if (this.renderInstanceSegmentation || this.renderSemanticSegmentation) {
             Debug.Log($"imageSynthesis null {agent.ImageSynthesis == null}");
             if (!agent.ImageSynthesis.hasCapturePass("_id")) {
@@ -1233,6 +1234,7 @@ ref MetadataWrapper metadata
             payload.Add(new KeyValuePair<string, byte[]>("image_ids", bytes));
 
             List<ColorId> colors = new List<ColorId>();
+            Debug.Log($"Add object image befor color keys is null {agent.ImageSynthesis.colorIds == null}?");
             foreach (Color key in agent.ImageSynthesis.colorIds.Keys) {
                 ColorId cid = new ColorId();
                 cid.color = new ushort[]

unity/Assets/Scripts/ImageSynthesis/ImageSynthesis.cs

@@ -273,6 +273,7 @@ public void OnEnable() {
         // RenderCapture s;
         // This initialization code MUST live in OnEnable and not Start as we instantiate ThirdPartyCameras
         // programatically in other functions and need them to be initialized immediately.
+        Debug.Log("OnEnable image synth");
         if (!initialized) {
             // XXXXXXXXXXX************
             // Remember, adding any new Shaders requires them to be included in Project Settings->Graphics->Always Included Shaders
@@ -316,6 +317,7 @@ public void OnEnable() {
             md5 = System.Security.Cryptography.MD5.Create();
 
             OnCameraChange();
+             Debug.Log("OnEnable image synth scenechange");
             OnSceneChange();
         }
         initialized = true;

unity/ProjectSettings/GraphicsSettings.asset

@@ -47,7 +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}
+  - {fileID: 4800000, guid: 72c7b921c024344cdb3b252b6a4399b0, type: 3}
   m_PreloadedShaders: []
   m_SpritesDefaultMaterial: {fileID: 10754, guid: 0000000000000000f000000000000000,
     type: 0}