Refactor the segmentation code