Adding splitter method

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 = (
+        polygon = np.array(coords[4:], dtype=np.float32).reshape(-1, 2) if len(coords) > 4 else None
-            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(
+        buffered = line.buffer(distance=distance, cap_style=cap_style, join_style=join_style)
-            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)
-        self.polygon = np.array(buffered.exterior.coords, dtype=np.int32)
+        xmx, ymx = self.polygon.max(axis=0)
-        self.bbox = np.array(
+        xc = (xmn + xmx) / 2
-            [np.min(self.polygon[:, 0]), np.min(self.polygon[:, 1])], dtype=np.int32
+        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[
+                tile = self.image[i * hstep : (i + 1) * hstep, j * wstep : (j + 1) * wstep]
-                    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(
+            print(tile_reference, tile.shape, labels)  # len(labels) if labels else None)
-                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(
+                    cv2.cvtColor((tile / tile.max() * 255).astype(np.uint8), cv2.COLOR_GRAY2BGR),
-                    (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)

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)
+            # line = LineString(pts)
-            # Buffer distance in pixels
+            # # Buffer distance in pixels
-            buffered = line.buffer(3, cap_style=2, join_style=2)
+            # buffered = line.buffer(3, cap_style=2, join_style=2)
-            coords = np.array(buffered.exterior.coords, dtype=np.int32)
+            # coords = np.array(buffered.exterior.coords, dtype=np.int32)
-            cv2.fillPoly(overlay, [coords], color=(255, 255, 255))
+            # cv2.fillPoly(overlay, [coords], color=(255, 255, 255))
 
             # fill on overlay
             cv2.fillPoly(overlay, [pts], color)