adding files

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)

src/utils/show_yolo_seg.py

@@ -0,0 +1 @@
+../../tests/show_yolo_seg.py