Image Segmentation with Quickshift and Shapefile Creation

def segmentation(image_file, mask_file, shapefile): with rasterio.open(image_file) as dataset: data = dataset.read() data = data[[4, 3, 2], :, :].transpose(1, 2, 0) segments_quick = quickshift(data, kernel_size=3, max_dist=6, ratio=0.5)

    # Write the output file with georeference information
    with rasterio.open(
            mask_file,
            'w',
            driver='GTiff',
            height=data.shape[1],
            width=data.shape[2],
            count=1,
            dtype=segments_quick.dtype,
            crs=dataset.crs,
            transform=dataset.transform,
    ) as dst:
        dst.write(segments_quick[np.newaxis, :, :])

The function you provided performs image segmentation using the Quickshift algorithm and writes the segmented output to a GeoTIFF file. However, it does not currently handle the creation of a shapefile. Below is an enhanced version of the function that includes the creation of a shapefile from the segmented image.

To achieve this, we will use the rasterio.features.shapes function to extract the shapes from the segmented image and the fiona library to write these shapes to a shapefile.

Here's the complete function:

import rasterio
from skimage.segmentation import quickshift
import numpy as np
import fiona
from fiona.crs import from_epsg
from shapely.geometry import shape, mapping

def segmentation(image_file, mask_file, shapefile):
    with rasterio.open(image_file) as dataset:
        data = dataset.read()
        data = data[[4, 3, 2], :, :].transpose(1, 2, 0)
        segments_quick = quickshift(data, kernel_size=3, max_dist=6, ratio=0.5)

        # Write the output file with georeference information
        with rasterio.open(
                mask_file,
                'w',
                driver='GTiff',
                height=data.shape[0],
                width=data.shape[1],
                count=1,
                dtype=segments_quick.dtype,
                crs=dataset.crs,
                transform=dataset.transform,
        ) as dst:
            dst.write(segments_quick[np.newaxis, :, :])

        # Extract shapes from the segmented image
        shapes = rasterio.features.shapes(segments_quick, transform=dataset.transform)

        # Define the schema for the shapefile
        schema = {
            'geometry': 'Polygon',
            'properties': {'value': 'int'},
        }

        # Write the shapes to a shapefile
        with fiona.open(shapefile, 'w', driver='ESRI Shapefile', crs=from_epsg(dataset.crs.to_epsg()), schema=schema) as shp:
            for geom, value in shapes:
                shp.write({
                    'geometry': mapping(shape(geom)),
                    'properties': {'value': int(value)},
                })

# Example usage
segmentation('input_image.tif', 'output_mask.tif', 'output_shapefile.shp')

Explanation:

1. Reading the Image: The image is read using rasterio.open, and the bands are selected and transposed to match the expected input format for the Quickshift algorithm.
2. Segmentation: The quickshift function from skimage.segmentation is used to segment the image.
3. Writing the Mask File: The segmented image is written to a GeoTIFF file with the same georeference information as the input image.
4. Extracting Shapes: The rasterio.features.shapes function is used to extract the shapes (polygons) from the segmented image.
5. Writing the Shapefile: The shapes are written to a shapefile using the fiona library. The schema defines the geometry type as 'Polygon' and includes a property 'value' to store the segment value.

Make sure you have the necessary libraries installed:

pip install rasterio scikit-image fiona shapely

This function will create both a segmented GeoTIFF file and a shapefile containing the polygons of the segmented regions.