diff --git a/aliby/tile/traps.py b/aliby/tile/traps.py index 08b5484778fbf5a372b5e9d66a12bc388a937c96..5d3e25fee5a01829fb4640a77971b40219a01e72 100644 --- a/aliby/tile/traps.py +++ b/aliby/tile/traps.py @@ -91,10 +91,28 @@ def segment_traps(image, tile_size, downscale=0.4): x - tile_size // 2 : x + tile_size // 2, y - tile_size // 2 : y + tile_size // 2 ] + # add template as mean of found traps + mean_template = ( + np.dstack( + [ + image[ + x - tile_size // 2 : x + tile_size // 2, + y - tile_size // 2 : y + tile_size // 2, + ] + for x, y in traps + ] + ) + .astype(int) + .mean(axis=-1) + ) + traps = identify_trap_locations(image, template) + mean_traps = identify_trap_locations(image, mean_template) + + traps = traps if len(traps) > len(mean_traps) else mean_traps - if len(traps) < 10 and downscale != 1: - print("Trying again.") + if len(traps) < 30 and downscale != 1: + print("Tiler:TrapIdentification: Trying again.") return segment_traps(image, tile_size, downscale=1) return traps