From ff2f45e671781e26d02e263848e0c50644ac2aec Mon Sep 17 00:00:00 2001 From: weixin_46229132 Date: Sat, 5 Jul 2025 19:44:40 +0800 Subject: [PATCH] =?UTF-8?q?=E9=95=B6=E5=B5=8C=E7=9A=84=E5=83=8F=E7=B4=A0?= =?UTF-8?q?=E9=80=89=E6=8B=A9=E7=AD=96=E7=95=A5?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- .gitignore | 1 + stac_mosaic_api.py | 12 +++- tools/stac_tiles_mosaic.py | 129 +++++++++++++++++++++++++++++++++++++ 3 files changed, 139 insertions(+), 3 deletions(-) create mode 100644 tools/stac_tiles_mosaic.py diff --git a/.gitignore b/.gitignore index 3ca8085..8dee0d2 100644 --- a/.gitignore +++ b/.gitignore @@ -1,6 +1,7 @@ # ---> Python # Custom dir /test_geosat1_imgs +/stac-fastapi-pgstac/pgdata # Byte-compiled / optimized / DLL files __pycache__/ diff --git a/stac_mosaic_api.py b/stac_mosaic_api.py index 6d0108c..92df1f2 100644 --- a/stac_mosaic_api.py +++ b/stac_mosaic_api.py @@ -4,6 +4,7 @@ from pydantic import BaseModel from shapely.geometry import Polygon, mapping from rio_tiler.io import Reader from rio_tiler.mosaic import mosaic_reader +from rio_tiler.mosaic.methods.defaults import MeanMethod # or LowestMethod import requests from fastapi.responses import StreamingResponse from io import BytesIO @@ -53,7 +54,8 @@ def fetch_items_by_polygon(geojson: dict) -> List[str]: def render_mosaic(image_paths: List[str], reader_fn, geo: object) -> BytesIO: - img, _ = mosaic_reader(image_paths, reader_fn, geo) + img, _ = mosaic_reader(image_paths, reader_fn, geo, + pixel_selection=MeanMethod()) buf = BytesIO(img.render(img_format="PNG")) buf.seek(0) return buf @@ -70,7 +72,9 @@ def bbox_mosaic(query: BBoxQuery): # TODO 这个函数的输入可以用*args, **kwargs，指定了max_size，后面记得改 def part_reader(src_path, part): with Reader(src_path) as cog: - return cog.part(part, max_size=1024) + img = cog.part(part, max_size=1024) + img.array.mask = (img.data == 0).all(axis=0) + return img # TODO 目前部分读取和镶嵌一起做做的，以后做实验需要解耦 img_buf = render_mosaic(image_paths, part_reader, query.bbox) @@ -88,7 +92,9 @@ def polygon_mosaic(polygon: GeoJSONPolygon = Body(...)): def feature_reader(src_path, feat): with Reader(src_path) as cog: - return cog.feature(feat, max_size=1024) + img = cog.feature(feat, max_size=1024) + img.array.mask = (img.data == 0).all(axis=0) + return img img_buf = render_mosaic( image_paths, feature_reader, polygon.model_dump()) diff --git a/tools/stac_tiles_mosaic.py b/tools/stac_tiles_mosaic.py new file mode 100644 index 0000000..b6c2e12 --- /dev/null +++ b/tools/stac_tiles_mosaic.py @@ -0,0 +1,129 @@ +from fastapi import FastAPI, Query, Body, HTTPException +from typing import List +from pydantic import BaseModel +from shapely.geometry import Polygon, mapping, shape +from rio_tiler.io import Reader +from rio_tiler.mosaic import mosaic_reader +from rio_tiler.mosaic.methods.defaults import MeanMethod # or LowestMethod +import requests +from fastapi.responses import StreamingResponse +from io import BytesIO +import morecantile +from starlette.responses import StreamingResponse +from PIL import Image +import numpy as np +import os + + +STAC_API_URL = "http://localhost:8082" +COLLECTION_ID = "geosat1" + + +# ---------- 实用函数 ---------- +def fetch_tiles_by_polygon(polygon) -> List[tuple]: + """ + 根据多边形获取相交的瓦片列表 + """ + zoom = 12 + tms = morecantile.tms.get("WebMercatorQuad") + tiles = list(tms.tiles(*polygon.bounds, zoom)) + + def tile_polygon(x, y, z): + bounds = tms.bounds(x, y, z) + return Polygon([ + (bounds.left, bounds.bottom), + (bounds.left, bounds.top), + (bounds.right, bounds.top), + (bounds.right, bounds.bottom), + (bounds.left, bounds.bottom) + ]) + + intersecting_tiles = [] + for x, y, z in tiles: + if polygon.intersects(tile_polygon(x, y, z)): + intersecting_tiles.append((x, y, z)) + + return intersecting_tiles + + +def fetch_items_by_tile(x, y, z) -> List[str]: + """ + 根据瓦片坐标 (x, y, z) 查询与该瓦片 bbox 相交的影像 + """ + tms = morecantile.tms.get("WebMercatorQuad") + bounds = tms.bounds(x, y, z) + bbox = [bounds.left, bounds.bottom, bounds.right, bounds.top] + url = f"{STAC_API_URL}/collections/{COLLECTION_ID}/items" + params = {"bbox": ",".join(map(str, bbox))} + r = requests.get(url, params=params) + r.raise_for_status() + data = r.json() + return [feat["assets"]["image"]["href"] for feat in data.get("features", [])] + + +def tile_mosaic_buf(i, x, y, z): + image_paths = fetch_items_by_tile(x, y, z) + if not image_paths: + # 返回一张全透明的空tile（PNG格式） + empty_img = Image.new("RGBA", (256, 256), (0, 0, 0, 0)) + buf = BytesIO() + empty_img.save(buf, format="PNG") + return buf.getvalue() + + def tile_reader(src_path, x, y, z): + with Reader(src_path) as cog: + img = cog.tile(x, y, z, tilesize=256) + img.array.mask = (img.data == 0).all(axis=0) + return img + img, _ = mosaic_reader(image_paths, tile_reader, x, y, z, pixel_selection=MeanMethod()) + # with open(f"imgs/output{i}.tif", "wb") as f: + # f.write(img.render(img_format="GTiff")) + return img.render(img_format="GTiff") + + +def main(): + polygon_coords = [ + [111.76, 29.25], + [111.89, 29.25], + [111.89, 29.62], + [111.76, 29.62], + [111.76, 29.25] + ] + polygon = Polygon(polygon_coords) + intersecting_tiles = fetch_tiles_by_polygon(polygon) + + tile_size = 256 # 每个瓦片的像素尺寸 + tiles = [] + xs, ys = [], [] + + # 先收集所有瓦片及其坐标 + for i, (x, y, z) in enumerate(intersecting_tiles): + tile_bytes = tile_mosaic_buf(i, x, y, z) + img = Image.open(BytesIO(tile_bytes)) + tiles.append((x, y, img)) + xs.append(x) + ys.append(y) + + # 计算拼接大图的尺寸 + min_x, max_x = min(xs), max(xs) + min_y, max_y = min(ys), max(ys) + width = (max_x - min_x + 1) * tile_size + height = (max_y - min_y + 1) * tile_size + + # 创建空白大图 + mosaic_img = Image.new("RGBA", (width, height)) + + # 按瓦片坐标拼接 + for x, y, img in tiles: + px = (x - min_x) * tile_size + py = (y - min_y) * tile_size # 不需要翻转，直接递增 + mosaic_img.paste(img, (px, py)) + + # 保存或展示大图 + output_path = "mosaic_output.png" + mosaic_img.save(output_path) + print(f"拼接完成，已保存为 {output_path}") + + +if __name__ == "__main__": + main()