diff --git a/im_generator.py b/im_generator.py
index 8560805..8446adf 100644
--- a/im_generator.py
+++ b/im_generator.py
@@ -136,7 +136,7 @@ def get_bounding_box(img):
     ms_colors = [(0, 187, 255), (0, 187, 124),
                  (241, 161, 0), (20, 83, 246)]  # in BGR
     color_idx = np.random.randint(0, 4)  # [0, 4)
-    output[output[:, :, 3] == 0] = (*ms_colors[color_idx], 255)
+    output[output[:, :, 3] <= 100] = (*ms_colors[color_idx], 255)
     return output, annos