134 lines
4.0 KiB
Python
134 lines
4.0 KiB
Python
import os
|
|
import sys
|
|
import shutil
|
|
from pathlib import Path
|
|
import matplotlib.pyplot as plt
|
|
from matplotlib.widgets import RectangleSelector
|
|
import pandas as pd
|
|
|
|
# 添加项目根目录到系统路径
|
|
project_root = str(Path(__file__).parent.parent)
|
|
sys.path.append(project_root)
|
|
|
|
from utils.gps_extractor import GPSExtractor
|
|
|
|
class GPSSelector:
|
|
def __init__(self, image_dir: str, output_dir: str = None):
|
|
self.image_dir = image_dir
|
|
self.output_dir = output_dir
|
|
self.gps_points = None
|
|
self.selected_points = []
|
|
self.fig, self.ax = plt.subplots(figsize=(12, 8))
|
|
self.scatter = None
|
|
self.rs = None
|
|
self.setup_plot()
|
|
|
|
def extract_gps(self):
|
|
"""提取GPS数据"""
|
|
extractor = GPSExtractor(self.image_dir)
|
|
self.gps_points = extractor.extract_all_gps()
|
|
print(f"成功提取 {len(self.gps_points)} 个GPS点")
|
|
|
|
def setup_plot(self):
|
|
"""设置绘图"""
|
|
self.ax.set_title('GPS Points - 使用鼠标拖动选择要删除的点')
|
|
self.ax.set_xlabel('Longitude')
|
|
self.ax.set_ylabel('Latitude')
|
|
self.ax.grid(True)
|
|
|
|
# 设置矩形选择器
|
|
self.rs = RectangleSelector(
|
|
self.ax, self.on_select,
|
|
interactive=True,
|
|
useblit=True,
|
|
button=[1], # 只响应左键
|
|
props=dict(facecolor='red', alpha=0.3)
|
|
)
|
|
|
|
# 添加按钮回调
|
|
self.fig.canvas.mpl_connect('key_press_event', self.on_key_press)
|
|
|
|
def plot_gps_points(self):
|
|
"""绘制GPS点"""
|
|
if self.scatter is not None:
|
|
self.scatter.remove()
|
|
|
|
self.scatter = self.ax.scatter(
|
|
self.gps_points['lon'],
|
|
self.gps_points['lat'],
|
|
c='blue',
|
|
s=20,
|
|
alpha=0.6
|
|
)
|
|
self.fig.canvas.draw_idle()
|
|
|
|
def on_select(self, eclick, erelease):
|
|
"""矩形选择回调"""
|
|
x1, y1 = eclick.xdata, eclick.ydata
|
|
x2, y2 = erelease.xdata, erelease.ydata
|
|
|
|
# 获取选中区域内的点
|
|
mask = (
|
|
(self.gps_points['lon'] >= min(x1, x2)) &
|
|
(self.gps_points['lon'] <= max(x1, x2)) &
|
|
(self.gps_points['lat'] >= min(y1, y2)) &
|
|
(self.gps_points['lat'] <= max(y1, y2))
|
|
)
|
|
|
|
selected = self.gps_points[mask]
|
|
self.selected_points.extend(selected['file'].tolist())
|
|
|
|
# 从数据中移除选中的点
|
|
self.gps_points = self.gps_points[~mask]
|
|
|
|
# 更新绘图
|
|
self.plot_gps_points()
|
|
print(f"选中 {len(selected)} 个点，剩余 {len(self.gps_points)} 个点")
|
|
|
|
def on_key_press(self, event):
|
|
"""键盘事件回调"""
|
|
if event.key == 'enter':
|
|
self.save_results()
|
|
plt.close()
|
|
elif event.key == 'escape':
|
|
plt.close()
|
|
|
|
def save_results(self):
|
|
"""保存结果"""
|
|
if not self.output_dir:
|
|
return
|
|
|
|
# 创建输出目录
|
|
os.makedirs(self.output_dir, exist_ok=True)
|
|
removed_dir = os.path.join(self.output_dir, "removed_images")
|
|
os.makedirs(removed_dir, exist_ok=True)
|
|
|
|
# 移动被删除的图像
|
|
for img_name in self.selected_points:
|
|
src = os.path.join(self.image_dir, img_name)
|
|
dst = os.path.join(removed_dir, img_name)
|
|
shutil.move(src, dst)
|
|
|
|
# 保存剩余点的信息
|
|
self.gps_points.to_csv(
|
|
os.path.join(self.output_dir, "remaining_points.csv"),
|
|
index=False
|
|
)
|
|
|
|
print(f"已移动 {len(self.selected_points)} 张图片到 {removed_dir}")
|
|
print(f"保留 {len(self.gps_points)} 个点")
|
|
|
|
def run(self):
|
|
"""运行选择器"""
|
|
self.extract_gps()
|
|
self.plot_gps_points()
|
|
plt.show()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
# 使用示例
|
|
selector = GPSSelector(
|
|
image_dir=r"E:\datasets\UAV\1619\project\images",
|
|
output_dir=r"E:\datasets\UAV\1619\filtered"
|
|
)
|
|
selector.run() |