cuibinge/typical-marine-ecological-feature-recognition-code

1"""Post-process seaweed masks with no-data and conservative water constraints."""2 3from __future__ import annotations4 5import argparse6from pathlib import Path7 8import numpy as np9import rasterio10from rasterio.windows import Window11 12 13def parse_args() -> argparse.Namespace:14    parser = argparse.ArgumentParser(description=__doc__)15    parser.add_argument("--image", required=True, help="Source 4-band fused/multispectral raster.")16    parser.add_argument("--mask", required=True, help="Predicted binary mask raster.")17    parser.add_argument("--output", required=True, help="Filtered output mask raster.")18    parser.add_argument("--stripe-height", type=int, default=1024)19    parser.add_argument("--black-threshold", type=float, default=32.0, help="Pixels with all bands <= this are no-data.")20    parser.add_argument(21        "--land-ndwi-threshold",22        type=float,23        default=-0.35,24        help="Very conservative NDWI cutoff for obvious dry land. Lower is safer for floating algae.",25    )26    parser.add_argument(27        "--land-nir-ratio",28        type=float,29        default=1.6,30        help="Only suppress NDWI-low pixels when NIR is this many times brighter than green.",31    )32    parser.add_argument("--no-land-filter", action="store_true", help="Only remove black/no-data regions.")33    return parser.parse_args()34 35 36def obvious_land_mask(tile: np.ndarray, ndwi_threshold: float, nir_ratio: float) -> np.ndarray:37    """Return a conservative land mask from B,G,R,NIR-like 4-band data.38 39    This is not a substitute for an official coastline/water mask. It only removes40    strongly land-like pixels to avoid deleting real floating algae.41    """42    if tile.shape[0] < 4:43        return np.zeros(tile.shape[1:], dtype=bool)44    green = tile[1].astype(np.float32, copy=False)45    nir = tile[3].astype(np.float32, copy=False)46    ndwi = (green - nir) / (green + nir + 1e-6)47    return (ndwi < ndwi_threshold) & (nir > green * nir_ratio)48 49 50def main() -> None:51    args = parse_args()52    image_path = Path(args.image)53    mask_path = Path(args.mask)54    output_path = Path(args.output)55    output_path.parent.mkdir(parents=True, exist_ok=True)56 57    with rasterio.open(image_path) as image_src, rasterio.open(mask_path) as mask_src:58        if (image_src.width, image_src.height) != (mask_src.width, mask_src.height):59            raise ValueError(60                f"Image and mask sizes differ: image={image_src.width}x{image_src.height}, "61                f"mask={mask_src.width}x{mask_src.height}"62            )63        profile = mask_src.profile.copy()64        profile.update(count=1, dtype="uint8", compress="lzw", nodata=0)65 66        total_pixels = image_src.width * image_src.height67        input_fg = 068        output_fg = 069        invalid_pixels = 070        land_pixels = 071 72        with rasterio.open(output_path, "w", **profile) as dst:73            for y in range(0, image_src.height, args.stripe_height):74                height = min(args.stripe_height, image_src.height - y)75                window = Window(0, y, image_src.width, height)76                image = image_src.read(window=window)77                mask = mask_src.read(1, window=window)78 79                predicted = mask > 080                valid = np.max(image, axis=0) > args.black_threshold81                land = np.zeros(valid.shape, dtype=bool)82                if not args.no_land_filter:83                    land = obvious_land_mask(image, args.land_ndwi_threshold, args.land_nir_ratio)84 85                filtered = predicted & valid & ~land86                dst.write((filtered.astype(np.uint8) * 255), 1, window=window)87 88                input_fg += int(predicted.sum())89                output_fg += int(filtered.sum())90                invalid_pixels += int((~valid).sum())91                land_pixels += int(land.sum())92 93    print(f"image={image_path}")94    print(f"mask={mask_path}")95    print(f"output={output_path}")96    print(f"total_pixels={total_pixels}")97    print(f"input_foreground={input_fg} ratio={input_fg / total_pixels:.6f}")98    print(f"output_foreground={output_fg} ratio={output_fg / total_pixels:.6f}")99    print(f"removed_foreground={input_fg - output_fg} ratio={(input_fg - output_fg) / total_pixels:.6f}")100    print(f"invalid_or_black_pixels={invalid_pixels} ratio={invalid_pixels / total_pixels:.6f}")101    print(f"conservative_land_pixels={land_pixels} ratio={land_pixels / total_pixels:.6f}")102 103 104if __name__ == "__main__":105    main()106