How to add or remove a color map from a raster dataset

static void addRemoveColormap(IRasterDataset rasterDataset, IRasterDataset
    colormapDataset)throws Exception{
    //If the colormapDataset is null, it will remove the colormap.
    //If the colormapDataset is not null, it will add the colormap of the colormapDataset to the input raster.

    //Cast to IRasterDatasetEdit for the colormap operation.
    IRasterDatasetEdit rasterDatasetEdit = new IRasterDatasetEditProxy(rasterDataset)
        ;

    //Add a colormap from a given raster dataset.
    if (colormapDataset != null){
        //Get raster colormap from the given raster dataset.
        IRasterBandCollection rasterBandCollection = new IRasterBandCollectionProxy
            (colormapDataset);
        IRasterBand rasterBand = rasterBandCollection.item(0);
        IRasterColormap rasterColormap = rasterBand.getColormap();
        if (rasterColormap != null){
            rasterDatasetEdit.alterColormap(rasterColormap);
        }
    }
    else{
        //Remove the raster colormap.
        rasterDatasetEdit.deleteColormap();
    }
}