lidR Vegetation Classification

I’m new to R and to lidR.

I have LAS files that are already ground classified and I’m attempting to use the method below to classify the vegetation and then write a new LAS file retaining the ground and vegetation classifications. The resulting LAS files end up with every point classified as vegetation.

What am I doing wrong?

Here’s my script:

library(lidR)
las <- readLAS("D:/tree-segmentation/test.las", select = "xyzc", filter = "-drop_z_below 0")
las <- segment_trees(las, li2012(R = 3, speed_up = 5))

plot(las, color = "treeID")


vegetation_class = function(Z)
{
  mz = max(Z) 
  if (mz < 10)
    return(LASLOWVEGETATION)
  else if (mz >= 10 & mz < 22)
    return(LASMEDIUMVEGETATION)
  else
    return(LASHIGHVEGETATION)
}



las@data[!is.na(treeID), Classification := vegetation_class(Z), by = treeID]
plot(las, color = "Classification")
writeLAS(las, "D:/tree-segmentation/test-veg-classified.las", index = FALSE)

las@data[!is.na(treeID) & Classification != LASGROUND, Classification := vegetation_class(Z), by = treeID]