File size: 7,639 Bytes
bdb955e |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 |
import pickle
import torch
from torch.utils.data import Dataset
import torchvision
import torchvision.transforms as T
from PIL import Image
import cv2
import numpy as np
import matplotlib.pyplot as plt
import pandas as pd
import h5py
from tqdm.auto import tqdm
from collections import defaultdict
from pathlib import Path
import json
import os
from argparse import ArgumentParser
lines_classes = [
"Big rect. left bottom",
"Big rect. left main",
"Big rect. left top",
"Big rect. right bottom",
"Big rect. right main",
"Big rect. right top",
"Circle central",
"Circle left",
"Circle right",
# "Goal left crossbar",
# "Goal left post left ",
# "Goal left post right",
# "Goal right crossbar",
# "Goal right post left",
# "Goal right post right",
"Goal unknown",
"Line unknown",
"Middle line",
"Side line bottom",
"Side line left",
"Side line right",
"Side line top",
"Small rect. left bottom",
"Small rect. left main",
"Small rect. left top",
"Small rect. right bottom",
"Small rect. right main",
"Small rect. right top",
]
# RGB values
palette = {
"Big rect. left bottom": (127, 0, 0),
"Big rect. left main": (102, 102, 102),
"Big rect. left top": (0, 0, 127),
"Big rect. right bottom": (86, 32, 39),
"Big rect. right main": (48, 77, 0),
"Big rect. right top": (14, 97, 100),
"Circle central": (0, 0, 255),
"Circle left": (255, 127, 0),
"Circle right": (0, 255, 255),
# "Goal left crossbar": (255, 255, 200),
# "Goal left post left ": (165, 255, 0),
# "Goal left post right": (155, 119, 45),
# "Goal right crossbar": (86, 32, 139),
# "Goal right post left": (196, 120, 153),
# "Goal right post right": (166, 36, 52),
"Goal unknown": (0, 0, 0),
"Line unknown": (0, 0, 0),
"Middle line": (255, 255, 0),
"Side line bottom": (255, 0, 255),
"Side line left": (0, 255, 150),
"Side line right": (0, 230, 0),
"Side line top": (230, 0, 0),
"Small rect. left bottom": (0, 150, 255),
"Small rect. left main": (254, 173, 225),
"Small rect. left top": (87, 72, 39),
"Small rect. right bottom": (122, 0, 255),
"Small rect. right main": (255, 255, 255),
"Small rect. right top": (153, 23, 153),
}
def create_target_from_annotation(width, height, annotation, classes, linewidth=4):
"""Draw one-hot encoded segments according to the annotation.
Creates target that matches image size ([C+1]xHxW).
"""
annotation_abs = defaultdict(list)
# unnormalize every point in every class k
for k in annotation.keys():
if k not in lines_classes:
continue
start = annotation[k][0].copy()
end = annotation[k][-1].copy()
for annotation_point in annotation[k]:
tup = annotation_point.copy()
tup["x"] *= width
tup["x"] = int(tup["x"])
tup["y"] *= height
tup["y"] = int(tup["y"])
annotation_abs[k].append(tup)
# draw lines between annotated points for each segment
# offset class +1 such that no classes detected will end in argmax 0
# otherwise argmax 0 will be another class
classes_segments = np.zeros(shape=(len(classes) + 1, height, width))
for cls, points in annotation_abs.items():
class_segments = np.zeros(shape=(height, width, 3))
for start, end in zip(points, points[1:]):
startxy = (start["x"], start["y"])
endxy = [end["x"], end["y"]]
class_segments = cv2.line(
class_segments, startxy, endxy, (1, 1, 1), linewidth
)
classes_segments[classes.index(cls) + 1] = class_segments[:, :, 1]
classes_segments = torch.Tensor(classes_segments)
return classes_segments
class ExtremitiesDataset(Dataset):
def __init__(
self, root, split, annotations, filter_cam=None, extremities_prefix="", classes=lines_classes, palette=palette
):
self.data_root = Path(root)
self.split = split
self.annotations_path = annotations
if filter_cam is None:
files = os.listdir(self.data_root / self.split)
self.annotations = sorted([fn for fn in files if fn.endswith("json")])
self.images = sorted([fn for fn in files if fn.endswith("jpg")])
else:
df = pd.read_json(self.data_root / self.split / "match_info_cam_gt.json").T
df = df.loc[df.camera == filter_cam]
assert len(df.index) > 0
df["image_file"] = df.index
df = df.sort_values(by=["image_file"])
df["annotation_file"] = df["image_file"].apply(
lambda s: extremities_prefix + s.split(".jpg")[0] + ".json"
)
self.annotations = df["annotation_file"].tolist()
self.images = df["image_file"].tolist()
self.classes = classes
def __len__(self):
return len(self.images)
def __getitem__(self, idx):
# see https://learnopencv.com/pytorch-for-beginners-semantic-segmentation-using-torchvision/
impath = self.data_root / self.split / self.images[idx]
annotation_path = self.annotations_path / self.annotations[idx]
with open(annotation_path, "r") as f:
annotation = json.load(f)
img = Image.open(impath) # .resize((1280, 720))
trf = T.Compose(
[
T.ToTensor(),
T.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
]
)
# prepare batches
img = trf(img)
# see https://git.tib.eu/vid2pos/sccvsd/-/blob/master/utils/synthetic_util.py
# draw lines (linewidth=4 for 720p) -> hence we rescale first
target = create_target_from_annotation(1280, 720, annotation, self.classes)
target = target.long().argmax(dim=0).unsqueeze(0)
# to binary mask
target = target.bool().float()
# rescale target equivalent to cv2.resize() with default args (interpolation bilinear) -> same as in torchvision
# bilinear -> [0, 0.25, 0.5, 0.7, 1.0] are entries
target = torchvision.transforms.Resize((180, 320))(target)
# print(torch.unique(target))
# to uint8
target = (target * 255.0).to(torch.uint8)
return img, target, impath.name
if __name__ == "__main__":
args = ArgumentParser()
args.add_argument("--data_dir", type=Path)
args.add_argument("--annotations", type=Path)
args.add_argument("--output_dir", type=Path)
args.add_argument("--extremities_prefix", type=str, default="")
args = args.parse_args()
data_dir = args.data_dir.parent
split = args.data_dir.name
output_dir = args.output_dir
if not output_dir.exists():
raise FileNotFoundError
dataset = ExtremitiesDataset(data_dir, split, args.annotations, filter_cam="Main camera center", extremities_prefix=args.extremities_prefix)
# img, edge_map, img_id = dataset[0]
# edge_map = edge_map.squeeze(0).numpy()
# Image.fromarray(edge_map).show()
image_src = []
edge_maps = np.zeros((len(dataset), 1, 180, 320), dtype=np.uint8)
for i, (_, edge_map, img_id) in enumerate(tqdm(dataset)):
edge_map = edge_map.numpy()
# Image.fromarray(edge_map).show()
edge_maps[i] = edge_map
image_src.append(img_id)
with h5py.File(output_dir / "seg_edge_maps.h5", "w") as f:
f.create_dataset("edge_map", data=edge_maps)
with open(output_dir / "seg_image_paths.pkl", "wb") as f:
pickle.dump(image_src, f)
|