Skip to content

Commit 691cda6

Browse files
arnabGuduarnabGudu
committed
added
1 parent 57db93e commit 691cda6

21 files changed

+2102
-2101
lines changed

‎.gitignore

+4-3
Original file line numberDiff line numberDiff line change
@@ -1,3 +1,4 @@
1-
*.pyc
2-
__pycache__
3-
.idea/
1+
*.pyc
2+
__pycache__
3+
.idea/
4+
checkpoints

‎LICENSE

+201-201
Large diffs are not rendered by default.

‎README.md

+123-123
Large diffs are not rendered by default.

‎TRAIN-ON-CUSTOM-DATASET.md

+75-75
Large diffs are not rendered by default.

‎datasets/coco.py

+179-179
Original file line numberDiff line numberDiff line change
@@ -1,179 +1,179 @@
1-
import copy
2-
import json
3-
import math
4-
import os
5-
import pickle
6-
7-
import cv2
8-
import numpy as np
9-
import pycocotools
10-
11-
from torch.utils.data.dataset import Dataset
12-
13-
BODY_PARTS_KPT_IDS = [[1, 8], [8, 9], [9, 10], [1, 11], [11, 12], [12, 13], [1, 2], [2, 3], [3, 4], [2, 16],
14-
[1, 5], [5, 6], [6, 7], [5, 17], [1, 0], [0, 14], [0, 15], [14, 16], [15, 17]]
15-
16-
17-
def get_mask(segmentations, mask):
18-
for segmentation in segmentations:
19-
rle = pycocotools.mask.frPyObjects(segmentation, mask.shape[0], mask.shape[1])
20-
mask[pycocotools.mask.decode(rle) > 0.5] = 0
21-
return mask
22-
23-
24-
class CocoTrainDataset(Dataset):
25-
def __init__(self, labels, images_folder, stride, sigma, paf_thickness, transform=None):
26-
super().__init__()
27-
self._images_folder = images_folder
28-
self._stride = stride
29-
self._sigma = sigma
30-
self._paf_thickness = paf_thickness
31-
self._transform = transform
32-
with open(labels, 'rb') as f:
33-
self._labels = pickle.load(f)
34-
35-
def __getitem__(self, idx):
36-
label = copy.deepcopy(self._labels[idx]) # label modified in transform
37-
image = cv2.imread(os.path.join(self._images_folder, label['img_paths']), cv2.IMREAD_COLOR)
38-
mask = np.ones(shape=(label['img_height'], label['img_width']), dtype=np.float32)
39-
mask = get_mask(label['segmentations'], mask)
40-
sample = {
41-
'label': label,
42-
'image': image,
43-
'mask': mask
44-
}
45-
if self._transform:
46-
sample = self._transform(sample)
47-
48-
mask = cv2.resize(sample['mask'], dsize=None, fx=1/self._stride, fy=1/self._stride, interpolation=cv2.INTER_AREA)
49-
keypoint_maps = self._generate_keypoint_maps(sample)
50-
sample['keypoint_maps'] = keypoint_maps
51-
keypoint_mask = np.zeros(shape=keypoint_maps.shape, dtype=np.float32)
52-
for idx in range(keypoint_mask.shape[0]):
53-
keypoint_mask[idx] = mask
54-
sample['keypoint_mask'] = keypoint_mask
55-
56-
paf_maps = self._generate_paf_maps(sample)
57-
sample['paf_maps'] = paf_maps
58-
paf_mask = np.zeros(shape=paf_maps.shape, dtype=np.float32)
59-
for idx in range(paf_mask.shape[0]):
60-
paf_mask[idx] = mask
61-
sample['paf_mask'] = paf_mask
62-
63-
image = sample['image'].astype(np.float32)
64-
image = (image - 128) / 256
65-
sample['image'] = image.transpose((2, 0, 1))
66-
del sample['label']
67-
return sample
68-
69-
def __len__(self):
70-
return len(self._labels)
71-
72-
def _generate_keypoint_maps(self, sample):
73-
n_keypoints = 18
74-
n_rows, n_cols, _ = sample['image'].shape
75-
keypoint_maps = np.zeros(shape=(n_keypoints + 1,
76-
n_rows // self._stride, n_cols // self._stride), dtype=np.float32) # +1 for bg
77-
78-
label = sample['label']
79-
for keypoint_idx in range(n_keypoints):
80-
keypoint = label['keypoints'][keypoint_idx]
81-
if keypoint[2] <= 1:
82-
self._add_gaussian(keypoint_maps[keypoint_idx], keypoint[0], keypoint[1], self._stride, self._sigma)
83-
for another_annotation in label['processed_other_annotations']:
84-
keypoint = another_annotation['keypoints'][keypoint_idx]
85-
if keypoint[2] <= 1:
86-
self._add_gaussian(keypoint_maps[keypoint_idx], keypoint[0], keypoint[1], self._stride, self._sigma)
87-
keypoint_maps[-1] = 1 - keypoint_maps.max(axis=0)
88-
return keypoint_maps
89-
90-
def _add_gaussian(self, keypoint_map, x, y, stride, sigma):
91-
n_sigma = 4
92-
tl = [int(x - n_sigma * sigma), int(y - n_sigma * sigma)]
93-
tl[0] = max(tl[0], 0)
94-
tl[1] = max(tl[1], 0)
95-
96-
br = [int(x + n_sigma * sigma), int(y + n_sigma * sigma)]
97-
map_h, map_w = keypoint_map.shape
98-
br[0] = min(br[0], map_w * stride)
99-
br[1] = min(br[1], map_h * stride)
100-
101-
shift = stride / 2 - 0.5
102-
for map_y in range(tl[1] // stride, br[1] // stride):
103-
for map_x in range(tl[0] // stride, br[0] // stride):
104-
d2 = (map_x * stride + shift - x) * (map_x * stride + shift - x) + \
105-
(map_y * stride + shift - y) * (map_y * stride + shift - y)
106-
exponent = d2 / 2 / sigma / sigma
107-
if exponent > 4.6052: # threshold, ln(100), ~0.01
108-
continue
109-
keypoint_map[map_y, map_x] += math.exp(-exponent)
110-
if keypoint_map[map_y, map_x] > 1:
111-
keypoint_map[map_y, map_x] = 1
112-
113-
def _generate_paf_maps(self, sample):
114-
n_pafs = len(BODY_PARTS_KPT_IDS)
115-
n_rows, n_cols, _ = sample['image'].shape
116-
paf_maps = np.zeros(shape=(n_pafs * 2, n_rows // self._stride, n_cols // self._stride), dtype=np.float32)
117-
118-
label = sample['label']
119-
for paf_idx in range(n_pafs):
120-
keypoint_a = label['keypoints'][BODY_PARTS_KPT_IDS[paf_idx][0]]
121-
keypoint_b = label['keypoints'][BODY_PARTS_KPT_IDS[paf_idx][1]]
122-
if keypoint_a[2] <= 1 and keypoint_b[2] <= 1:
123-
self._set_paf(paf_maps[paf_idx * 2:paf_idx * 2 + 2],
124-
keypoint_a[0], keypoint_a[1], keypoint_b[0], keypoint_b[1],
125-
self._stride, self._paf_thickness)
126-
for another_annotation in label['processed_other_annotations']:
127-
keypoint_a = another_annotation['keypoints'][BODY_PARTS_KPT_IDS[paf_idx][0]]
128-
keypoint_b = another_annotation['keypoints'][BODY_PARTS_KPT_IDS[paf_idx][1]]
129-
if keypoint_a[2] <= 1 and keypoint_b[2] <= 1:
130-
self._set_paf(paf_maps[paf_idx * 2:paf_idx * 2 + 2],
131-
keypoint_a[0], keypoint_a[1], keypoint_b[0], keypoint_b[1],
132-
self._stride, self._paf_thickness)
133-
return paf_maps
134-
135-
def _set_paf(self, paf_map, x_a, y_a, x_b, y_b, stride, thickness):
136-
x_a /= stride
137-
y_a /= stride
138-
x_b /= stride
139-
y_b /= stride
140-
x_ba = x_b - x_a
141-
y_ba = y_b - y_a
142-
_, h_map, w_map = paf_map.shape
143-
x_min = int(max(min(x_a, x_b) - thickness, 0))
144-
x_max = int(min(max(x_a, x_b) + thickness, w_map))
145-
y_min = int(max(min(y_a, y_b) - thickness, 0))
146-
y_max = int(min(max(y_a, y_b) + thickness, h_map))
147-
norm_ba = (x_ba * x_ba + y_ba * y_ba) ** 0.5
148-
if norm_ba < 1e-7: # Same points, no paf
149-
return
150-
x_ba /= norm_ba
151-
y_ba /= norm_ba
152-
153-
for y in range(y_min, y_max):
154-
for x in range(x_min, x_max):
155-
x_ca = x - x_a
156-
y_ca = y - y_a
157-
d = math.fabs(x_ca * y_ba - y_ca * x_ba)
158-
if d <= thickness:
159-
paf_map[0, y, x] = x_ba
160-
paf_map[1, y, x] = y_ba
161-
162-
163-
class CocoValDataset(Dataset):
164-
def __init__(self, labels, images_folder):
165-
super().__init__()
166-
with open(labels, 'r') as f:
167-
self._labels = json.load(f)
168-
self._images_folder = images_folder
169-
170-
def __getitem__(self, idx):
171-
file_name = self._labels['images'][idx]['file_name']
172-
img = cv2.imread(os.path.join(self._images_folder, file_name), cv2.IMREAD_COLOR)
173-
return {
174-
'img': img,
175-
'file_name': file_name
176-
}
177-
178-
def __len__(self):
179-
return len(self._labels['images'])
1+
import copy
2+
import json
3+
import math
4+
import os
5+
import pickle
6+
7+
import cv2
8+
import numpy as np
9+
import pycocotools
10+
11+
from torch.utils.data.dataset import Dataset
12+
13+
BODY_PARTS_KPT_IDS = [[1, 8], [8, 9], [9, 10], [1, 11], [11, 12], [12, 13], [1, 2], [2, 3], [3, 4], [2, 16],
14+
[1, 5], [5, 6], [6, 7], [5, 17], [1, 0], [0, 14], [0, 15], [14, 16], [15, 17]]
15+
16+
17+
def get_mask(segmentations, mask):
18+
for segmentation in segmentations:
19+
rle = pycocotools.mask.frPyObjects(segmentation, mask.shape[0], mask.shape[1])
20+
mask[pycocotools.mask.decode(rle) > 0.5] = 0
21+
return mask
22+
23+
24+
class CocoTrainDataset(Dataset):
25+
def __init__(self, labels, images_folder, stride, sigma, paf_thickness, transform=None):
26+
super().__init__()
27+
self._images_folder = images_folder
28+
self._stride = stride
29+
self._sigma = sigma
30+
self._paf_thickness = paf_thickness
31+
self._transform = transform
32+
with open(labels, 'rb') as f:
33+
self._labels = pickle.load(f)
34+
35+
def __getitem__(self, idx):
36+
label = copy.deepcopy(self._labels[idx]) # label modified in transform
37+
image = cv2.imread(os.path.join(self._images_folder, label['img_paths']), cv2.IMREAD_COLOR)
38+
mask = np.ones(shape=(label['img_height'], label['img_width']), dtype=np.float32)
39+
mask = get_mask(label['segmentations'], mask)
40+
sample = {
41+
'label': label,
42+
'image': image,
43+
'mask': mask
44+
}
45+
if self._transform:
46+
sample = self._transform(sample)
47+
48+
mask = cv2.resize(sample['mask'], dsize=None, fx=1/self._stride, fy=1/self._stride, interpolation=cv2.INTER_AREA)
49+
keypoint_maps = self._generate_keypoint_maps(sample)
50+
sample['keypoint_maps'] = keypoint_maps
51+
keypoint_mask = np.zeros(shape=keypoint_maps.shape, dtype=np.float32)
52+
for idx in range(keypoint_mask.shape[0]):
53+
keypoint_mask[idx] = mask
54+
sample['keypoint_mask'] = keypoint_mask
55+
56+
paf_maps = self._generate_paf_maps(sample)
57+
sample['paf_maps'] = paf_maps
58+
paf_mask = np.zeros(shape=paf_maps.shape, dtype=np.float32)
59+
for idx in range(paf_mask.shape[0]):
60+
paf_mask[idx] = mask
61+
sample['paf_mask'] = paf_mask
62+
63+
image = sample['image'].astype(np.float32)
64+
image = (image - 128) / 256
65+
sample['image'] = image.transpose((2, 0, 1))
66+
del sample['label']
67+
return sample
68+
69+
def __len__(self):
70+
return len(self._labels)
71+
72+
def _generate_keypoint_maps(self, sample):
73+
n_keypoints = 18
74+
n_rows, n_cols, _ = sample['image'].shape
75+
keypoint_maps = np.zeros(shape=(n_keypoints + 1,
76+
n_rows // self._stride, n_cols // self._stride), dtype=np.float32) # +1 for bg
77+
78+
label = sample['label']
79+
for keypoint_idx in range(n_keypoints):
80+
keypoint = label['keypoints'][keypoint_idx]
81+
if keypoint[2] <= 1:
82+
self._add_gaussian(keypoint_maps[keypoint_idx], keypoint[0], keypoint[1], self._stride, self._sigma)
83+
for another_annotation in label['processed_other_annotations']:
84+
keypoint = another_annotation['keypoints'][keypoint_idx]
85+
if keypoint[2] <= 1:
86+
self._add_gaussian(keypoint_maps[keypoint_idx], keypoint[0], keypoint[1], self._stride, self._sigma)
87+
keypoint_maps[-1] = 1 - keypoint_maps.max(axis=0)
88+
return keypoint_maps
89+
90+
def _add_gaussian(self, keypoint_map, x, y, stride, sigma):
91+
n_sigma = 4
92+
tl = [int(x - n_sigma * sigma), int(y - n_sigma * sigma)]
93+
tl[0] = max(tl[0], 0)
94+
tl[1] = max(tl[1], 0)
95+
96+
br = [int(x + n_sigma * sigma), int(y + n_sigma * sigma)]
97+
map_h, map_w = keypoint_map.shape
98+
br[0] = min(br[0], map_w * stride)
99+
br[1] = min(br[1], map_h * stride)
100+
101+
shift = stride / 2 - 0.5
102+
for map_y in range(tl[1] // stride, br[1] // stride):
103+
for map_x in range(tl[0] // stride, br[0] // stride):
104+
d2 = (map_x * stride + shift - x) * (map_x * stride + shift - x) + \
105+
(map_y * stride + shift - y) * (map_y * stride + shift - y)
106+
exponent = d2 / 2 / sigma / sigma
107+
if exponent > 4.6052: # threshold, ln(100), ~0.01
108+
continue
109+
keypoint_map[map_y, map_x] += math.exp(-exponent)
110+
if keypoint_map[map_y, map_x] > 1:
111+
keypoint_map[map_y, map_x] = 1
112+
113+
def _generate_paf_maps(self, sample):
114+
n_pafs = len(BODY_PARTS_KPT_IDS)
115+
n_rows, n_cols, _ = sample['image'].shape
116+
paf_maps = np.zeros(shape=(n_pafs * 2, n_rows // self._stride, n_cols // self._stride), dtype=np.float32)
117+
118+
label = sample['label']
119+
for paf_idx in range(n_pafs):
120+
keypoint_a = label['keypoints'][BODY_PARTS_KPT_IDS[paf_idx][0]]
121+
keypoint_b = label['keypoints'][BODY_PARTS_KPT_IDS[paf_idx][1]]
122+
if keypoint_a[2] <= 1 and keypoint_b[2] <= 1:
123+
self._set_paf(paf_maps[paf_idx * 2:paf_idx * 2 + 2],
124+
keypoint_a[0], keypoint_a[1], keypoint_b[0], keypoint_b[1],
125+
self._stride, self._paf_thickness)
126+
for another_annotation in label['processed_other_annotations']:
127+
keypoint_a = another_annotation['keypoints'][BODY_PARTS_KPT_IDS[paf_idx][0]]
128+
keypoint_b = another_annotation['keypoints'][BODY_PARTS_KPT_IDS[paf_idx][1]]
129+
if keypoint_a[2] <= 1 and keypoint_b[2] <= 1:
130+
self._set_paf(paf_maps[paf_idx * 2:paf_idx * 2 + 2],
131+
keypoint_a[0], keypoint_a[1], keypoint_b[0], keypoint_b[1],
132+
self._stride, self._paf_thickness)
133+
return paf_maps
134+
135+
def _set_paf(self, paf_map, x_a, y_a, x_b, y_b, stride, thickness):
136+
x_a /= stride
137+
y_a /= stride
138+
x_b /= stride
139+
y_b /= stride
140+
x_ba = x_b - x_a
141+
y_ba = y_b - y_a
142+
_, h_map, w_map = paf_map.shape
143+
x_min = int(max(min(x_a, x_b) - thickness, 0))
144+
x_max = int(min(max(x_a, x_b) + thickness, w_map))
145+
y_min = int(max(min(y_a, y_b) - thickness, 0))
146+
y_max = int(min(max(y_a, y_b) + thickness, h_map))
147+
norm_ba = (x_ba * x_ba + y_ba * y_ba) ** 0.5
148+
if norm_ba < 1e-7: # Same points, no paf
149+
return
150+
x_ba /= norm_ba
151+
y_ba /= norm_ba
152+
153+
for y in range(y_min, y_max):
154+
for x in range(x_min, x_max):
155+
x_ca = x - x_a
156+
y_ca = y - y_a
157+
d = math.fabs(x_ca * y_ba - y_ca * x_ba)
158+
if d <= thickness:
159+
paf_map[0, y, x] = x_ba
160+
paf_map[1, y, x] = y_ba
161+
162+
163+
class CocoValDataset(Dataset):
164+
def __init__(self, labels, images_folder):
165+
super().__init__()
166+
with open(labels, 'r') as f:
167+
self._labels = json.load(f)
168+
self._images_folder = images_folder
169+
170+
def __getitem__(self, idx):
171+
file_name = self._labels['images'][idx]['file_name']
172+
img = cv2.imread(os.path.join(self._images_folder, file_name), cv2.IMREAD_COLOR)
173+
return {
174+
'img': img,
175+
'file_name': file_name
176+
}
177+
178+
def __len__(self):
179+
return len(self._labels['images'])

0 commit comments

Comments
 (0)