split_dota.py
ultralytics\data\split_dota.py
目录
2.def bbox_iof(polygon1: np.ndarray, bbox2: np.ndarray, eps: float = 1e-6) -> np.ndarray:
3.def load_yolo_dota(data_root: str, split: str = "train") -> List[Dict[str, Any]]:
1.所需的库和模块
# Ultralytics 🚀 AGPL-3.0 License - https://ultralytics.com/license
import itertools
from glob import glob
from math import ceil
from pathlib import Path
from typing import Any, Dict, List, Tuple
import cv2
import numpy as np
from PIL import Image
from ultralytics.data.utils import exif_size, img2label_paths
from ultralytics.utils import TQDM
from ultralytics.utils.checks import check_requirements
2.def bbox_iof(polygon1: np.ndarray, bbox2: np.ndarray, eps: float = 1e-6) -> np.ndarray:
# 这段代码定义了一个函数 bbox_iof ,用于计算两个几何形状(多边形和边界框)之间的交并比(Intersection over Foreground,IOF)。
# 定义了一个函数 bbox_iof ,接受三个参数:
# 1.polygon1 :一个 NumPy 数组,表示一组多边形的顶点坐标。
# 2.bbox2 :一个 NumPy 数组,表示一组边界框的坐标,格式为 [x1, y1, x2, y2] ,分别表示左上角和右下角的坐标。
# 3.eps :一个浮点数,默认值为 1e-6 ,用于数值稳定性,避免除以零的情况。
# 函数返回一个 NumPy 数组,表示每个多边形与每个边界框之间的 IOF 值。
def bbox_iof(polygon1: np.ndarray, bbox2: np.ndarray, eps: float = 1e-6) -> np.ndarray:
# 计算多边形与边界框之间的前景交集 (IoF)。
# 参数:
# polygon1 (np.ndarray):多边形坐标,形状为 (N, 8)。
# bbox2 (np.ndarray):边界框,形状为 (N, 4)。
# eps (float,可选):较小的值,用于避免除以零。
# 返回:
# (np.ndarray):IoF 分数,形状为 (N, 1);如果 bbox2 为 (M, 4),则返回 (N, M)。
# 说明:
# 多边形格式:[x1, y1, x2, y2, x3, y3, x4, y4]。
# 边界框格式:[x_min, y_min, x_max, y_max]。
"""
Calculate Intersection over Foreground (IoF) between polygons and bounding boxes.
Args:
polygon1 (np.ndarray): Polygon coordinates with shape (N, 8).
bbox2 (np.ndarray): Bounding boxes with shape (N, 4).
eps (float, optional): Small value to prevent division by zero.
Returns:
(np.ndarray): IoF scores with shape (N, 1) or (N, M) if bbox2 is (M, 4).
Notes:
Polygon format: [x1, y1, x2, y2, x3, y3, x4, y4].
Bounding box format: [x_min, y_min, x_max, y_max].
"""
# 检查是否安装了 shapely 库,且版本不低于 2.0.0 。如果未安装或版本不符合要求,会抛出错误。
check_requirements("shapely>=2.0.0")
# 从 shapely.geometry 模块中导入 Polygon 类,用于后续创建多边形对象并进行几何运算。
from shapely.geometry import Polygon
# 将输入的 polygon1 数组重新整形为形状为 (n, 4, 2) 的数组,其中 n 是多边形的数量, 4 表示每个多边形有 4 个顶点, 2 表示每个顶点的坐标为 (x, y) 。
polygon1 = polygon1.reshape(-1, 4, 2)
# 计算每个多边形的左上角点坐标。通过沿着顶点维度( axis=-2 )取最小值,得到每个多边形的左上角点 (x_min, y_min) 。
lt_point = np.min(polygon1, axis=-2) # left-top
# 计算每个多边形的右下角点坐标。通过沿着顶点维度( axis=-2 )取最大值,得到每个多边形的右下角点 (x_max, y_max) 。
rb_point = np.max(polygon1, axis=-2) # right-bottom
# 将左上角点和右下角点拼接起来,形成每个多边形对应的边界框 bbox1 ,格式为 [x_min, y_min, x_max, y_max] 。
bbox1 = np.concatenate([lt_point, rb_point], axis=-1)
# 计算 bbox1 和 bbox2 的左上角点的逐元素最大值,得到重叠区域的左上角点坐标。
lt = np.maximum(bbox1[:, None, :2], bbox2[..., :2])
# 计算 bbox1 和 bbox2 的右下角点的逐元素最小值,得到重叠区域的右下角点坐标。
rb = np.minimum(bbox1[:, None, 2:], bbox2[..., 2:])
# 计算重叠区域的宽度和高度,即 rb - lt 。使用 np.clip 将负值裁剪为 0,避免出现负的宽度或高度。
wh = np.clip(rb - lt, 0, np.inf)
# 计算重叠区域的面积,即宽度乘以高度。
h_overlaps = wh[..., 0] * wh[..., 1]
# 将 bbox2 的坐标分解为左、上、右、下四个边界。
left, top, right, bottom = (bbox2[..., i] for i in range(4))
# 根据 bbox2 的边界坐标构造对应的多边形顶点坐标。每个边界框被转换为一个矩形多边形,顶点顺序为左上、右上、右下、左下。
polygon2 = np.stack([left, top, right, top, right, bottom, left, bottom], axis=-1).reshape(-1, 4, 2)
# 将 polygon1 中的每个多边形转换为 shapely.geometry.Polygon 对象,用于后续的几何运算。
sg_polys1 = [Polygon(p) for p in polygon1]
# 将 polygon2 中的每个多边形转换为 shapely.geometry.Polygon 对象。
sg_polys2 = [Polygon(p) for p in polygon2]
# 初始化一个与 h_overlaps 形状相同的零矩阵,用于存储多边形之间的实际重叠面积。
overlaps = np.zeros(h_overlaps.shape)
# 遍历 h_overlaps 中非零元素的位置(表示有重叠的区域)。对于每一对多边形,计算它们的交集面积,并将结果存储到 overlaps 中。
for p in zip(*np.nonzero(h_overlaps)):
overlaps[p] = sg_polys1[p[0]].intersection(sg_polys2[p[-1]]).area
# 计算 polygon1 中每个多边形的面积,并将其存储为一个 NumPy 数组。
unions = np.array([p.area for p in sg_polys1], dtype=np.float32)
# 将 unions 的形状扩展为 (n, 1) ,以便后续进行广播运算。
unions = unions[..., None]
# 使用 np.clip 将 unions 中的值裁剪到 [eps, inf) 范围内,避免除以零的情况。
unions = np.clip(unions, eps, np.inf)
# 计算每个多边形与每个边界框之间的 IOF 值,即重叠面积除以多边形的面积。
outputs = overlaps / unions
# 如果 outputs 是一维数组(即只有一个多边形或一个边界框),将其扩展为二维数组,形状为 (n, 1) 。
if outputs.ndim == 1:
outputs = outputs[..., None]
# 返回计算得到的 IOF 值数组。
return outputs
# 这段代码实现了一个计算多边形与边界框之间交并比(IOF)的功能。它首先将多边形和边界框转换为几何对象,然后通过计算重叠面积与多边形面积的比值来得到 IOF。代码中使用了 shapely 库来处理几何运算,并通过 NumPy 进行高效的数组操作。这种计算在计算机视觉和目标检测任务中非常有用,例如评估多边形目标与边界框预测之间的匹配程度。
3.def load_yolo_dota(data_root: str, split: str = "train") -> List[Dict[str, Any]]:
# 这段代码定义了一个函数 load_yolo_dota ,用于加载 YOLO 格式的 DOTA 数据集(一种用于目标检测的遥感图像数据集)。它读取图像文件和对应的文件标签,并将它们的元数据和标注信息整理成一个列表,每个元素是一个包含图像尺寸、标签和文件路径的字典。
# 定义了一个函数 load_yolo_dota ,接受两个参数:
# 1.data_root :一个字符串,表示数据集的根目录路径。
# 2.split :一个字符串,默认值为 "train" ,表示要加载的数据集划分(训练集或验证集)。
# 函数返回一个列表,列表中的每个元素是一个字典,包含图像的原始尺寸、标签信息和文件路径。
def load_yolo_dota(data_root: str, split: str = "train") -> List[Dict[str, Any]]:
# 加载 DOTA 数据集注释和图像信息。
# 参数:
# data_root (str):数据根目录。
# split (str,可选):拆分后的数据集,可以是“train”或“val”。
# 返回:
# (List[Dict[str, Any]]):包含图像信息的注释字典列表。
# 说明:
# DOTA 数据集的目录结构如下:
# - data_root
# - images
# - train
# - val
# - labels
# - train
# - val
"""
Load DOTA dataset annotations and image information.
Args:
data_root (str): Data root directory.
split (str, optional): The split data set, could be 'train' or 'val'.
Returns:
(List[Dict[str, Any]]): List of annotation dictionaries containing image information.
Notes:
The directory structure assumed for the DOTA dataset:
- data_root
- images
- train
- val
- labels
- train
- val
"""
# 断言 split 参数的值必须是 "train" 或 "val" ,否则抛出异常。这确保了函数只加载训练集或验证集。
assert split in {"train", "val"}, f"Split must be 'train' or 'val', not {split}." # 拆分必须是“train”或“val”,而不是{split}。
# 使用 pathlib.Path 构造图像文件夹的路径,路径格式为 [data_root]/images/[split] 。
im_dir = Path(data_root) / "images" / split
# 断言图像文件夹路径存在,否则抛出异常。这确保了数据集路径正确。
assert im_dir.exists(), f"Can't find {im_dir}, please check your data root." # 找不到 {im_dir},请检查您的数据根目录。
# 使用 glob 函数获取图像文件夹中所有文件的路径。 glob 会匹配文件夹中的所有文件(包括图像文件和可能的其他文件)。
im_files = glob(str(Path(data_root) / "images" / split / "*"))
# 调用函数 img2label_paths (未在代码中定义,但可以推测其功能),将图像文件路径转换为对应的标签文件路径。通常,标签文件的命名与图像文件一致,但扩展名可能不同。
lb_files = img2label_paths(im_files)
# 初始化一个空列表 annos ,用于存储加载的标注信息。
annos = []
# 遍历图像文件路径和标签文件路径的配对。
for im_file, lb_file in zip(im_files, lb_files):
# 使用 PIL.Image.open 打开图像文件,并调用 exif_size 函数(未在代码中定义,但可以推测其功能)获取图像的原始宽度和高度。
w, h = exif_size(Image.open(im_file))
# 打开对应的标签文件,指定编码为 UTF-8。
with open(lb_file, encoding="utf-8") as f:
# 读取标签文件的内容,按行分割,过滤掉空行,然后将每行按空格分割为一个列表。这通常用于解析 YOLO 格式的标签文件,其中每行表示一个目标的标注信息。
lb = [x.split() for x in f.read().strip().splitlines() if len(x)]
# 将解析后的标签信息转换为 NumPy 数组,数据类型为 float32 。
lb = np.array(lb, dtype=np.float32)
# 将当前图像的元数据和标注信息整理成一个字典,并将其添加到 annos 列表中。字典包含以下键值对:
# "ori_size" :图像的原始尺寸 (高度, 宽度) 。
# "label" :标注信息数组。
# "filepath" :图像文件的路径。
annos.append(dict(ori_size=(h, w), label=lb, filepath=im_file))
# 返回包含所有图像标注信息的列表 annos 。
return annos
# 这段代码实现了一个用于加载 YOLO 格式的 DOTA 数据集的函数。它通过以下步骤完成任务: 验证输入参数的合法性(数据集划分必须是 "train" 或 "val" )。 构造图像文件夹路径,并验证路径是否存在。 获取图像文件路径和对应的标签文件路径。 遍历每对图像和标签文件,读取图像尺寸和标签信息。 将图像的元数据和标注信息整理成字典,并存储到列表中。最终返回的列表 annos 包含了数据集中每张图像的标注信息,可用于后续的目标检测模型训练或验证。
4.def get_windows(im_size: Tuple[int, int], crop_sizes: Tuple[int, ...] = (1024,), gaps: Tuple[int, ...] = (200,), im_rate_thr: float = 0.6, eps: float = 0.01,) -> np.ndarray:
# 这段代码定义了一个函数 get_windows ,用于生成图像的滑动窗口分割方案。它根据给定的图像尺寸、窗口大小、步长(通过窗口大小和间隙计算)以及图像覆盖率阈值,计算出符合条件的滑动窗口,并返回这些窗口的坐标范围。
# 定义了一个函数 get_windows ,接受以下参数:
# 1.im_size :一个元组 (h, w) ,表示图像的高度和宽度。
# 2.crop_sizes :一个元组,表示每个滑动窗口的大小,默认为 (1024,) 。
# 3.gaps :一个元组,表示每个滑动窗口之间的间隙大小,默认为 (200,) 。
# 4.im_rate_thr :一个浮点数,表示图像覆盖率阈值,默认为 0.6 。
# 5.eps :一个浮点数,用于处理数值精度问题,默认为 0.01 。
# 函数返回一个 NumPy 数组,表示符合条件的滑动窗口的坐标范围,格式为 [x_start, y_start, x_end, y_end] 。
def get_windows(
im_size: Tuple[int, int],
crop_sizes: Tuple[int, ...] = (1024,),
gaps: Tuple[int, ...] = (200,),
im_rate_thr: float = 0.6,
eps: float = 0.01,
) -> np.ndarray:
# 获取用于图像裁剪的滑动窗口坐标。
# 参数:
# im_size (Tuple[int, int]):原始图像尺寸,(H, W)。
# crop_sizes (Tuple[int, ...],可选):裁剪窗口的大小。
# gaps (Tuple[int, ...],可选):裁剪之间的间隙。
# im_rate_thr (float,可选):窗口面积除以图像面积的阈值。
# eps (float,可选):用于数学运算的 Epsilon 值。
# 返回:
# (np.ndarray):形状为 (N, 4) 的窗口坐标数组,每行包含 [x_start, y_start, x_stop, y_stop]。
"""
Get the coordinates of sliding windows for image cropping.
Args:
im_size (Tuple[int, int]): Original image size, (H, W).
crop_sizes (Tuple[int, ...], optional): Crop size of windows.
gaps (Tuple[int, ...], optional): Gap between crops.
im_rate_thr (float, optional): Threshold of windows areas divided by image areas.
eps (float, optional): Epsilon value for math operations.
Returns:
(np.ndarray): Array of window coordinates with shape (N, 4) where each row is [x_start, y_start, x_stop, y_stop].
"""
# 解包 im_size ,分别获取图像的高度 h 和宽度 w 。
h, w = im_size
# 初始化一个空列表 windows ,用于存储所有滑动窗口的坐标范围。
windows = []
# 遍历 crop_sizes 和 gaps ,分别获取每个滑动窗口的大小 crop_size 和间隙 gap 。
for crop_size, gap in zip(crop_sizes, gaps):
# 断言窗口大小必须大于间隙大小,否则抛出异常。
assert crop_size > gap, f"invalid crop_size gap pair [{crop_size} {gap}]" # crop_size 间隙对无效 [{crop_size} {gap}] 。
# 计算滑动窗口的步长,即窗口大小减去间隙。
step = crop_size - gap
# 计算在宽度方向上需要的滑动窗口数量: 如果图像宽度小于或等于窗口大小,则只需要一个窗口。 否则,计算需要的窗口数量,向上取整。
xn = 1 if w <= crop_size else ceil((w - crop_size) / step + 1)
# 计算每个滑动窗口在宽度方向上的起始坐标。
xs = [step * i for i in range(xn)]
# 如果最后一个窗口的结束位置超出图像宽度,则将其起始坐标调整为图像宽度减去窗口大小。
if len(xs) > 1 and xs[-1] + crop_size > w:
xs[-1] = w - crop_size
# 计算在高度方向上需要的滑动窗口数量,逻辑与宽度方向相同。
yn = 1 if h <= crop_size else ceil((h - crop_size) / step + 1)
# 计算每个滑动窗口在高度方向上的起始坐标。
ys = [step * i for i in range(yn)]
# 如果最后一个窗口的结束位置超出图像高度,则将其起始坐标调整为图像高度减去窗口大小。
if len(ys) > 1 and ys[-1] + crop_size > h:
ys[-1] = h - crop_size
# 使用 itertools.product 生成所有可能的滑动窗口的起始坐标 (x_start, y_start) ,并将其转换为 NumPy 数组。
start = np.array(list(itertools.product(xs, ys)), dtype=np.int64)
# 计算每个滑动窗口的结束坐标 (x_end, y_end) 。
stop = start + crop_size
# 将每个滑动窗口的起始和结束坐标拼接起来,形成 [x_start, y_start, x_end, y_end] 的格式,并将其添加到 windows 列表中。
windows.append(np.concatenate([start, stop], axis=1))
# 将所有滑动窗口的坐标范围合并为一个 NumPy 数组。
windows = np.concatenate(windows, axis=0)
# 复制 windows 数组,用于后续计算每个窗口中实际包含的图像区域。
im_in_wins = windows.copy()
# 将窗口的起始和结束坐标在宽度方向上裁剪到 [0, w] 范围内。
im_in_wins[:, 0::2] = np.clip(im_in_wins[:, 0::2], 0, w)
# 将窗口的起始和结束坐标在高度方向上裁剪到 [0, h] 范围内。
im_in_wins[:, 1::2] = np.clip(im_in_wins[:, 1::2], 0, h)
# 计算每个窗口中实际包含的图像区域的面积。
im_areas = (im_in_wins[:, 2] - im_in_wins[:, 0]) * (im_in_wins[:, 3] - im_in_wins[:, 1])
# 计算每个窗口的总面积。
win_areas = (windows[:, 2] - windows[:, 0]) * (windows[:, 3] - windows[:, 1])
# 计算每个窗口中图像区域的覆盖率,即图像区域面积与窗口总面积的比值。
im_rates = im_areas / win_areas
# 如果没有任何窗口的覆盖率超过阈值 im_rate_thr ,则执行以下操作:
if not (im_rates > im_rate_thr).any():
# 找到最大覆盖率 max_rate ,并将与最大覆盖率相差小于 eps 的窗口的覆盖率设置为 1,以确保至少有一个窗口被保留。
max_rate = im_rates.max()
im_rates[abs(im_rates - max_rate) < eps] = 1
# 返回覆盖率超过阈值 im_rate_thr 的滑动窗口的坐标范围。
return windows[im_rates > im_rate_thr]
# 这段代码实现了一个用于生成图像滑动窗口分割方案的函数。它根据图像尺寸、窗口大小、间隙大小和覆盖率阈值,计算出符合条件的滑动窗口。主要步骤包括: 根据窗口大小和间隙计算步长。 在图像的宽度和高度方向上计算滑动窗口的数量和起始坐标。 生成所有滑动窗口的坐标范围。 计算每个窗口中实际包含的图像区域的覆盖率。 筛选出覆盖率超过阈值的窗口。最终返回的滑动窗口坐标范围可用于图像分割、目标检测等任务,确保每个窗口中包含足够的图像内容。
5.def get_window_obj(anno: Dict[str, Any], windows: np.ndarray, iof_thr: float = 0.7) -> List[np.ndarray]:
# 这段代码定义了一个函数 get_window_obj ,用于将标注信息(如目标检测中的边界框)分配到不同的滑动窗口中。它通过计算标注框与滑动窗口的交并比(IOF),筛选出与每个窗口的 IOF 大于阈值的标注框,并返回每个窗口对应的标注信息。
# 定义了一个函数 get_window_obj ,接受以下参数:
# 1.anno :一个字典,包含图像的标注信息,例如图像的原始尺寸和标注框。
# 2.windows :一个 NumPy 数组,表示滑动窗口的坐标范围,格式为 [x_start, y_start, x_end, y_end] 。
# 3.iof_thr :一个浮点数,默认值为 0.7 ,表示交并比阈值。
# 函数返回一个列表,列表中的每个元素是一个 NumPy 数组,表示每个滑动窗口对应的标注信息。
def get_window_obj(anno: Dict[str, Any], windows: np.ndarray, iof_thr: float = 0.7) -> List[np.ndarray]:
# 根据 IoF 阈值获取每个窗口的对象。
"""Get objects for each window based on IoF threshold."""
# 从 anno 字典中提取图像的原始高度 h 和宽度 w 。
h, w = anno["ori_size"]
# 从 anno 字典中提取标注信息 label ,通常是一个 NumPy 数组,包含标注框的坐标和其他信息。
label = anno["label"]
# 判断标注信息是否为空。如果标注框的数量大于 0,则执行以下操作。
if len(label):
# 将标注框的宽度坐标(从第 2 列开始,每隔一列)乘以图像宽度 w ,将其从归一化坐标转换为绝对像素坐标。
label[:, 1::2] *= w
# 将标注框的高度坐标(从第 3 列开始,每隔一列)乘以图像高度 h ,将其从归一化坐标转换为绝对像素坐标。
label[:, 2::2] *= h
# 调用 bbox_iof 函数,计算标注框与每个滑动窗口的交并比(IOF)。 label[:, 1:] 表示标注框的坐标部分, windows 是滑动窗口的坐标范围。
iofs = bbox_iof(label[:, 1:], windows)
# Unnormalized and misaligned coordinates
# 遍历每个滑动窗口,筛选出与当前窗口的 IOF 大于阈值 iof_thr 的标注框。 使用列表推导式返回一个列表,列表中的每个元素是一个 NumPy 数组,表示每个滑动窗口对应的标注信息。
return [(label[iofs[:, i] >= iof_thr]) for i in range(len(windows))] # window_anns
# 如果标注信息为空(即 len(label) == 0 ),则执行以下操作。
else:
# 返回一个列表,列表中的每个元素是一个形状为 (0, 9) 的空 NumPy 数组,表示每个滑动窗口没有标注信息。 9 是标注框的坐标和其他信息的总维度。
return [np.zeros((0, 9), dtype=np.float32) for _ in range(len(windows))] # window_anns
# 这段代码实现了一个用于将标注信息分配到滑动窗口的函数。它通过以下步骤完成任务: 提取图像的原始尺寸和标注信息。 将标注框的归一化坐标转换为绝对像素坐标。 计算标注框与每个滑动窗口的交并比(IOF)。 筛选出与每个窗口的 IOF 大于阈值的标注框。 返回每个滑动窗口对应的标注信息。如果标注信息为空,则为每个窗口返回一个空数组。这种处理方式适用于目标检测任务中的滑动窗口策略,确保每个窗口的标注信息是独立的且符合要求。
6.def crop_and_save(anno: Dict[str, Any], windows: np.ndarray, window_objs: List[np.ndarray], im_dir: str, lb_dir: str, allow_background_images: bool = True,) -> None:
# 这段代码定义了一个函数 crop_and_save ,用于根据滑动窗口的坐标范围从图像中裁剪出子图像(patch),并将对应的标注信息保存到指定的目录中。它还支持保存背景图像(即没有标注信息的图像),并根据需要调整标注信息的坐标。
# 定义了一个函数 crop_and_save ,接受以下参数:
# 1.anno :一个字典,包含图像的标注信息,例如图像路径和标注框。
# 2.windows :一个 NumPy 数组,表示滑动窗口的坐标范围,格式为 [x_start, y_start, x_end, y_end] 。
# 3.window_objs :一个列表,包含每个滑动窗口对应的标注信息。
# 4.im_dir :一个字符串,表示裁剪后的子图像保存路径。
# 5.lb_dir :一个字符串,表示标注信息保存路径。
# 6.allow_background_images :一个布尔值,默认为 True ,表示是否允许保存没有标注信息的背景图像。
# 函数没有返回值,其主要功能是裁剪图像并保存标注信息。
def crop_and_save(
anno: Dict[str, Any],
windows: np.ndarray,
window_objs: List[np.ndarray],
im_dir: str,
lb_dir: str,
allow_background_images: bool = True,
) -> None:
# 裁剪图像并为每个窗口保存新的标签。
# 注意:
# DOTA 数据集的目录结构如下:
# - data_root
# - images
# - train
# - val
# - labels
# - train
# - val
"""
Crop images and save new labels for each window.
Args:
anno (Dict[str, Any]): Annotation dict, including 'filepath', 'label', 'ori_size' as its keys.
windows (np.ndarray): Array of windows coordinates with shape (N, 4).
window_objs (List[np.ndarray]): A list of labels inside each window.
im_dir (str): The output directory path of images.
lb_dir (str): The output directory path of labels.
allow_background_images (bool, optional): Whether to include background images without labels.
Notes:
The directory structure assumed for the DOTA dataset:
- data_root
- images
- train
- val
- labels
- train
- val
"""
# 使用 OpenCV 的 cv2.imread 函数读取图像文件,图像路径从 anno 字典中的 "filepath" 键获取。
im = cv2.imread(anno["filepath"])
# 使用 pathlib.Path 获取图像文件的文件名(不包含扩展名),用于后续生成裁剪后的图像和标注文件的名称。
name = Path(anno["filepath"]).stem
# 遍历每个滑动窗口的坐标范围。
for i, window in enumerate(windows):
# 将当前滑动窗口的坐标范围解包为起始和结束坐标。
x_start, y_start, x_stop, y_stop = window.tolist()
# 根据滑动窗口的坐标范围生成裁剪后的图像和标注文件的名称。格式为 [原文件名]__[窗口宽度]__[x_start]__[y_start] 。
new_name = f"{name}__{x_stop - x_start}__{x_start}___{y_start}"
# 使用 NumPy 的切片操作从原始图像中裁剪出当前滑动窗口对应的子图像。
patch_im = im[y_start:y_stop, x_start:x_stop]
# 获取裁剪后的子图像的高度 ph 和宽度 pw 。
ph, pw = patch_im.shape[:2]
# 获取当前滑动窗口对应的标注信息。
label = window_objs[i]
# 如果当前滑动窗口有标注信息( len(label) > 0 ),或者允许保存背景图像( allow_background_images == True ),则将裁剪后的子图像保存到指定的目录中。
if len(label) or allow_background_images:
cv2.imwrite(str(Path(im_dir) / f"{new_name}.jpg"), patch_im)
# 如果当前滑动窗口有标注信息,则执行以下操作。
if len(label):
# 将标注框的坐标从全局坐标系转换为裁剪后的子图像坐标系。具体来说,标注框的 x 坐标减去 x_start , y 坐标减去 y_start 。
label[:, 1::2] -= x_start
label[:, 2::2] -= y_start
# 将标注框的坐标从绝对像素坐标转换为归一化坐标,分别除以裁剪后的子图像的宽度 pw 和高度 ph 。
label[:, 1::2] /= pw
label[:, 2::2] /= ph
# 打开一个文件用于保存当前滑动窗口的标注信息,文件路径为 lb_dir ,文件名为 [new_name].txt 。
with open(Path(lb_dir) / f"{new_name}.txt", "w", encoding="utf-8") as f:
# 遍历当前滑动窗口的标注信息,将每个标注框的类别和坐标格式化为字符串,并写入文件。类别标签 lb[0] 转换为整数,坐标值保留 6 位有效数字。
for lb in label:
formatted_coords = [f"{coord:.6g}" for coord in lb[1:]]
f.write(f"{int(lb[0])} {' '.join(formatted_coords)}\n")
# 这段代码实现了一个用于裁剪图像并保存标注信息的函数。它通过以下步骤完成任务: 读取原始图像文件。 遍历每个滑动窗口的坐标范围。 从原始图像中裁剪出子图像。 如果滑动窗口有标注信息或允许保存背景图像,则保存裁剪后的子图像。 如果滑动窗口有标注信息,将标注框的坐标从全局坐标系转换为裁剪后的子图像坐标系,并保存标注信息到文件中。这种处理方式适用于目标检测任务中的图像预处理,特别是在处理大图像时,通过滑动窗口策略将大图像分割为多个子图像,并为每个子图像生成对应的标注信息。
7.def split_images_and_labels(data_root: str, save_dir: str, split: str = "train", crop_sizes: Tuple[int, ...] = (1024,), gaps: Tuple[int, ...] = (200,),) -> None:
# 这段代码定义了一个函数 split_images_and_labels ,用于将大图像分割成小块(crop),并根据滑动窗口的坐标范围裁剪图像和对应的标注信息,最后将裁剪后的图像和标注信息保存到指定的目录中。它适用于目标检测任务中的数据预处理,特别是在处理大图像时。
# 定义了一个函数 split_images_and_labels ,接受以下参数:
# 1.data_root :一个字符串,表示原始数据集的根目录路径。
# 2.save_dir :一个字符串,表示保存裁剪后的图像和标注信息的目标目录。
# 3.split :一个字符串,默认值为 "train" ,表示数据集的划分(训练集或验证集)。
# 4.crop_sizes :一个元组,默认值为 (1024,) ,表示滑动窗口的大小。
# 5.gaps :一个元组,默认值为 (200,) ,表示滑动窗口之间的间隙大小。
# 函数没有返回值,其主要功能是裁剪图像并保存标注信息。
def split_images_and_labels(
data_root: str,
save_dir: str,
split: str = "train",
crop_sizes: Tuple[int, ...] = (1024,),
gaps: Tuple[int, ...] = (200,),
) -> None:
# 根据给定的数据集拆分图像和标签。
# 注意:
# DOTA 数据集的目录结构假设如下:
# - data_root
# - images
# - split
# - labels
# - split
# 输出目录结构如下:
# - save_dir
# - images
# - split
# - labels
# - split
"""
Split both images and labels for a given dataset split.
Args:
data_root (str): Root directory of the dataset.
save_dir (str): Directory to save the split dataset.
split (str, optional): The split data set, could be 'train' or 'val'.
crop_sizes (Tuple[int, ...], optional): Tuple of crop sizes.
gaps (Tuple[int, ...], optional): Tuple of gaps between crops.
Notes:
The directory structure assumed for the DOTA dataset:
- data_root
- images
- split
- labels
- split
and the output directory structure is:
- save_dir
- images
- split
- labels
- split
"""
# 使用 pathlib.Path 构造保存裁剪后图像的目标目录路径,路径格式为 [save_dir]/images/[split] 。
im_dir = Path(save_dir) / "images" / split
# 如果目标目录不存在,则创建该目录。 parents=True 表示创建所有必要的父目录, exist_ok=True 表示如果目录已存在,不会抛出异常。
im_dir.mkdir(parents=True, exist_ok=True)
# 使用 pathlib.Path 构造保存标注信息的目标目录路径,路径格式为 [save_dir]/labels/[split] 。
lb_dir = Path(save_dir) / "labels" / split
# 如果目标目录不存在,则创建该目录。
lb_dir.mkdir(parents=True, exist_ok=True)
# 调用 load_yolo_dota 函数,加载指定数据集划分( split )的标注信息。 load_yolo_dota 函数返回一个列表,每个元素是一个字典,包含图像的原始尺寸、标注信息和文件路径。
annos = load_yolo_dota(data_root, split=split)
# 使用 TQDM 来显示进度条,遍历所有标注信息。 total=len(annos) 表示总进度, desc=split 表示进度条的描述信息。
for anno in TQDM(annos, total=len(annos), desc=split):
# 调用 get_windows 函数,根据图像的原始尺寸、滑动窗口大小和间隙,计算滑动窗口的坐标范围。
windows = get_windows(anno["ori_size"], crop_sizes, gaps)
# 调用 get_window_obj 函数,将标注信息分配到每个滑动窗口中,返回每个窗口对应的标注信息。
window_objs = get_window_obj(anno, windows)
# 调用 crop_and_save 函数,根据滑动窗口的坐标范围裁剪图像,并保存裁剪后的图像和标注信息到指定目录。
crop_and_save(anno, windows, window_objs, str(im_dir), str(lb_dir))
# 这段代码实现了一个用于将大图像分割成小块并保存裁剪后的图像和标注信息的函数。它通过以下步骤完成任务: 创建保存裁剪后图像和标注信息的目标目录。 加载指定数据集划分的标注信息。 遍历每张图像的标注信息: 计算滑动窗口的坐标范围。 将标注信息分配到每个滑动窗口。 裁剪图像并保存裁剪后的图像和标注信息到目标目录。这种处理方式适用于目标检测任务中的数据预处理,特别是在处理大图像时,通过滑动窗口策略将大图像分割为多个小块,并为每个小块生成对应的标注信息,便于后续的模型训练和验证。
8.def split_trainval(data_root: str, save_dir: str, crop_size: int = 1024, gap: int = 200, rates: Tuple[float, ...] = (1.0,)) -> None:
# 这段代码定义了一个函数 split_trainval ,用于对整个数据集进行训练集和验证集的分割,并对每个数据集划分执行图像分割操作。它通过调整滑动窗口的大小和间隙,根据给定的比例因子( rates )生成不同尺寸的滑动窗口,并调用 split_images_and_labels 函数来完成图像的裁剪和标注信息的保存。
# 定义了一个函数 split_trainval ,接受以下参数:
# 1.data_root :一个字符串,表示原始数据集的根目录路径。
# 2.save_dir :一个字符串,表示保存裁剪后图像和标注信息的目标目录。
# 3.crop_size :一个整数,默认值为 1024 ,表示基本滑动窗口的大小。
# 4.gap :一个整数,默认值为 200 ,表示基本滑动窗口之间的间隙大小。
# 5.rates :一个元组,默认值为 (1.0,) ,表示滑动窗口大小和间隙的缩放比例因子。
# 函数没有返回值,其主要功能是对训练集和验证集分别进行图像分割和标注信息保存。
def split_trainval(
data_root: str, save_dir: str, crop_size: int = 1024, gap: int = 200, rates: Tuple[float, ...] = (1.0,)
) -> None:
# 使用多种缩放比例对 DOTA 数据集进行训练集和验证集拆分。
# 备注:
# DOTA 数据集的目录结构如下:
# - data_root
# - images
# - train
# - val
# - labels
# - train
# - val
# 输出目录结构如下:
# - save_dir
# - images
# - train
# - val
# - labels
# - train
# - val
"""
Split train and val sets of DOTA dataset with multiple scaling rates.
Args:
data_root (str): Root directory of the dataset.
save_dir (str): Directory to save the split dataset.
crop_size (int, optional): Base crop size.
gap (int, optional): Base gap between crops.
rates (Tuple[float, ...], optional): Scaling rates for crop_size and gap.
Notes:
The directory structure assumed for the DOTA dataset:
- data_root
- images
- train
- val
- labels
- train
- val
and the output directory structure is:
- save_dir
- images
- train
- val
- labels
- train
- val
"""
# 初始化两个空列表 crop_sizes 和 gaps ,用于存储根据比例因子调整后的滑动窗口大小和间隙大小。
crop_sizes, gaps = [], []
# 遍历比例因子 rates 。
for r in rates:
# 根据比例因子 r 缩放基本滑动窗口大小 crop_size ,并将结果添加到 crop_sizes 列表中。 int() 用于将结果转换为整数。
crop_sizes.append(int(crop_size / r))
# 根据比例因子 r 缩放基本滑动窗口间隙 gap ,并将结果添加到 gaps 列表中。
gaps.append(int(gap / r))
# 遍历数据集划分,分别处理训练集( "train" )和验证集( "val" )。
for split in ["train", "val"]:
# 调用 split_images_and_labels 函数,对当前数据集划分( split )执行图像分割操作。传递调整后的滑动窗口大小 crop_sizes 和间隙 gaps ,以及数据集的根目录 data_root 和保存目录 save_dir 。
split_images_and_labels(data_root, save_dir, split, crop_sizes, gaps)
# 这段代码实现了一个用于对数据集进行训练集和验证集分割,并对每个划分执行图像分割操作的函数。它通过以下步骤完成任务: 根据给定的比例因子( rates ),调整基本滑动窗口的大小和间隙。 遍历训练集和验证集: 调用 split_images_and_labels 函数,对每个数据集划分执行图像分割操作。 裁剪图像并保存裁剪后的图像和标注信息到指定目录。这种处理方式适用于目标检测任务中的数据预处理,特别是在处理大图像数据集时,通过滑动窗口策略将大图像分割为多个小块,并为每个小块生成对应的标注信息,便于后续的模型训练和验证。此外,通过调整滑动窗口的大小和间隙,可以生成不同尺度的图像块,从而提高模型对不同尺寸目标的检测能力。
9.def split_test(data_root: str, save_dir: str, crop_size: int = 1024, gap: int = 200, rates: Tuple[float, ...] = (1.0,)) -> None:
# 这段代码定义了一个函数 split_test ,用于处理测试集图像,将每张大图像分割成多个小块(crop),并保存这些小块到指定的目录中。它通过调整滑动窗口的大小和间隙,根据给定的比例因子( rates )生成不同尺寸的滑动窗口,并使用 OpenCV 裁剪图像。
# 定义了一个函数 split_test ,接受以下参数:
# 1.data_root :一个字符串,表示测试集图像的根目录路径。
# 2.save_dir :一个字符串,表示保存裁剪后图像的目标目录。
# 3.crop_size :一个整数,默认值为 1024 ,表示基本滑动窗口的大小。
# 4.gap :一个整数,默认值为 200 ,表示基本滑动窗口之间的间隙大小。
# 5.rates :一个元组,默认值为 (1.0,) ,表示滑动窗口大小和间隙的缩放比例因子。
# 函数没有返回值,其主要功能是裁剪测试集图像并保存裁剪后的图像。
def split_test(
data_root: str, save_dir: str, crop_size: int = 1024, gap: int = 200, rates: Tuple[float, ...] = (1.0,)
) -> None:
# DOTA 数据集的拆分测试集,标签不包含在该集合中。
# 注:
# DOTA 数据集的目录结构假设如下:
# - data_root
# - images
# - test
# 输出目录结构如下:
# - save_dir
# - images
# - test
"""
Split test set of DOTA dataset, labels are not included within this set.
Args:
data_root (str): Root directory of the dataset.
save_dir (str): Directory to save the split dataset.
crop_size (int, optional): Base crop size.
gap (int, optional): Base gap between crops.
rates (Tuple[float, ...], optional): Scaling rates for crop_size and gap.
Notes:
The directory structure assumed for the DOTA dataset:
- data_root
- images
- test
and the output directory structure is:
- save_dir
- images
- test
"""
# 初始化两个空列表 crop_sizes 和 gaps ,用于存储根据比例因子调整后的滑动窗口大小和间隙大小。
crop_sizes, gaps = [], []
# 遍历比例因子 rates 。
for r in rates:
# 根据比例因子 r 缩放基本滑动窗口大小 crop_size ,并将结果添加到 crop_sizes 列表中。 int() 用于将结果转换为整数。
crop_sizes.append(int(crop_size / r))
# 根据比例因子 r 缩放基本滑动窗口间隙 gap ,并将结果添加到 gaps 列表中。
gaps.append(int(gap / r))
# 使用 pathlib.Path 构造保存裁剪后图像的目标目录路径,路径格式为 [save_dir]/images/test 。
save_dir = Path(save_dir) / "images" / "test"
# 如果目标目录不存在,则创建该目录。 parents=True 表示创建所有必要的父目录, exist_ok=True 表示如果目录已存在,不会抛出异常。
save_dir.mkdir(parents=True, exist_ok=True)
# 使用 pathlib.Path 构造测试集图像的目录路径,路径格式为 [data_root]/images/test 。
im_dir = Path(data_root) / "images" / "test"
# 断言测试集图像目录存在,否则抛出异常。这确保了路径正确。
assert im_dir.exists(), f"Can't find {im_dir}, please check your data root." # 找不到 {im_dir},请检查您的数据根目录。
# 使用 glob 函数获取测试集图像目录中所有文件的路径。
im_files = glob(str(im_dir / "*"))
# 使用 TQDM 来显示进度条,遍历所有测试集图像文件。 total=len(im_files) 表示总进度, desc="test" 表示进度条的描述信息。
for im_file in TQDM(im_files, total=len(im_files), desc="test"):
# 使用 PIL.Image.open 打开图像文件,并调用 exif_size 函数获取图像的宽度 w 和高度 h 。
w, h = exif_size(Image.open(im_file))
# 调用 get_windows 函数,根据图像的尺寸、滑动窗口大小和间隙,计算滑动窗口的坐标范围。
windows = get_windows((h, w), crop_sizes=crop_sizes, gaps=gaps)
# 使用 OpenCV 的 cv2.imread 函数读取当前图像文件。
im = cv2.imread(im_file)
# 使用 pathlib.Path 获取图像文件的文件名(不包含扩展名),用于后续生成裁剪后的图像文件名。
name = Path(im_file).stem
# 遍历每个滑动窗口的坐标范围。
for window in windows:
# 将当前滑动窗口的坐标范围解包为起始和结束坐标。
x_start, y_start, x_stop, y_stop = window.tolist()
# 根据滑动窗口的坐标范围生成裁剪后的图像文件名。格式为 [原文件名]__[窗口宽度]__[x_start]__[y_start] 。
new_name = f"{name}__{x_stop - x_start}__{x_start}___{y_start}"
# 使用 NumPy 的切片操作从原始图像中裁剪出当前滑动窗口对应的子图像。
patch_im = im[y_start:y_stop, x_start:x_stop]
# 使用 OpenCV 的 cv2.imwrite 函数将裁剪后的子图像保存到目标目录中。
cv2.imwrite(str(save_dir / f"{new_name}.jpg"), patch_im)
# 这段代码实现了一个用于处理测试集图像的函数,通过以下步骤完成任务: 根据给定的比例因子( rates ),调整基本滑动窗口的大小和间隙。 创建保存裁剪后图像的目标目录。 遍历测试集图像目录中的所有图像文件: 获取图像的尺寸。 计算滑动窗口的坐标范围。 裁剪图像并保存裁剪后的子图像到目标目录。这种处理方式适用于目标检测任务中的测试集预处理,特别是在处理大图像时,通过滑动窗口策略将大图像分割为多个小块,便于后续的模型推理和结果评估。
10.if __name__ == "__main__":
# 这是一个 Python 中常见的入口点判断。当脚本被直接运行时(而不是被其他模块导入时), __name__ 的值为 "__main__" ,此时会执行缩进代码块中的内容。
if __name__ == "__main__":
# 调用 split_trainval 函数,对数据集进行训练集和验证集的分割,并对每个划分执行图像分割操作。
# data_root="DOTAv2" :指定原始数据集的根目录为 "DOTAv2" 。
# save_dir="DOTAv2-split" :指定保存裁剪后图像和标注信息的目标目录为 "DOTAv2-split" 。
# 这一步会处理训练集和验证集的图像,生成不同尺寸的滑动窗口,并保存裁剪后的图像和标注信息到目标目录。
split_trainval(data_root="DOTAv2", save_dir="DOTAv2-split")
# 调用 split_test 函数,处理测试集图像,将每张大图像分割成多个小块,并保存这些小块到指定的目录中。
# data_root="DOTAv2" :指定测试集图像的根目录为 "DOTAv2" 。
# save_dir="DOTAv2-split" :指定保存裁剪后图像的目标目录为 "DOTAv2-split" 。
# 这一步会处理测试集图像,生成不同尺寸的滑动窗口,并保存裁剪后的图像到目标目录。
split_test(data_root="DOTAv2", save_dir="DOTAv2-split")