209 lines
7.2 KiB
Python
209 lines
7.2 KiB
Python
import os
|
||
import pandas as pd
|
||
from tqdm import tqdm
|
||
import rasterio
|
||
from PIL import Image
|
||
from collections import defaultdict
|
||
import numpy as np
|
||
Image.MAX_IMAGE_PIXELS = None # 禁用像素限制警告
|
||
class PredictionAggregator:
|
||
def __init__(self, img_dir=None):
|
||
"""
|
||
img_dir: 原始图像路径（用于自动获取 full_size）
|
||
"""
|
||
self.preds = defaultdict(list)
|
||
self.sizes = dict()
|
||
self.img_dir = img_dir
|
||
|
||
def _get_image_size(self, fname):
|
||
"""
|
||
自动根据图像格式获取 (H, W)
|
||
"""
|
||
path = os.path.join(self.img_dir, fname)
|
||
ext = os.path.splitext(fname)[1].lower()
|
||
if ext in ['.tif', '.tiff']:
|
||
with rasterio.open(path) as src:
|
||
return src.height, src.width # rasterio: (H, W)
|
||
else:
|
||
with Image.open(path) as img:
|
||
return img.size[1], img.size[0] # PIL: (W, H) → (H, W)
|
||
|
||
def add_patch(self, fname, x, y, pred_patch, valid_h, valid_w):
|
||
"""
|
||
添加一个 patch 到对应位置（自动裁去 padding）
|
||
fname: 原始图像名（如 image_1.png）
|
||
x, y: patch 左上角坐标
|
||
pred_patch: numpy array, shape=[H_pad, W_pad]
|
||
valid_h, valid_w: 原始补零前 patch 的真实高度和宽度
|
||
"""
|
||
if fname not in self.sizes:
|
||
if self.img_dir is None:
|
||
raise ValueError("img_dir 未指定，无法自动获取图像尺寸")
|
||
self.sizes[fname] = self._get_image_size(fname)
|
||
|
||
# 裁剪有效区域（去掉 padding）
|
||
patch_crop = pred_patch[:valid_h, :valid_w]
|
||
self.preds[fname].append((x, y, patch_crop))
|
||
|
||
def save_all(self, save_dir):
|
||
os.makedirs(save_dir, exist_ok=True)
|
||
out_save_path = ""
|
||
for fname, patches in self.preds.items():
|
||
H, W = self.sizes[fname]
|
||
canvas = np.zeros((H, W), dtype=np.uint8)
|
||
|
||
for x, y, patch in patches:
|
||
h, w = patch.shape
|
||
canvas[y:y+h, x:x+w] = patch
|
||
|
||
save_name = f"{fname}" # 保留原图扩展名
|
||
save_path = os.path.join(save_dir, save_name)
|
||
Image.fromarray(canvas).save(save_path)
|
||
out_save_path=save_path
|
||
return out_save_path
|
||
def save_one_pic(self, save_dir):
|
||
os.makedirs(save_dir, exist_ok=True)
|
||
out_save_path=""
|
||
for fname, patches in self.preds.items():
|
||
H, W = self.sizes[fname]
|
||
canvas = np.zeros((H, W), dtype=np.uint8)
|
||
|
||
for x, y, patch in patches:
|
||
h, w = patch.shape
|
||
canvas[y:y+h, x:x+w] = patch
|
||
|
||
save_name = f"{fname}" # 保留原图扩展名
|
||
save_path = os.path.join(save_dir, save_name)
|
||
Image.fromarray(canvas).save(save_path)
|
||
out_save_path=save_path
|
||
self.preds.clear()
|
||
return out_save_path
|
||
|
||
|
||
|
||
|
||
# class PredictionAggregator:
|
||
# def __init__(self, img_dir=None):
|
||
# """
|
||
# img_dir: 原始图像路径（用于自动获取 full_size）
|
||
# """
|
||
# self.preds = defaultdict(list)
|
||
# self.sizes = dict()
|
||
# self.img_dir = img_dir
|
||
|
||
# def _get_image_size(self, fname):
|
||
# """
|
||
# 自动根据图像格式获取 (H, W)
|
||
# """
|
||
# path = os.path.join(self.img_dir, fname)
|
||
# ext = os.path.splitext(fname)[1].lower()
|
||
# if ext in ['.tif', '.tiff']:
|
||
# with rasterio.open(path) as src:
|
||
# return src.height, src.width # rasterio: (H, W)
|
||
# else:
|
||
# with Image.open(path) as img:
|
||
# return img.size[1], img.size[0] # PIL: (W, H) → (H, W)
|
||
|
||
# def add_patch(self, fname, x, y, pred_patch):
|
||
# """
|
||
# 添加一个 patch 到对应位置
|
||
# fname: 原始图像名（如 image_1.png）
|
||
# x, y: patch 左上角坐标
|
||
# pred_patch: numpy array, shape=[H, W]
|
||
# """
|
||
# if fname not in self.sizes:
|
||
# if self.img_dir is None:
|
||
# raise ValueError("img_dir 未指定，无法自动获取图像尺寸")
|
||
# self.sizes[fname] = self._get_image_size(fname)
|
||
|
||
# self.preds[fname].append((x, y, pred_patch))
|
||
|
||
# def save_all(self, save_dir):
|
||
# os.makedirs(save_dir, exist_ok=True)
|
||
# for fname, patches in self.preds.items():
|
||
# H, W = self.sizes[fname]
|
||
# canvas = np.zeros((H, W), dtype=np.uint8)
|
||
|
||
# for x, y, patch in patches:
|
||
# h, w = patch.shape
|
||
# canvas[y:y+h, x:x+w] = patch
|
||
|
||
# save_name = f"pred_{fname}" # 保留原图扩展名
|
||
# save_path = os.path.join(save_dir, save_name)
|
||
# Image.fromarray(canvas).save(save_path)
|
||
|
||
|
||
class PatchIndexer:
|
||
def __init__(self, img_dir, patch_size=256, stride=256, save_dir=None):
|
||
self.img_dir = img_dir
|
||
self.patch_size = patch_size
|
||
self.stride = stride
|
||
self.coord_records = []
|
||
|
||
if save_dir is None:
|
||
self.save_dir = os.path.dirname(os.path.abspath(img_dir))
|
||
else:
|
||
self.save_dir = save_dir
|
||
|
||
def index_all(self, save_name=None):
|
||
filelist = sorted(f for f in os.listdir(self.img_dir)
|
||
if f.lower().endswith(('.tif', '.tiff', '.png', '.jpg', '.jpeg')))
|
||
print(f"[INFO] Found {len(filelist)} images.")
|
||
|
||
for fname in tqdm(filelist, desc='Indexing patches'):
|
||
self.index_one(fname)
|
||
|
||
if save_name is None:
|
||
folder_name = os.path.basename(self.save_dir.rstrip('/\\'))
|
||
save_name = f'coords.csv'
|
||
|
||
save_path = os.path.join(self.save_dir, save_name)
|
||
|
||
df = pd.DataFrame(self.coord_records)
|
||
if df.empty:
|
||
print("[⚠️] No patch coordinates were generated!")
|
||
else:
|
||
df.to_csv(save_path, index=False)
|
||
print(f"[✅] Saved {len(df)} patch records to {save_path}")
|
||
|
||
def index_one(self, fname):
|
||
path = os.path.join(self.img_dir, fname)
|
||
name, ext = os.path.splitext(fname)
|
||
ext = ext.lower()
|
||
|
||
try:
|
||
if ext in ['.tif', '.tiff']:
|
||
with rasterio.open(path) as src:
|
||
H, W = src.height, src.width
|
||
else:
|
||
with Image.open(path) as img:
|
||
W, H = img.size # 注意 PIL 是 (W, H)
|
||
except Exception as e:
|
||
print(f"[ERROR] Failed to open {fname}: {e}")
|
||
return
|
||
|
||
ps, st = self.patch_size, self.stride
|
||
|
||
if H < ps or W < ps:
|
||
print(f"[WARNING] Skip {fname}: too small ({W}x{H})")
|
||
return
|
||
|
||
y_list = list(range(0, H - ps + 1, st))
|
||
x_list = list(range(0, W - ps + 1, st))
|
||
|
||
if (H - ps) % st != 0 and (H - ps) not in y_list:
|
||
y_list.append(H - ps)
|
||
if (W - ps) % st != 0 and (W - ps) not in x_list:
|
||
x_list.append(W - ps)
|
||
|
||
for y in y_list:
|
||
for x in x_list:
|
||
self.coord_records.append({
|
||
'orig_name': fname,
|
||
'x': x,
|
||
'y': y,
|
||
'h': ps,
|
||
'w': ps,
|
||
'patch_name': f"{name}_{x}_{y}{ext}"
|
||
})
|