adding files
This commit is contained in:
226
src/utils/image_splitter.py
Normal file
226
src/utils/image_splitter.py
Normal file
@@ -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)
|
||||
1
src/utils/show_yolo_seg.py
Symbolic link
1
src/utils/show_yolo_seg.py
Symbolic link
@@ -0,0 +1 @@
|
||||
../../tests/show_yolo_seg.py
|
||||
Reference in New Issue
Block a user