diff --git a/src/utils/image_splitter.py b/src/utils/image_splitter.py new file mode 100644 index 0000000..df1573b --- /dev/null +++ b/src/utils/image_splitter.py @@ -0,0 +1,226 @@ +import numpy as np + +from pathlib import Path +from tifffile import imread, imwrite +from shapely.geometry import LineString +from copy import deepcopy + +# debug +from src.utils.image import Image +from show_yolo_seg import draw_annotations +import pylab as plt +import cv2 + + +class Label: + def __init__(self, yolo_annotation: str): + class_id, bbox, polygon = self.parse_yolo_annotation(yolo_annotation) + self.class_id = class_id + self.bbox = bbox + self.polygon = polygon + + def parse_yolo_annotation(self, yolo_annotation: str): + 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 + ) + return class_id, bbox, polygon + + def offset_label( + self, distance: float = 3.0, cap_style: int = 2, join_style: int = 2 + ): + if self.polygon is None: + self.bbox = np.array( + [ + self.bbox[0] - distance if self.bbox[0] - distance > 0 else 0, + self.bbox[1] - distance if self.bbox[1] - distance > 0 else 0, + self.bbox[2] + distance if self.bbox[2] + distance < 1 else 1, + self.bbox[3] + distance if self.bbox[3] + distance < 1 else 1, + ], + dtype=np.float32, + ) + return self.bbox + + line = LineString(self.polygon) + # 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 + ) + + return self.bbox, self.polygon + + def translate(self, x, y, scale_x, scale_y): + self.bbox[0] -= x + self.bbox[0] *= scale_x + self.bbox[1] -= y + self.bbox[1] *= scale_y + self.bbox[2] *= scale_x + self.bbox[3] *= scale_y + if self.polygon is not None: + self.polygon[:, 0] -= x + self.polygon[:, 0] *= scale_x + self.polygon[:, 1] -= y + self.polygon[:, 1] *= scale_y + + def in_range(self, hrange, wrange): + xc, yc, h, w = self.bbox + x1 = xc - w / 2 + y1 = yc - h / 2 + x2 = xc + w / 2 + y2 = yc + h / 2 + truth_val = ( + xc >= wrange[0] + and x1 <= wrange[1] + and x2 >= wrange[0] + and x2 <= wrange[1] + and y1 >= hrange[0] + and y1 <= hrange[1] + and y2 >= hrange[0] + and y2 <= hrange[1] + ) + + print(x1, x2, wrange, y1, y2, hrange, truth_val) + return truth_val + + def to_string(self, bbox: list = None, polygon: list = None): + if bbox is None: + bbox = self.bbox + if polygon is None: + polygon = self.polygon + coords = " ".join([f"{x:.6f}" for x in self.bbox]) + if self.polygon is not None: + coords += " " + " ".join([f"{x:.6f} {y:.6f}" for x, y in self.polygon]) + return f"{self.class_id} {coords}" + + def __str__(self): + return f"Class: {self.class_id}, BBox: {self.bbox}, Polygon: {self.polygon}" + + +class YoloLabelReader: + def __init__(self, label_path: Path): + self.label_path = label_path + self.labels = self._read_labels() + + def _read_labels(self): + with open(self.label_path, "r") as f: + labels = [Label(line) for line in f.readlines()] + + return labels + + def get_labels(self, hrange, wrange): + """hrange and wrange are tuples of (start, end) normalized to [0, 1]""" + labels = [] + # print(hrange, wrange) + for lbl in self.labels: + # print(lbl) + if lbl.in_range(hrange, wrange): + labels.append(lbl) + return labels if len(labels) > 0 else None + + +class ImageSplitter: + def __init__(self, image_path: Path, label_path: Path): + self.image = imread(image_path) + self.image_path = image_path + self.label_path = label_path + self.labels = YoloLabelReader(label_path) + + def split(self, patch_size: tuple = (2, 2)): + """Split image into patches of size patch_size""" + hstep, wstep = ( + self.image.shape[0] // patch_size[0], + self.image.shape[1] // patch_size[1], + ) + h, w = self.image.shape[:2] + + for i in range(patch_size[0]): + for j in range(patch_size[1]): + tile_reference = f"i{i}j{j}" + 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 + ] + print(id(labels)) + + if labels is not None: + print(hrange[0], wrange[0]) + for l in labels: + print(l.bbox) + [l.translate(wrange[0], hrange[0], 2, 2) for l in labels] + print("translated") + for l in labels: + print(l.bbox) + + # print(labels) + yield tile_reference, tile, labels + + +def main(args): + + 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): + print() + 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") + 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 __name__ == "__main__": + import argparse + + parser = argparse.ArgumentParser() + parser.add_argument("-i", "--input", type=Path) + parser.add_argument("-o", "--output", type=Path) + parser.add_argument( + "-p", + "--patch-size", + nargs=2, + type=int, + default=[2, 2], + help="Number of patches along height and width, rows and columns, respectively", + ) + args = parser.parse_args() + + main(args) diff --git a/src/utils/show_yolo_seg.py b/src/utils/show_yolo_seg.py new file mode 120000 index 0000000..ba41988 --- /dev/null +++ b/src/utils/show_yolo_seg.py @@ -0,0 +1 @@ +../../tests/show_yolo_seg.py \ No newline at end of file