From 510eabfa942b9f3bab513c9fe877236971d2df8d Mon Sep 17 00:00:00 2001 From: Martin Laasmaa Date: Mon, 5 Jan 2026 13:56:57 +0200 Subject: [PATCH] Adding splitter method --- src/utils/image_splitter.py | 196 ++++++++++++++++++++++++++++-------- tests/show_yolo_seg.py | 10 +- 2 files changed, 157 insertions(+), 49 deletions(-) diff --git a/src/utils/image_splitter.py b/src/utils/image_splitter.py index df1573b..3dacf6e 100644 --- a/src/utils/image_splitter.py +++ b/src/utils/image_splitter.py @@ -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,48 +189,119 @@ 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 - plt.figure(figsize=(10, 10 * tile.shape[0] / tile.shape[1])) - if labels is None: - plt.imshow(tile, cmap="gray") + debug = False + if debug: + plt.figure(figsize=(10, 10 * tile.shape[0] / tile.shape[1])) + if labels is None: + plt.imshow(tile, cmap="gray") + plt.axis("off") + plt.title(f"{image_path.name} ({tile_reference})") + plt.show() + continue + + print(labels[0].bbox) + # Draw annotations + out = draw_annotations( + cv2.cvtColor((tile / tile.max() * 255).astype(np.uint8), cv2.COLOR_GRAY2BGR), + [l.to_string() for l in labels], + alpha=0.1, + ) + + # Convert BGR -> RGB for matplotlib display + out_rgb = cv2.cvtColor(out, cv2.COLOR_BGR2RGB) + plt.imshow(out_rgb) plt.axis("off") plt.title(f"{image_path.name} ({tile_reference})") plt.show() - continue - - print(labels[0].bbox) - # Draw annotations - out = draw_annotations( - cv2.cvtColor( - (tile / tile.max() * 255).astype(np.uint8), cv2.COLOR_GRAY2BGR - ), - [l.to_string() for l in labels], - alpha=0.1, - ) - - # Convert BGR -> RGB for matplotlib display - out_rgb = cv2.cvtColor(out, cv2.COLOR_BGR2RGB) - plt.imshow(out_rgb) - 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) diff --git a/tests/show_yolo_seg.py b/tests/show_yolo_seg.py index a698dee..e1deb8c 100644 --- a/tests/show_yolo_seg.py +++ b/tests/show_yolo_seg.py @@ -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)