martin/object-segmentation
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Adding splitter method

Browse Source
This commit is contained in:
Martin Laasmaa
2026-01-05 13:56:57 +02:00
parent 395d263900
commit 510eabfa94
2 changed files with 157 additions and 49 deletions
Download Patch File Download Diff File Expand all files Collapse all files

160
src/utils/image_splitter.py
View File

@@ -23,15 +23,16 @@ class Label:
class_id, *coords = yolo_annotation.split()
class_id = int(class_id)
bbox = np.array(coords[:4], dtype=np.float32)
polygon = (
np.array(coords[4:], dtype=np.float32).reshape(-1, 2)
if len(coords) > 4
else None
)
polygon = np.array(coords[4:], dtype=np.float32).reshape(-1, 2) if len(coords) > 4 else None
return class_id, bbox, polygon
def offset_label(
self, distance: float = 3.0, cap_style: int = 2, join_style: int = 2
self,
img_w,
img_h,
distance: float = 3.0,
cap_style: int = 2,
join_style: int = 2,
):
if self.polygon is None:
self.bbox = np.array(
@@ -45,15 +46,32 @@ class Label:
)
return self.bbox
line = LineString(self.polygon)
def coords_are_normalized(coords):
# If every coordinate is between 0 and 1 (inclusive-ish), assume normalized
print(coords)
# if not coords:
# return False
return all(max(coords.flatten)) <= 1.001
def poly_to_pts(coords, img_w, img_h):
# coords: [x1 y1 x2 y2 ...] either normalized or absolute
# if coords_are_normalized(coords):
coords = [coords[i] * (img_w if i % 2 == 0 else img_h) for i in range(len(coords))]
pts = np.array(coords, dtype=np.int32).reshape(-1, 2)
return pts
pts = poly_to_pts(self.polygon, img_w, img_h)
line = LineString(pts)
# Buffer distance in pixels
buffered = line.buffer(
distance=distance, cap_style=cap_style, join_style=join_style
)
self.polygon = np.array(buffered.exterior.coords, dtype=np.int32)
self.bbox = np.array(
[np.min(self.polygon[:, 0]), np.min(self.polygon[:, 1])], dtype=np.int32
)
buffered = line.buffer(distance=distance, cap_style=cap_style, join_style=join_style)
self.polygon = np.array(buffered.exterior.coords, dtype=np.float32)
xmn, ymn = self.polygon.min(axis=0)
xmx, ymx = self.polygon.max(axis=0)
xc = (xmn + xmx) / 2
yc = (ymn + ymx) / 2
bw = xmx - xmn
bh = ymx - ymn
self.bbox = np.array([xc, yc, bw, bh], dtype=np.float32)
return self.bbox, self.polygon
@@ -125,6 +143,15 @@ class YoloLabelReader:
labels.append(lbl)
return labels if len(labels) > 0 else None
def __get_item__(self, index):
return self.labels[index]
def __len__(self):
return len(self.labels)
def __iter__(self):
return iter(self.labels)
class ImageSplitter:
def __init__(self, image_path: Path, label_path: Path):
@@ -133,7 +160,7 @@ class ImageSplitter:
self.label_path = label_path
self.labels = YoloLabelReader(label_path)
def split(self, patch_size: tuple = (2, 2)):
def split_into_tiles(self, patch_size: tuple = (2, 2)):
"""Split image into patches of size patch_size"""
hstep, wstep = (
self.image.shape[0] // patch_size[0],
@@ -147,9 +174,7 @@ class ImageSplitter:
hrange = (i * hstep / h, (i + 1) * hstep / h)
wrange = (j * wstep / w, (j + 1) * wstep / w)
labels = deepcopy(self.labels.get_labels(hrange, wrange))
tile = self.image[
i * hstep : (i + 1) * hstep, j * wstep : (j + 1) * wstep
]
tile = self.image[i * hstep : (i + 1) * hstep, j * wstep : (j + 1) * wstep]
print(id(labels))
if labels is not None:
@@ -164,21 +189,88 @@ class ImageSplitter:
# print(labels)
yield tile_reference, tile, labels
def split_respective_to_label(self, padding: int = 67):
if self.labels is None:
raise ValueError("No labels found. Only images having labels can be split.")
for i, label in enumerate(self.labels):
tile_reference = f"_lbl-{i+1:02d}"
# print(label.bbox)
xc_norm, yc_norm, h_norm, w_norm = label.bbox # normalized coords
xc, yc, h, w = [
int(np.round(f))
for f in [
xc_norm * self.image.shape[1],
yc_norm * self.image.shape[0],
h_norm * self.image.shape[0],
w_norm * self.image.shape[1],
]
] # image coords
# print("img coords:", xc, yc, h, w)
pad_xneg = padding + 1 # int(w / 2) + padding
pad_xpos = padding # int(w / 2) + padding
pad_yneg = padding + 1 # int(h / 2) + padding
pad_ypos = padding # int(h / 2) + padding
if xc - pad_xneg < 0:
pad_xneg = xc
if pad_xpos + xc > self.image.shape[1]:
pad_xpos = self.image.shape[1] - xc
if yc - pad_yneg < 0:
pad_yneg = yc
if pad_ypos + yc > self.image.shape[0]:
pad_ypos = self.image.shape[0] - yc
# print("pads:", pad_xneg, pad_xpos, pad_yneg, pad_ypos)
tile = self.image[
yc - pad_yneg : yc + pad_ypos,
xc - pad_xneg : xc + pad_xpos,
]
ny, nx = tile.shape
x_offset = pad_xneg
y_offset = pad_yneg
# print("tile shape:", tile.shape)
yolo_annotation = f"{label.class_id} {x_offset/nx} {y_offset/ny} {h/ny} {w/nx} " + " ".join(
[
f"{(x*self.image.shape[1]-(xc - x_offset))/nx:.6f} {(y*self.image.shape[0]-(yc-y_offset))/ny:.6f}"
for x, y in label.polygon
]
)
# print(yolo_annotation)
new_label = Label(yolo_annotation=yolo_annotation)
yield tile_reference, tile, [new_label]
def main(args):
if args.output:
args.output.mkdir(exist_ok=True, parents=True)
(args.output / "images").mkdir(exist_ok=True)
(args.output / "labels").mkdir(exist_ok=True)
for image_path in (args.input / "images").glob("*.tif"):
data = ImageSplitter(
image_path=image_path,
label_path=(args.input / "labels" / image_path.stem).with_suffix(".txt"),
)
for tile_reference, tile, labels in data.split(patch_size=args.patch_size):
if args.split_around_label:
data = data.split_respective_to_label(padding=args.padding)
else:
data = data.split_into_tiles(patch_size=args.patch_size)
for tile_reference, tile, labels in data:
print()
print(
tile_reference, tile.shape, labels
) # len(labels) if labels else None)
print(tile_reference, tile.shape, labels) # len(labels) if labels else None)
# { debug
debug = False
if debug:
plt.figure(figsize=(10, 10 * tile.shape[0] / tile.shape[1]))
if labels is None:
plt.imshow(tile, cmap="gray")
@@ -190,9 +282,7 @@ def main(args):
print(labels[0].bbox)
# Draw annotations
out = draw_annotations(
cv2.cvtColor(
(tile / tile.max() * 255).astype(np.uint8), cv2.COLOR_GRAY2BGR
),
cv2.cvtColor((tile / tile.max() * 255).astype(np.uint8), cv2.COLOR_GRAY2BGR),
[l.to_string() for l in labels],
alpha=0.1,
)
@@ -203,9 +293,15 @@ def main(args):
plt.axis("off")
plt.title(f"{image_path.name} ({tile_reference})")
plt.show()
# } debug
if args.output:
imwrite(args.output / "images" / f"{image_path.stem}_{tile_reference}.tif", tile)
with open(args.output / "labels" / f"{image_path.stem}_{tile_reference}.txt", "w") as f:
for label in labels:
label.offset_label(tile.shape[1], tile.shape[0])
f.write(label.to_string() + "\n")
if __name__ == "__main__":
import argparse
@@ -221,6 +317,18 @@ if __name__ == "__main__":
default=[2, 2],
help="Number of patches along height and width, rows and columns, respectively",
)
parser.add_argument(
"-sal",
"--split-around-label",
action="store_true",
help="If enabled, the image will be split around the label and for each label, a separate image will be created.",
)
parser.add_argument(
"--padding",
type=int,
default=67,
help="Padding around the label when splitting around the label.",
)
args = parser.parse_args()
main(args)

10
tests/show_yolo_seg.py
View File

@@ -91,11 +91,11 @@ def draw_annotations(img, labels, alpha=0.4, draw_bbox_for_poly=True):
cv2.rectangle(img, (x1, y1), (x2, y2), color, 1)
pts = poly_to_pts(coords[4:], w, h)
line = LineString(pts)
# Buffer distance in pixels
buffered = line.buffer(3, cap_style=2, join_style=2)
coords = np.array(buffered.exterior.coords, dtype=np.int32)
cv2.fillPoly(overlay, [coords], color=(255, 255, 255))
# line = LineString(pts)
# # Buffer distance in pixels
# buffered = line.buffer(3, cap_style=2, join_style=2)
# coords = np.array(buffered.exterior.coords, dtype=np.int32)
# cv2.fillPoly(overlay, [coords], color=(255, 255, 255))
# fill on overlay
cv2.fillPoly(overlay, [pts], color)