segmentation API

AirLib/include/api/RpcLibClientBase.hpp

@@ -33,6 +33,9 @@ class RpcLibClientBase {
     void enableApiControl(bool is_enabled);
     void reset();
 
+    bool simSetSegmentationObjectID(const std::string& mesh_name, int object_id, bool is_name_regex = false);
+    int simGetSegmentationObjectID(const std::string& mesh_name);
+
     virtual ~RpcLibClientBase();    //required for pimpl
 
 protected:

AirLib/include/api/VehicleApiBase.hpp

@@ -23,7 +23,7 @@ class VehicleApiBase {
     virtual void simSetPose(const Pose& pose, bool ignore_collison) = 0;
     virtual Pose simGetPose() = 0;
 
-    virtual void simSetSegmentationObjectID(const std::string& mesh_name, int object_id, bool is_name_regex = false) = 0;
+    virtual bool simSetSegmentationObjectID(const std::string& mesh_name, int object_id, bool is_name_regex = false) = 0;
     virtual int simGetSegmentationObjectID(const std::string& mesh_name) = 0;
 
     virtual ~VehicleApiBase() = default;

AirLib/include/controllers/VehicleConnectorBase.hpp

@@ -28,7 +28,7 @@ class VehicleConnectorBase : public UpdatableObject
     virtual VehicleCameraBase* getCamera(unsigned int index) = 0;
     virtual void setPose(const Pose& pose, bool ignore_collison) = 0;
     virtual Pose getPose() = 0;
-    virtual void setSegmentationObjectID(const std::string& mesh_name, int object_id,
+    virtual bool setSegmentationObjectID(const std::string& mesh_name, int object_id,
         bool is_name_regex = false) = 0;
     virtual int getSegmentationObjectID(const std::string& mesh_name) = 0;
 };

AirLib/include/vehicles/multirotor/api/DroneApi.hpp

@@ -178,10 +178,10 @@ class DroneApi : public VehicleApiBase {
         return vehicle_->getPose();
     }
 
-    virtual void simSetSegmentationObjectID(const std::string& mesh_name, int object_id,
+    virtual bool simSetSegmentationObjectID(const std::string& mesh_name, int object_id,
         bool is_name_regex = false) override
     {
-        vehicle_->setSegmentationObjectID(mesh_name, object_id, is_name_regex);
+        return vehicle_->setSegmentationObjectID(mesh_name, object_id, is_name_regex);
     }
 
     virtual int simGetSegmentationObjectID(const std::string& mesh_name) override

AirLib/src/api/RpcLibClientBase.cpp

@@ -65,6 +65,14 @@ RpcLibClientBase::ConnectionState RpcLibClientBase::getConnectionState()
         return ConnectionState::Unknown;
     }
 }
+bool RpcLibClientBase::simSetSegmentationObjectID(const std::string& mesh_name, int object_id, bool is_name_regex)
+{
+    return pimpl_->client.call("simSetSegmentationObjectID", mesh_name, object_id, is_name_regex).as<bool>();
+}
+int RpcLibClientBase::simGetSegmentationObjectID(const std::string& mesh_name)
+{
+    return pimpl_->client.call("simGetSegmentationObjectID", mesh_name).as<int>();
+}
 void RpcLibClientBase::enableApiControl(bool is_enabled)
 {
     pimpl_->client.call("enableApiControl", is_enabled);

AirLib/src/api/RpcLibServerBase.cpp

@@ -78,6 +78,15 @@ RpcLibServerBase::RpcLibServerBase(VehicleApiBase* vehicle, string server_addres
             RpcLibAdapatorsBase::Pose { return vehicle_->simGetPose();
     });
 
+    pimpl_->server.
+        bind("simSetSegmentationObjectID", [&](const std::string& mesh_name, int object_id, bool is_name_regex) -> bool {
+        return vehicle_->simSetSegmentationObjectID(mesh_name, object_id, is_name_regex);
+    });
+    pimpl_->server.
+        bind("simGetSegmentationObjectID", [&](const std::string& mesh_name) -> int {
+        return vehicle_->simGetSegmentationObjectID(mesh_name);
+    });    
+
     pimpl_->server.bind("reset", [&]() -> void {
         vehicle_->reset();
     });

PythonClient/AirSimClient.py

@@ -172,6 +172,11 @@ def enableApiControl(self, is_enabled):
     def isApiControlEnabled(self):
         return self.client.call('isApiControlEnabled')
 
+    def simSetSegmentationObjectID(self, mesh_name, object_id, is_name_regex = False):
+        return self.client.call('simSetSegmentationObjectID', mesh_name, object_id, is_name_regex)
+    def simGetSegmentationObjectID(self, mesh_name):
+        return self.client.call('simGetSegmentationObjectID', mesh_name)
+
     # camera control
     # simGetImage returns compressed png in array of bytes
     # image_type uses one of the AirSimImageType members

PythonClient/PythonClient.pyproj

@@ -5,7 +5,7 @@
     <SchemaVersion>2.0</SchemaVersion>
     <ProjectGuid>e2049e20-b6dd-474e-8bca-1c8dc54725aa</ProjectGuid>
     <ProjectHome>.</ProjectHome>
-    <StartupFile>hello_drone.py</StartupFile>
+    <StartupFile>segmentation.py</StartupFile>
     <SearchPath>
     </SearchPath>
     <WorkingDirectory>.</WorkingDirectory>
@@ -49,6 +49,9 @@
     </Compile>
     <Compile Include="pfm.py" />
     <Compile Include="AirSimClient.py" />
+    <Compile Include="segmentation.py">
+      <SubType>Code</SubType>
+    </Compile>
   </ItemGroup>
   <PropertyGroup>
     <VisualStudioVersion Condition="'$(VisualStudioVersion)' == ''">10.0</VisualStudioVersion>

PythonClient/segmentation.py

@@ -0,0 +1,31 @@
+# In settings.json first activate computer vision mode: 
+# https://github.com/Microsoft/AirSim/blob/master/docs/image_apis.md#computer-vision-mode
+
+from AirSimClient import *
+
+client = CarClient()
+client.confirmConnection()
+
+found = client.simSetSegmentationObjectID("Ground", 20);
+print("Done: %r" % (found))
+
+found = client.simSetSegmentationObjectID("Ground_2", 22);
+print("Done: %r" % (found))
+
+found = client.simSetSegmentationObjectID("Ground_3", 23);
+print("Done: %r" % (found))
+
+found = client.simSetSegmentationObjectID("Ground_4", 24);
+print("Done: %r" % (found))
+
+found = client.simSetSegmentationObjectID("Ground_5", 25);
+print("Done: %r" % (found))
+
+found = client.simSetSegmentationObjectID("Ground_6", 26);
+print("Done: %r" % (found))
+
+found = client.simSetSegmentationObjectID("Ground_7", 27);
+print("Done: %r" % (found))
+
+
+

Unreal/Plugins/AirSim/Source/AirBlueprintLib.cpp

@@ -7,6 +7,7 @@
 #include "Kismet/GameplayStatics.h"
 #include "GameFramework/RotatingMovementComponent.h"
 #include <exception>
+#include <regex>
 #include "common/common_utils/Utils.hpp"
 #include "Components/StaticMeshComponent.h"
 #include "EngineUtils.h"
@@ -126,39 +127,66 @@ void UAirBlueprintLib::FindAllActor(const UObject* context, TArray<AActor*>& fou
     UGameplayStatics::GetAllActorsOfClass(context == nullptr ? GEngine : context, T::StaticClass(), foundActors);
 }
 
+std::string UAirBlueprintLib::GetMeshName(UMeshComponent* mesh)
+{
+    return std::string( TCHAR_TO_UTF8(*(mesh->GetOwner() ? mesh->GetOwner()->GetName() :
+        ""))); //UKismetSystemLibrary::GetDisplayName(mesh)
+}
+
 void UAirBlueprintLib::InitializeMeshStencilIDs()
 {
     for (TObjectIterator<UMeshComponent> comp; comp; ++comp)
     {
         UMeshComponent *mesh = *comp;
         mesh->SetRenderCustomDepth(true);
-        FString name = mesh->GetName();
+        std::string mesh_name = GetMeshName(mesh);
+        if (mesh_name == "")
+            continue;
+        FString name(mesh_name.c_str());
         int hash = 0;
         for (int idx = 0; idx < name.Len() && idx < 3; ++idx) {
             hash += UKismetStringLibrary::GetCharacterAsNumber(name, idx);
         }
-
-        mesh->CustomDepthStencilValue = hash % 256;
+        mesh->CustomDepthStencilValue = (hash+1) % 256;
         mesh->MarkRenderStateDirty();
     }
 }
 
-void UAirBlueprintLib::SetMeshStencilID(const std::string& mesh_name, int object_id,
+bool UAirBlueprintLib::SetMeshStencilID(const std::string& mesh_name, int object_id,
     bool is_name_regex)
 {
-    FString fmesh_name(mesh_name.c_str());
+    std::regex name_regex;
+
+    if (is_name_regex)
+        name_regex.assign(mesh_name);
+
+    int changes = 0;
     for (TObjectIterator<UMeshComponent> comp; comp; ++comp)
     {
-        // Access the subclass instance with the * or -> operators.
         UMeshComponent *mesh = *comp;
-        if (mesh->GetName() == fmesh_name) {
+
+        std::string comp_mesh_name = GetMeshName(mesh);
+        if (comp_mesh_name == "")
+            continue;
+
+        bool is_match = (!is_name_regex && (comp_mesh_name == mesh_name))
+            || (is_name_regex && std::regex_match(comp_mesh_name, name_regex));
+
+        if (is_match) {
+            ++changes;
+            mesh->SetRenderCustomDepth(false);
             mesh->SetRenderCustomDepth(true);
             mesh->CustomDepthStencilValue = object_id;
+            //mesh->SetVisibility(false);
+            //mesh->SetVisibility(true);
             mesh->MarkRenderStateDirty();
 
-            break;
+            if (! is_name_regex)
+                return true;
         }
     }
+
+    return changes > 0;
 }
 
 int UAirBlueprintLib::GetMeshStencilID(const std::string& mesh_name)

Unreal/Plugins/AirSim/Source/AirBlueprintLib.h

@@ -44,10 +44,11 @@ class UAirBlueprintLib : public UBlueprintFunctionLibrary
     static bool GetLastObstaclePosition(const AActor* actor, const FVector& start, const FVector& end, FHitResult& hit, const AActor* ignore_actor = nullptr, ECollisionChannel collison_channel = ECC_Visibility);
     static void FollowActor(AActor* follower, const AActor* followee, const FVector& offset, bool fixed_z = false, float fixed_z_val = 2.0f);
 
-    static void SetMeshStencilID(const std::string& mesh_name, int object_id,
+    static bool SetMeshStencilID(const std::string& mesh_name, int object_id,
         bool is_name_regex = false);
     static int GetMeshStencilID(const std::string& mesh_name);
     static void InitializeMeshStencilIDs();
+    static std::string GetMeshName(UMeshComponent* mesh);
 
     template<class UserClass>
     static FInputActionBinding& BindActionToKey(const FName action_name, const FKey in_key, UserClass* actor,

Unreal/Plugins/AirSim/Source/Car/CarPawn.cpp

@@ -55,10 +55,10 @@ class ACarPawn::CarApi : public msr::airlib::CarApiBase {
         return response;
     }
 
-    virtual void simSetSegmentationObjectID(const std::string& mesh_name, int object_id, 
+    virtual bool simSetSegmentationObjectID(const std::string& mesh_name, int object_id, 
         bool is_name_regex = false) override
     {
-        UAirBlueprintLib::SetMeshStencilID(mesh_name, object_id, is_name_regex);
+        return UAirBlueprintLib::SetMeshStencilID(mesh_name, object_id, is_name_regex);
     }
 
     virtual int simGetSegmentationObjectID(const std::string& mesh_name) override

Unreal/Plugins/AirSim/Source/Multirotor/MultiRotorConnector.cpp

@@ -235,10 +235,10 @@ Pose MultiRotorConnector::getPose()
     return vehicle_.getPose();
 }
 
-void MultiRotorConnector::setSegmentationObjectID(const std::string& mesh_name, int object_id,
+bool MultiRotorConnector::setSegmentationObjectID(const std::string& mesh_name, int object_id,
     bool is_name_regex)
 {
-    UAirBlueprintLib::SetMeshStencilID(mesh_name, object_id, is_name_regex);
+    return UAirBlueprintLib::SetMeshStencilID(mesh_name, object_id, is_name_regex);
 }
 
 int MultiRotorConnector::getSegmentationObjectID(const std::string& mesh_name)

Unreal/Plugins/AirSim/Source/Multirotor/MultiRotorConnector.h

@@ -54,7 +54,7 @@ class MultiRotorConnector : public msr::airlib::VehicleConnectorBase
     virtual void setPose(const Pose& pose, bool ignore_collison) override;
     virtual Pose getPose() override;
 
-    virtual void setSegmentationObjectID(const std::string& mesh_name, int object_id,
+    virtual bool setSegmentationObjectID(const std::string& mesh_name, int object_id,
         bool is_name_regex = false) override;
     virtual int getSegmentationObjectID(const std::string& mesh_name) override;
 

Unreal/Plugins/AirSim/Source/SimHUD/SimHUD.cpp

@@ -223,8 +223,8 @@ void ASimHUD::initializeSubWindows()
     //setup defaults
     if (camera_count > 0) {
         subwindow_cameras_[0] = wrapper->getCamera(0);
-        subwindow_cameras_[1] = wrapper->getCamera(camera_count > 3 ? 3 : 0);
-        subwindow_cameras_[2] = wrapper->getCamera(camera_count > 4 ? 4 : 0);
+        subwindow_cameras_[1] = wrapper->getCamera(0); //camera_count > 3 ? 3 : 0
+        subwindow_cameras_[2] = wrapper->getCamera(0); //camera_count > 4 ? 4 : 0
     }
     else
         subwindow_cameras_[0] = subwindow_cameras_[1] = subwindow_cameras_[2] = nullptr;