From 37bef05d7032ca48913732906d9bfc0b217505b3 Mon Sep 17 00:00:00 2001
From: Michel Breyer <michel.p.breyer@gmail.com>
Date: Wed, 7 Jul 2021 10:44:57 +0200
Subject: [PATCH] Prevent the floor from being selected as a target

---
 simulation.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/simulation.py b/simulation.py
index 676ff8d..e3d460b 100644
--- a/simulation.py
+++ b/simulation.py
@@ -103,6 +103,8 @@ class Simulation(BtSim):
     def select_target(self):
         img = self.camera.get_image()
         uids, counts = np.unique(img.mask, return_counts=True)
+        mask = np.isin(uids, self.object_uids)  # remove ids of the floor, table, etc
+        uids, counts = uids[mask], counts[mask]
         target_uid = uids[np.argmin(counts)]
         p.changeVisualShape(target_uid, -1, rgbaColor=[1, 0, 0, 1])
         return target_uid