diff --git a/src/segmentation_a_direct.jl b/src/segmentation_a_direct.jl
index 527031b9..170a6023 100644
--- a/src/segmentation_a_direct.jl
+++ b/src/segmentation_a_direct.jl
@@ -9,22 +9,37 @@ Apply k-means segmentation to a gray image to isolate a cluster group representi
 
 """
 function kmeans_segmentation(
-    gray_image::Matrix{Gray{Float64}}, ice_labels::Vector{Int64}
+    gray_image::Matrix{Gray{Float64}},
+    ice_labels::Vector{Int64},
+    k::Int64=4,
+    maxiter::Int64=50,
 )::BitMatrix
-    Random.seed!(45) # this seed generates consistent clusters for the final output
-    gray_image_height, gray_image_width = size(gray_image)
-    gray_image_1d = vec(gray_image)
-    @info("Done with reshape")
+    segmented = kmeans_segmentation(gray_image; k=k, maxiter=maxiter)
 
-    ## NOTE(tjd): this clusters into 4 classes and solves iteratively with a max of 50 iterations
+    ## Isolate ice floes and contrast from background
+    segmented_ice = get_segmented_ice(segmented, ice_labels)
+    return segmented_ice
+end
+
+function kmeans_segmentation(
+    gray_image::Matrix{Gray{Float64}}; k::Int64=4, maxiter::Int64=50
+)
+    Random.seed!(45)
+
+    ## NOTE(tjd): this clusters into k classes and solves iteratively with a max of maxiter iterations
     feature_classes = Clustering.kmeans(
-        gray_image_1d, 4; maxiter=50, display=:none, init=:kmpp
+        vec(gray_image), k; maxiter=maxiter, display=:none, init=:kmpp
     )
+
     class_assignments = assignments(feature_classes)
 
-    ## NOTE(tjd): this reshapes column major vector of kmeans classes back into original image shape
-    segmented = reshape(class_assignments, gray_image_height, gray_image_width)
+    ## NOTE(tjd): this clusters into 4 classes and solves iteratively with a max of 50 iterations
+    segmented = reshape(class_assignments, size(gray_image))
+
+    return segmented
+end
 
+function get_segmented_ice(segmented::Matrix{Int64}, ice_labels::Vector{Int64})
     ## Isolate ice floes and contrast from background
     nlabel = StatsBase.mode(segmented[ice_labels])
     segmented_ice = segmented .== nlabel
@@ -76,9 +91,7 @@ function segmentation_A(
 
     segmented_bridged = IceFloeTracker.bridge(segmented_opened_branched)
 
-    #segmented_ice_filled = .!bwareamaxfilt(.!segmented_bridged)
     segmented_ice_filled = IceFloeTracker.MorphSE.fill_holes(segmented_bridged)
-    @info "Done filling segmented_ice"
 
     diff_matrix = segmented_ice_opened .!= segmented_ice_filled
 
diff --git a/test/test-segmentation-a.jl b/test/test-segmentation-a.jl
index 020050f9..e12c343c 100644
--- a/test/test-segmentation-a.jl
+++ b/test/test-segmentation-a.jl
@@ -21,15 +21,12 @@
     @time segmented_ice_cloudmasked = IceFloeTracker.segmented_ice_cloudmasking(
         ice_water_discriminated_image, cloudmask, ice_labels
     )
-    IceFloeTracker.@persist segmented_ice_cloudmasked "./test_outputs/segmented_a_ice_cloudmasked.png" true
 
     @time segmented_A = IceFloeTracker.segmentation_A(segmented_ice_cloudmasked)
-    IceFloeTracker.@persist segmented_A "./test_outputs/segmented_a.png" true
 
     @time segmented_ice = IceFloeTracker.kmeans_segmentation(
         ice_water_discriminated_image, ice_labels
     )
-    IceFloeTracker.@persist segmented_ice "./test_outputs/segmented_a_ice.png" true
 
     @test typeof(segmented_A) == typeof(matlab_segmented_A_bitmatrix)
     @test test_similarity(matlab_segmented_A_bitmatrix, segmented_A, 0.039)