'''
根据检测结果，获取检测框中心的经纬度信息，并可视化
'''
import json
import rasterio
from rasterio.warp import transform
import numpy as np
import cv2

# 输入输出路径
tif_path = r"C:/Users/xiaobao/Desktop/Router/img_88c4c436-7466-4b02-8c69-a55b1f71eb0a.tif"
json_path = r"C:/Users/xiaobao/Desktop/Router/detection_results.json"
output_tif = r"C:/Users/xiaobao/Desktop/Router/img_with_detections.tif"

# 打开影像
ds = rasterio.open(tif_path)
print("CRS:", ds.crs)  # EPSG:32651
print("Transform:", ds.transform)

# 读取前三波段
img = ds.read([1,2,3])
img = np.transpose(img, (1,2,0))
img = cv2.convertScaleAbs(img, alpha=(255.0 / img.max()))

# 读取检测结果
with open(json_path, "r", encoding="utf-8") as f:
    data = json.load(f)
detections = data["detections"]

for det in detections:
    bbox = det["bbox"]
    cx = float(det["center"]["x"])
    cy = float(det["center"]["y"])

    # 像素坐标 -> UTM 坐标
    x_utm, y_utm = ds.xy(cy, cx)

    # UTM -> 经纬度 (EPSG:4326)
    lon, lat = transform(ds.crs, "EPSG:4326", [x_utm], [y_utm])
    lon, lat = lon[0], lat[0]

    # 绘制检测框和经纬度
    cv2.rectangle(img,
                  (int(bbox["minx"]), int(bbox["miny"])),
                  (int(bbox["maxx"]), int(bbox["maxy"])),
                  (255, 0, 0), 30)
    cv2.circle(img, (int(cx), int(cy)), 20, (0, 255, 0), -1)

    cv2.putText(img, f"{lon:.6f},{lat:.6f}",
                (int(cx) -580, int(cy) - 55),
                cv2.FONT_HERSHEY_SIMPLEX, 3,
                (255, 255, 0), 8, cv2.LINE_AA)

# 保存带绘制的新 GeoTIFF
out_img = np.transpose(img, (2,0,1))
profile = ds.profile
profile.update(dtype=rasterio.uint8, count=3)

with rasterio.open(output_tif, "w", **profile) as dst:
    dst.write(out_img.astype(rasterio.uint8))

print(f"结果已保存: {output_tif}")