130 lines
3.9 KiB
Python
130 lines
3.9 KiB
Python
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()
|