Spaces:
Running
on
L40S
Running
on
L40S
franciszzj
commited on
Commit
β’
b213d84
1
Parent(s):
afadbd4
init code
Browse filesThis view is limited to 50 files because it contains too many changes. Β
See raw diff
- 3rdparty/SCHP/__init__.py +241 -0
- 3rdparty/SCHP/networks/AugmentCE2P.py +480 -0
- 3rdparty/SCHP/networks/__init__.py +13 -0
- 3rdparty/SCHP/utils/transforms.py +174 -0
- 3rdparty/densepose/__init__.py +20 -0
- 3rdparty/densepose/config.py +277 -0
- 3rdparty/densepose/converters/__init__.py +15 -0
- 3rdparty/densepose/converters/base.py +93 -0
- 3rdparty/densepose/converters/builtin.py +31 -0
- 3rdparty/densepose/converters/chart_output_hflip.py +71 -0
- 3rdparty/densepose/converters/chart_output_to_chart_result.py +188 -0
- 3rdparty/densepose/converters/hflip.py +34 -0
- 3rdparty/densepose/converters/segm_to_mask.py +150 -0
- 3rdparty/densepose/converters/to_chart_result.py +70 -0
- 3rdparty/densepose/converters/to_mask.py +49 -0
- 3rdparty/densepose/data/__init__.py +25 -0
- 3rdparty/densepose/data/build.py +736 -0
- 3rdparty/densepose/data/combined_loader.py +44 -0
- 3rdparty/densepose/data/dataset_mapper.py +168 -0
- 3rdparty/densepose/data/datasets/__init__.py +5 -0
- 3rdparty/densepose/data/datasets/builtin.py +16 -0
- 3rdparty/densepose/data/datasets/chimpnsee.py +29 -0
- 3rdparty/densepose/data/datasets/coco.py +432 -0
- 3rdparty/densepose/data/datasets/dataset_type.py +11 -0
- 3rdparty/densepose/data/datasets/lvis.py +257 -0
- 3rdparty/densepose/data/image_list_dataset.py +72 -0
- 3rdparty/densepose/data/inference_based_loader.py +172 -0
- 3rdparty/densepose/data/meshes/__init__.py +5 -0
- 3rdparty/densepose/data/meshes/builtin.py +101 -0
- 3rdparty/densepose/data/meshes/catalog.py +71 -0
- 3rdparty/densepose/data/samplers/__init__.py +8 -0
- 3rdparty/densepose/data/samplers/densepose_base.py +203 -0
- 3rdparty/densepose/data/samplers/densepose_confidence_based.py +108 -0
- 3rdparty/densepose/data/samplers/densepose_cse_base.py +139 -0
- 3rdparty/densepose/data/samplers/densepose_cse_confidence_based.py +119 -0
- 3rdparty/densepose/data/samplers/densepose_cse_uniform.py +12 -0
- 3rdparty/densepose/data/samplers/densepose_uniform.py +41 -0
- 3rdparty/densepose/data/samplers/mask_from_densepose.py +28 -0
- 3rdparty/densepose/data/samplers/prediction_to_gt.py +98 -0
- 3rdparty/densepose/data/transform/__init__.py +3 -0
- 3rdparty/densepose/data/transform/image.py +39 -0
- 3rdparty/densepose/data/utils.py +38 -0
- 3rdparty/densepose/data/video/__init__.py +17 -0
- 3rdparty/densepose/data/video/frame_selector.py +87 -0
- 3rdparty/densepose/data/video/video_keyframe_dataset.py +300 -0
- 3rdparty/densepose/engine/__init__.py +3 -0
- 3rdparty/densepose/engine/trainer.py +258 -0
- 3rdparty/densepose/evaluation/__init__.py +3 -0
- 3rdparty/densepose/evaluation/d2_evaluator_adapter.py +50 -0
- 3rdparty/densepose/evaluation/densepose_coco_evaluation.py +1303 -0
3rdparty/SCHP/__init__.py
ADDED
@@ -0,0 +1,241 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from collections import OrderedDict
|
2 |
+
|
3 |
+
import cv2
|
4 |
+
import numpy as np
|
5 |
+
import torch
|
6 |
+
from PIL import Image
|
7 |
+
from SCHP import networks
|
8 |
+
from SCHP.utils.transforms import get_affine_transform, transform_logits
|
9 |
+
from torchvision import transforms
|
10 |
+
|
11 |
+
|
12 |
+
def get_palette(num_cls):
|
13 |
+
"""Returns the color map for visualizing the segmentation mask.
|
14 |
+
Args:
|
15 |
+
num_cls: Number of classes
|
16 |
+
Returns:
|
17 |
+
The color map
|
18 |
+
"""
|
19 |
+
n = num_cls
|
20 |
+
palette = [0] * (n * 3)
|
21 |
+
for j in range(0, n):
|
22 |
+
lab = j
|
23 |
+
palette[j * 3 + 0] = 0
|
24 |
+
palette[j * 3 + 1] = 0
|
25 |
+
palette[j * 3 + 2] = 0
|
26 |
+
i = 0
|
27 |
+
while lab:
|
28 |
+
palette[j * 3 + 0] |= ((lab >> 0) & 1) << (7 - i)
|
29 |
+
palette[j * 3 + 1] |= ((lab >> 1) & 1) << (7 - i)
|
30 |
+
palette[j * 3 + 2] |= ((lab >> 2) & 1) << (7 - i)
|
31 |
+
i += 1
|
32 |
+
lab >>= 3
|
33 |
+
return palette
|
34 |
+
|
35 |
+
|
36 |
+
dataset_settings = {
|
37 |
+
"lip": {
|
38 |
+
"input_size": [473, 473],
|
39 |
+
"num_classes": 20,
|
40 |
+
"label": [
|
41 |
+
"Background",
|
42 |
+
"Hat",
|
43 |
+
"Hair",
|
44 |
+
"Glove",
|
45 |
+
"Sunglasses",
|
46 |
+
"Upper-clothes",
|
47 |
+
"Dress",
|
48 |
+
"Coat",
|
49 |
+
"Socks",
|
50 |
+
"Pants",
|
51 |
+
"Jumpsuits",
|
52 |
+
"Scarf",
|
53 |
+
"Skirt",
|
54 |
+
"Face",
|
55 |
+
"Left-arm",
|
56 |
+
"Right-arm",
|
57 |
+
"Left-leg",
|
58 |
+
"Right-leg",
|
59 |
+
"Left-shoe",
|
60 |
+
"Right-shoe",
|
61 |
+
],
|
62 |
+
},
|
63 |
+
"atr": {
|
64 |
+
"input_size": [512, 512],
|
65 |
+
"num_classes": 18,
|
66 |
+
"label": [
|
67 |
+
"Background",
|
68 |
+
"Hat",
|
69 |
+
"Hair",
|
70 |
+
"Sunglasses",
|
71 |
+
"Upper-clothes",
|
72 |
+
"Skirt",
|
73 |
+
"Pants",
|
74 |
+
"Dress",
|
75 |
+
"Belt",
|
76 |
+
"Left-shoe",
|
77 |
+
"Right-shoe",
|
78 |
+
"Face",
|
79 |
+
"Left-leg",
|
80 |
+
"Right-leg",
|
81 |
+
"Left-arm",
|
82 |
+
"Right-arm",
|
83 |
+
"Bag",
|
84 |
+
"Scarf",
|
85 |
+
],
|
86 |
+
},
|
87 |
+
"pascal": {
|
88 |
+
"input_size": [512, 512],
|
89 |
+
"num_classes": 7,
|
90 |
+
"label": [
|
91 |
+
"Background",
|
92 |
+
"Head",
|
93 |
+
"Torso",
|
94 |
+
"Upper Arms",
|
95 |
+
"Lower Arms",
|
96 |
+
"Upper Legs",
|
97 |
+
"Lower Legs",
|
98 |
+
],
|
99 |
+
},
|
100 |
+
}
|
101 |
+
|
102 |
+
|
103 |
+
class SCHP:
|
104 |
+
def __init__(self, ckpt_path, device):
|
105 |
+
dataset_type = None
|
106 |
+
if "lip" in ckpt_path:
|
107 |
+
dataset_type = "lip"
|
108 |
+
elif "atr" in ckpt_path:
|
109 |
+
dataset_type = "atr"
|
110 |
+
elif "pascal" in ckpt_path:
|
111 |
+
dataset_type = "pascal"
|
112 |
+
assert dataset_type is not None, "Dataset type not found in checkpoint path"
|
113 |
+
self.device = device
|
114 |
+
self.num_classes = dataset_settings[dataset_type]["num_classes"]
|
115 |
+
self.input_size = dataset_settings[dataset_type]["input_size"]
|
116 |
+
self.aspect_ratio = self.input_size[1] * 1.0 / self.input_size[0]
|
117 |
+
self.palette = get_palette(self.num_classes)
|
118 |
+
|
119 |
+
self.label = dataset_settings[dataset_type]["label"]
|
120 |
+
self.model = networks.init_model(
|
121 |
+
"resnet101", num_classes=self.num_classes, pretrained=None
|
122 |
+
).to(device)
|
123 |
+
self.load_ckpt(ckpt_path)
|
124 |
+
self.model.eval()
|
125 |
+
|
126 |
+
self.transform = transforms.Compose(
|
127 |
+
[
|
128 |
+
transforms.ToTensor(),
|
129 |
+
transforms.Normalize(
|
130 |
+
mean=[0.406, 0.456, 0.485], std=[0.225, 0.224, 0.229]
|
131 |
+
),
|
132 |
+
]
|
133 |
+
)
|
134 |
+
self.upsample = torch.nn.Upsample(
|
135 |
+
size=self.input_size, mode="bilinear", align_corners=True
|
136 |
+
)
|
137 |
+
|
138 |
+
def load_ckpt(self, ckpt_path):
|
139 |
+
rename_map = {
|
140 |
+
"decoder.conv3.2.weight": "decoder.conv3.3.weight",
|
141 |
+
"decoder.conv3.3.weight": "decoder.conv3.4.weight",
|
142 |
+
"decoder.conv3.3.bias": "decoder.conv3.4.bias",
|
143 |
+
"decoder.conv3.3.running_mean": "decoder.conv3.4.running_mean",
|
144 |
+
"decoder.conv3.3.running_var": "decoder.conv3.4.running_var",
|
145 |
+
"fushion.3.weight": "fushion.4.weight",
|
146 |
+
"fushion.3.bias": "fushion.4.bias",
|
147 |
+
}
|
148 |
+
state_dict = torch.load(ckpt_path, map_location="cpu")["state_dict"]
|
149 |
+
new_state_dict = OrderedDict()
|
150 |
+
for k, v in state_dict.items():
|
151 |
+
name = k[7:] # remove `module.`
|
152 |
+
new_state_dict[name] = v
|
153 |
+
new_state_dict_ = OrderedDict()
|
154 |
+
for k, v in list(new_state_dict.items()):
|
155 |
+
if k in rename_map:
|
156 |
+
new_state_dict_[rename_map[k]] = v
|
157 |
+
else:
|
158 |
+
new_state_dict_[k] = v
|
159 |
+
self.model.load_state_dict(new_state_dict_, strict=False)
|
160 |
+
|
161 |
+
def _box2cs(self, box):
|
162 |
+
x, y, w, h = box[:4]
|
163 |
+
return self._xywh2cs(x, y, w, h)
|
164 |
+
|
165 |
+
def _xywh2cs(self, x, y, w, h):
|
166 |
+
center = np.zeros((2), dtype=np.float32)
|
167 |
+
center[0] = x + w * 0.5
|
168 |
+
center[1] = y + h * 0.5
|
169 |
+
if w > self.aspect_ratio * h:
|
170 |
+
h = w * 1.0 / self.aspect_ratio
|
171 |
+
elif w < self.aspect_ratio * h:
|
172 |
+
w = h * self.aspect_ratio
|
173 |
+
scale = np.array([w, h], dtype=np.float32)
|
174 |
+
return center, scale
|
175 |
+
|
176 |
+
def preprocess(self, image):
|
177 |
+
if isinstance(image, str):
|
178 |
+
img = cv2.imread(image, cv2.IMREAD_COLOR)
|
179 |
+
elif isinstance(image, Image.Image):
|
180 |
+
# to cv2 format
|
181 |
+
img = np.array(image)
|
182 |
+
|
183 |
+
h, w, _ = img.shape
|
184 |
+
# Get person center and scale
|
185 |
+
person_center, s = self._box2cs([0, 0, w - 1, h - 1])
|
186 |
+
r = 0
|
187 |
+
trans = get_affine_transform(person_center, s, r, self.input_size)
|
188 |
+
input = cv2.warpAffine(
|
189 |
+
img,
|
190 |
+
trans,
|
191 |
+
(int(self.input_size[1]), int(self.input_size[0])),
|
192 |
+
flags=cv2.INTER_LINEAR,
|
193 |
+
borderMode=cv2.BORDER_CONSTANT,
|
194 |
+
borderValue=(0, 0, 0),
|
195 |
+
)
|
196 |
+
|
197 |
+
input = self.transform(input).to(self.device).unsqueeze(0)
|
198 |
+
meta = {
|
199 |
+
"center": person_center,
|
200 |
+
"height": h,
|
201 |
+
"width": w,
|
202 |
+
"scale": s,
|
203 |
+
"rotation": r,
|
204 |
+
}
|
205 |
+
return input, meta
|
206 |
+
|
207 |
+
def __call__(self, image_or_path):
|
208 |
+
if isinstance(image_or_path, list):
|
209 |
+
image_list = []
|
210 |
+
meta_list = []
|
211 |
+
for image in image_or_path:
|
212 |
+
image, meta = self.preprocess(image)
|
213 |
+
image_list.append(image)
|
214 |
+
meta_list.append(meta)
|
215 |
+
image = torch.cat(image_list, dim=0)
|
216 |
+
else:
|
217 |
+
image, meta = self.preprocess(image_or_path)
|
218 |
+
meta_list = [meta]
|
219 |
+
|
220 |
+
output = self.model(image)
|
221 |
+
# upsample_outputs = self.upsample(output[0][-1])
|
222 |
+
upsample_outputs = self.upsample(output)
|
223 |
+
upsample_outputs = upsample_outputs.permute(0, 2, 3, 1) # BCHW -> BHWC
|
224 |
+
|
225 |
+
output_img_list = []
|
226 |
+
for upsample_output, meta in zip(upsample_outputs, meta_list):
|
227 |
+
c, s, w, h = meta["center"], meta["scale"], meta["width"], meta["height"]
|
228 |
+
logits_result = transform_logits(
|
229 |
+
upsample_output.data.cpu().numpy(),
|
230 |
+
c,
|
231 |
+
s,
|
232 |
+
w,
|
233 |
+
h,
|
234 |
+
input_size=self.input_size,
|
235 |
+
)
|
236 |
+
parsing_result = np.argmax(logits_result, axis=2)
|
237 |
+
output_img = Image.fromarray(np.asarray(parsing_result, dtype=np.uint8))
|
238 |
+
output_img.putpalette(self.palette)
|
239 |
+
output_img_list.append(output_img)
|
240 |
+
|
241 |
+
return output_img_list[0] if len(output_img_list) == 1 else output_img_list
|
3rdparty/SCHP/networks/AugmentCE2P.py
ADDED
@@ -0,0 +1,480 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
#!/usr/bin/env python
|
2 |
+
# -*- encoding: utf-8 -*-
|
3 |
+
|
4 |
+
"""
|
5 |
+
@Author : Peike Li
|
6 |
+
@Contact : [email protected]
|
7 |
+
@File : AugmentCE2P.py
|
8 |
+
@Time : 8/4/19 3:35 PM
|
9 |
+
@Desc :
|
10 |
+
@License : This source code is licensed under the license found in the
|
11 |
+
LICENSE file in the root directory of this source tree.
|
12 |
+
"""
|
13 |
+
|
14 |
+
import torch
|
15 |
+
import torch.nn as nn
|
16 |
+
|
17 |
+
from torch.nn import BatchNorm2d, functional as F, LeakyReLU
|
18 |
+
|
19 |
+
affine_par = True
|
20 |
+
pretrained_settings = {
|
21 |
+
"resnet101": {
|
22 |
+
"imagenet": {
|
23 |
+
"input_space": "BGR",
|
24 |
+
"input_size": [3, 224, 224],
|
25 |
+
"input_range": [0, 1],
|
26 |
+
"mean": [0.406, 0.456, 0.485],
|
27 |
+
"std": [0.225, 0.224, 0.229],
|
28 |
+
"num_classes": 1000,
|
29 |
+
}
|
30 |
+
},
|
31 |
+
}
|
32 |
+
|
33 |
+
|
34 |
+
def conv3x3(in_planes, out_planes, stride=1):
|
35 |
+
"3x3 convolution with padding"
|
36 |
+
return nn.Conv2d(
|
37 |
+
in_planes, out_planes, kernel_size=3, stride=stride, padding=1, bias=False
|
38 |
+
)
|
39 |
+
|
40 |
+
|
41 |
+
class Bottleneck(nn.Module):
|
42 |
+
expansion = 4
|
43 |
+
|
44 |
+
def __init__(
|
45 |
+
self,
|
46 |
+
inplanes,
|
47 |
+
planes,
|
48 |
+
stride=1,
|
49 |
+
dilation=1,
|
50 |
+
downsample=None,
|
51 |
+
fist_dilation=1,
|
52 |
+
multi_grid=1,
|
53 |
+
):
|
54 |
+
super(Bottleneck, self).__init__()
|
55 |
+
self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
|
56 |
+
self.bn1 = BatchNorm2d(planes)
|
57 |
+
self.conv2 = nn.Conv2d(
|
58 |
+
planes,
|
59 |
+
planes,
|
60 |
+
kernel_size=3,
|
61 |
+
stride=stride,
|
62 |
+
padding=dilation * multi_grid,
|
63 |
+
dilation=dilation * multi_grid,
|
64 |
+
bias=False,
|
65 |
+
)
|
66 |
+
self.bn2 = BatchNorm2d(planes)
|
67 |
+
self.conv3 = nn.Conv2d(planes, planes * 4, kernel_size=1, bias=False)
|
68 |
+
self.bn3 = BatchNorm2d(planes * 4)
|
69 |
+
self.relu = nn.ReLU(inplace=False)
|
70 |
+
self.relu_inplace = nn.ReLU(inplace=True)
|
71 |
+
self.downsample = downsample
|
72 |
+
self.dilation = dilation
|
73 |
+
self.stride = stride
|
74 |
+
|
75 |
+
def forward(self, x):
|
76 |
+
residual = x
|
77 |
+
|
78 |
+
out = self.conv1(x)
|
79 |
+
out = self.bn1(out)
|
80 |
+
out = self.relu(out)
|
81 |
+
|
82 |
+
out = self.conv2(out)
|
83 |
+
out = self.bn2(out)
|
84 |
+
out = self.relu(out)
|
85 |
+
|
86 |
+
out = self.conv3(out)
|
87 |
+
out = self.bn3(out)
|
88 |
+
|
89 |
+
if self.downsample is not None:
|
90 |
+
residual = self.downsample(x)
|
91 |
+
|
92 |
+
out = out + residual
|
93 |
+
out = self.relu_inplace(out)
|
94 |
+
|
95 |
+
return out
|
96 |
+
|
97 |
+
|
98 |
+
class PSPModule(nn.Module):
|
99 |
+
"""
|
100 |
+
Reference:
|
101 |
+
Zhao, Hengshuang, et al. *"Pyramid scene parsing network."*
|
102 |
+
"""
|
103 |
+
|
104 |
+
def __init__(self, features, out_features=512, sizes=(1, 2, 3, 6)):
|
105 |
+
super(PSPModule, self).__init__()
|
106 |
+
|
107 |
+
self.stages = []
|
108 |
+
self.stages = nn.ModuleList(
|
109 |
+
[self._make_stage(features, out_features, size) for size in sizes]
|
110 |
+
)
|
111 |
+
self.bottleneck = nn.Sequential(
|
112 |
+
nn.Conv2d(
|
113 |
+
features + len(sizes) * out_features,
|
114 |
+
out_features,
|
115 |
+
kernel_size=3,
|
116 |
+
padding=1,
|
117 |
+
dilation=1,
|
118 |
+
bias=False,
|
119 |
+
),
|
120 |
+
BatchNorm2d(out_features),
|
121 |
+
LeakyReLU(),
|
122 |
+
)
|
123 |
+
|
124 |
+
def _make_stage(self, features, out_features, size):
|
125 |
+
prior = nn.AdaptiveAvgPool2d(output_size=(size, size))
|
126 |
+
conv = nn.Conv2d(features, out_features, kernel_size=1, bias=False)
|
127 |
+
return nn.Sequential(
|
128 |
+
prior,
|
129 |
+
conv,
|
130 |
+
# bn
|
131 |
+
BatchNorm2d(out_features),
|
132 |
+
LeakyReLU(),
|
133 |
+
)
|
134 |
+
|
135 |
+
def forward(self, feats):
|
136 |
+
h, w = feats.size(2), feats.size(3)
|
137 |
+
priors = [
|
138 |
+
F.interpolate(
|
139 |
+
input=stage(feats), size=(h, w), mode="bilinear", align_corners=True
|
140 |
+
)
|
141 |
+
for stage in self.stages
|
142 |
+
] + [feats]
|
143 |
+
bottle = self.bottleneck(torch.cat(priors, 1))
|
144 |
+
return bottle
|
145 |
+
|
146 |
+
|
147 |
+
class ASPPModule(nn.Module):
|
148 |
+
"""
|
149 |
+
Reference:
|
150 |
+
Chen, Liang-Chieh, et al. *"Rethinking Atrous Convolution for Semantic Image Segmentation."*
|
151 |
+
"""
|
152 |
+
|
153 |
+
def __init__(
|
154 |
+
self, features, inner_features=256, out_features=512, dilations=(12, 24, 36)
|
155 |
+
):
|
156 |
+
super(ASPPModule, self).__init__()
|
157 |
+
|
158 |
+
self.conv1 = nn.Sequential(
|
159 |
+
nn.AdaptiveAvgPool2d((1, 1)),
|
160 |
+
nn.Conv2d(
|
161 |
+
features,
|
162 |
+
inner_features,
|
163 |
+
kernel_size=1,
|
164 |
+
padding=0,
|
165 |
+
dilation=1,
|
166 |
+
bias=False,
|
167 |
+
),
|
168 |
+
# InPlaceABNSync(inner_features)
|
169 |
+
BatchNorm2d(inner_features),
|
170 |
+
LeakyReLU(),
|
171 |
+
)
|
172 |
+
self.conv2 = nn.Sequential(
|
173 |
+
nn.Conv2d(
|
174 |
+
features,
|
175 |
+
inner_features,
|
176 |
+
kernel_size=1,
|
177 |
+
padding=0,
|
178 |
+
dilation=1,
|
179 |
+
bias=False,
|
180 |
+
),
|
181 |
+
BatchNorm2d(inner_features),
|
182 |
+
LeakyReLU(),
|
183 |
+
)
|
184 |
+
self.conv3 = nn.Sequential(
|
185 |
+
nn.Conv2d(
|
186 |
+
features,
|
187 |
+
inner_features,
|
188 |
+
kernel_size=3,
|
189 |
+
padding=dilations[0],
|
190 |
+
dilation=dilations[0],
|
191 |
+
bias=False,
|
192 |
+
),
|
193 |
+
BatchNorm2d(inner_features),
|
194 |
+
LeakyReLU(),
|
195 |
+
)
|
196 |
+
self.conv4 = nn.Sequential(
|
197 |
+
nn.Conv2d(
|
198 |
+
features,
|
199 |
+
inner_features,
|
200 |
+
kernel_size=3,
|
201 |
+
padding=dilations[1],
|
202 |
+
dilation=dilations[1],
|
203 |
+
bias=False,
|
204 |
+
),
|
205 |
+
BatchNorm2d(inner_features),
|
206 |
+
LeakyReLU(),
|
207 |
+
)
|
208 |
+
self.conv5 = nn.Sequential(
|
209 |
+
nn.Conv2d(
|
210 |
+
features,
|
211 |
+
inner_features,
|
212 |
+
kernel_size=3,
|
213 |
+
padding=dilations[2],
|
214 |
+
dilation=dilations[2],
|
215 |
+
bias=False,
|
216 |
+
),
|
217 |
+
BatchNorm2d(inner_features),
|
218 |
+
LeakyReLU(),
|
219 |
+
)
|
220 |
+
|
221 |
+
self.bottleneck = nn.Sequential(
|
222 |
+
nn.Conv2d(
|
223 |
+
inner_features * 5,
|
224 |
+
out_features,
|
225 |
+
kernel_size=1,
|
226 |
+
padding=0,
|
227 |
+
dilation=1,
|
228 |
+
bias=False,
|
229 |
+
),
|
230 |
+
BatchNorm2d(inner_features),
|
231 |
+
LeakyReLU(),
|
232 |
+
nn.Dropout2d(0.1),
|
233 |
+
)
|
234 |
+
|
235 |
+
def forward(self, x):
|
236 |
+
_, _, h, w = x.size()
|
237 |
+
|
238 |
+
feat1 = F.interpolate(
|
239 |
+
self.conv1(x), size=(h, w), mode="bilinear", align_corners=True
|
240 |
+
)
|
241 |
+
|
242 |
+
feat2 = self.conv2(x)
|
243 |
+
feat3 = self.conv3(x)
|
244 |
+
feat4 = self.conv4(x)
|
245 |
+
feat5 = self.conv5(x)
|
246 |
+
out = torch.cat((feat1, feat2, feat3, feat4, feat5), 1)
|
247 |
+
|
248 |
+
bottle = self.bottleneck(out)
|
249 |
+
return bottle
|
250 |
+
|
251 |
+
|
252 |
+
class Edge_Module(nn.Module):
|
253 |
+
"""
|
254 |
+
Edge Learning Branch
|
255 |
+
"""
|
256 |
+
|
257 |
+
def __init__(self, in_fea=[256, 512, 1024], mid_fea=256, out_fea=2):
|
258 |
+
super(Edge_Module, self).__init__()
|
259 |
+
|
260 |
+
self.conv1 = nn.Sequential(
|
261 |
+
nn.Conv2d(
|
262 |
+
in_fea[0], mid_fea, kernel_size=1, padding=0, dilation=1, bias=False
|
263 |
+
),
|
264 |
+
BatchNorm2d(mid_fea),
|
265 |
+
LeakyReLU(),
|
266 |
+
)
|
267 |
+
self.conv2 = nn.Sequential(
|
268 |
+
nn.Conv2d(
|
269 |
+
in_fea[1], mid_fea, kernel_size=1, padding=0, dilation=1, bias=False
|
270 |
+
),
|
271 |
+
BatchNorm2d(mid_fea),
|
272 |
+
LeakyReLU(),
|
273 |
+
)
|
274 |
+
self.conv3 = nn.Sequential(
|
275 |
+
nn.Conv2d(
|
276 |
+
in_fea[2], mid_fea, kernel_size=1, padding=0, dilation=1, bias=False
|
277 |
+
),
|
278 |
+
BatchNorm2d(mid_fea),
|
279 |
+
LeakyReLU(),
|
280 |
+
)
|
281 |
+
self.conv4 = nn.Conv2d(
|
282 |
+
mid_fea, out_fea, kernel_size=3, padding=1, dilation=1, bias=True
|
283 |
+
)
|
284 |
+
# self.conv5 = nn.Conv2d(out_fea * 3, out_fea, kernel_size=1, padding=0, dilation=1, bias=True)
|
285 |
+
|
286 |
+
def forward(self, x1, x2, x3):
|
287 |
+
_, _, h, w = x1.size()
|
288 |
+
|
289 |
+
edge1_fea = self.conv1(x1)
|
290 |
+
# edge1 = self.conv4(edge1_fea)
|
291 |
+
edge2_fea = self.conv2(x2)
|
292 |
+
edge2 = self.conv4(edge2_fea)
|
293 |
+
edge3_fea = self.conv3(x3)
|
294 |
+
edge3 = self.conv4(edge3_fea)
|
295 |
+
|
296 |
+
edge2_fea = F.interpolate(
|
297 |
+
edge2_fea, size=(h, w), mode="bilinear", align_corners=True
|
298 |
+
)
|
299 |
+
edge3_fea = F.interpolate(
|
300 |
+
edge3_fea, size=(h, w), mode="bilinear", align_corners=True
|
301 |
+
)
|
302 |
+
edge2 = F.interpolate(edge2, size=(h, w), mode="bilinear", align_corners=True)
|
303 |
+
edge3 = F.interpolate(edge3, size=(h, w), mode="bilinear", align_corners=True)
|
304 |
+
|
305 |
+
# edge = torch.cat([edge1, edge2, edge3], dim=1)
|
306 |
+
edge_fea = torch.cat([edge1_fea, edge2_fea, edge3_fea], dim=1)
|
307 |
+
# edge = self.conv5(edge)
|
308 |
+
|
309 |
+
# return edge, edge_fea
|
310 |
+
return edge_fea
|
311 |
+
|
312 |
+
|
313 |
+
class Decoder_Module(nn.Module):
|
314 |
+
"""
|
315 |
+
Parsing Branch Decoder Module.
|
316 |
+
"""
|
317 |
+
|
318 |
+
def __init__(self, num_classes):
|
319 |
+
super(Decoder_Module, self).__init__()
|
320 |
+
self.conv1 = nn.Sequential(
|
321 |
+
nn.Conv2d(512, 256, kernel_size=1, padding=0, dilation=1, bias=False),
|
322 |
+
BatchNorm2d(256),
|
323 |
+
LeakyReLU(),
|
324 |
+
)
|
325 |
+
self.conv2 = nn.Sequential(
|
326 |
+
nn.Conv2d(
|
327 |
+
256, 48, kernel_size=1, stride=1, padding=0, dilation=1, bias=False
|
328 |
+
),
|
329 |
+
BatchNorm2d(48),
|
330 |
+
LeakyReLU(),
|
331 |
+
)
|
332 |
+
self.conv3 = nn.Sequential(
|
333 |
+
nn.Conv2d(304, 256, kernel_size=1, padding=0, dilation=1, bias=False),
|
334 |
+
BatchNorm2d(256),
|
335 |
+
LeakyReLU(),
|
336 |
+
nn.Conv2d(256, 256, kernel_size=1, padding=0, dilation=1, bias=False),
|
337 |
+
BatchNorm2d(256),
|
338 |
+
LeakyReLU(),
|
339 |
+
)
|
340 |
+
|
341 |
+
# self.conv4 = nn.Conv2d(256, num_classes, kernel_size=1, padding=0, dilation=1, bias=True)
|
342 |
+
|
343 |
+
def forward(self, xt, xl):
|
344 |
+
_, _, h, w = xl.size()
|
345 |
+
xt = F.interpolate(
|
346 |
+
self.conv1(xt), size=(h, w), mode="bilinear", align_corners=True
|
347 |
+
)
|
348 |
+
xl = self.conv2(xl)
|
349 |
+
x = torch.cat([xt, xl], dim=1)
|
350 |
+
x = self.conv3(x)
|
351 |
+
# seg = self.conv4(x)
|
352 |
+
# return seg, x
|
353 |
+
return x
|
354 |
+
|
355 |
+
|
356 |
+
class ResNet(nn.Module):
|
357 |
+
def __init__(self, block, layers, num_classes):
|
358 |
+
self.inplanes = 128
|
359 |
+
super(ResNet, self).__init__()
|
360 |
+
self.conv1 = conv3x3(3, 64, stride=2)
|
361 |
+
self.bn1 = BatchNorm2d(64)
|
362 |
+
self.relu1 = nn.ReLU(inplace=False)
|
363 |
+
self.conv2 = conv3x3(64, 64)
|
364 |
+
self.bn2 = BatchNorm2d(64)
|
365 |
+
self.relu2 = nn.ReLU(inplace=False)
|
366 |
+
self.conv3 = conv3x3(64, 128)
|
367 |
+
self.bn3 = BatchNorm2d(128)
|
368 |
+
self.relu3 = nn.ReLU(inplace=False)
|
369 |
+
|
370 |
+
self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
|
371 |
+
|
372 |
+
self.layer1 = self._make_layer(block, 64, layers[0])
|
373 |
+
self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
|
374 |
+
self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
|
375 |
+
self.layer4 = self._make_layer(
|
376 |
+
block, 512, layers[3], stride=1, dilation=2, multi_grid=(1, 1, 1)
|
377 |
+
)
|
378 |
+
|
379 |
+
self.context_encoding = PSPModule(2048, 512)
|
380 |
+
|
381 |
+
self.edge = Edge_Module()
|
382 |
+
self.decoder = Decoder_Module(num_classes)
|
383 |
+
|
384 |
+
self.fushion = nn.Sequential(
|
385 |
+
nn.Conv2d(1024, 256, kernel_size=1, padding=0, dilation=1, bias=False),
|
386 |
+
BatchNorm2d(256),
|
387 |
+
LeakyReLU(),
|
388 |
+
nn.Dropout2d(0.1),
|
389 |
+
nn.Conv2d(
|
390 |
+
256, num_classes, kernel_size=1, padding=0, dilation=1, bias=True
|
391 |
+
),
|
392 |
+
)
|
393 |
+
|
394 |
+
def _make_layer(self, block, planes, blocks, stride=1, dilation=1, multi_grid=1):
|
395 |
+
downsample = None
|
396 |
+
if stride != 1 or self.inplanes != planes * block.expansion:
|
397 |
+
downsample = nn.Sequential(
|
398 |
+
nn.Conv2d(
|
399 |
+
self.inplanes,
|
400 |
+
planes * block.expansion,
|
401 |
+
kernel_size=1,
|
402 |
+
stride=stride,
|
403 |
+
bias=False,
|
404 |
+
),
|
405 |
+
BatchNorm2d(planes * block.expansion, affine=affine_par),
|
406 |
+
)
|
407 |
+
|
408 |
+
layers = []
|
409 |
+
generate_multi_grid = lambda index, grids: (
|
410 |
+
grids[index % len(grids)] if isinstance(grids, tuple) else 1
|
411 |
+
)
|
412 |
+
layers.append(
|
413 |
+
block(
|
414 |
+
self.inplanes,
|
415 |
+
planes,
|
416 |
+
stride,
|
417 |
+
dilation=dilation,
|
418 |
+
downsample=downsample,
|
419 |
+
multi_grid=generate_multi_grid(0, multi_grid),
|
420 |
+
)
|
421 |
+
)
|
422 |
+
self.inplanes = planes * block.expansion
|
423 |
+
for i in range(1, blocks):
|
424 |
+
layers.append(
|
425 |
+
block(
|
426 |
+
self.inplanes,
|
427 |
+
planes,
|
428 |
+
dilation=dilation,
|
429 |
+
multi_grid=generate_multi_grid(i, multi_grid),
|
430 |
+
)
|
431 |
+
)
|
432 |
+
|
433 |
+
return nn.Sequential(*layers)
|
434 |
+
|
435 |
+
def forward(self, x):
|
436 |
+
x = self.relu1(self.bn1(self.conv1(x)))
|
437 |
+
x = self.relu2(self.bn2(self.conv2(x)))
|
438 |
+
x = self.relu3(self.bn3(self.conv3(x)))
|
439 |
+
x = self.maxpool(x)
|
440 |
+
x2 = self.layer1(x)
|
441 |
+
x3 = self.layer2(x2)
|
442 |
+
x4 = self.layer3(x3)
|
443 |
+
x5 = self.layer4(x4)
|
444 |
+
x = self.context_encoding(x5)
|
445 |
+
# parsing_result, parsing_fea = self.decoder(x, x2)
|
446 |
+
parsing_fea = self.decoder(x, x2)
|
447 |
+
# Edge Branch
|
448 |
+
# edge_result, edge_fea = self.edge(x2, x3, x4)
|
449 |
+
edge_fea = self.edge(x2, x3, x4)
|
450 |
+
# Fusion Branch
|
451 |
+
x = torch.cat([parsing_fea, edge_fea], dim=1)
|
452 |
+
fusion_result = self.fushion(x)
|
453 |
+
# return [[parsing_result, fusion_result], [edge_result]]
|
454 |
+
return fusion_result
|
455 |
+
|
456 |
+
|
457 |
+
def initialize_pretrained_model(
|
458 |
+
model, settings, pretrained="./models/resnet101-imagenet.pth"
|
459 |
+
):
|
460 |
+
model.input_space = settings["input_space"]
|
461 |
+
model.input_size = settings["input_size"]
|
462 |
+
model.input_range = settings["input_range"]
|
463 |
+
model.mean = settings["mean"]
|
464 |
+
model.std = settings["std"]
|
465 |
+
|
466 |
+
if pretrained is not None:
|
467 |
+
saved_state_dict = torch.load(pretrained)
|
468 |
+
new_params = model.state_dict().copy()
|
469 |
+
for i in saved_state_dict:
|
470 |
+
i_parts = i.split(".")
|
471 |
+
if not i_parts[0] == "fc":
|
472 |
+
new_params[".".join(i_parts[0:])] = saved_state_dict[i]
|
473 |
+
model.load_state_dict(new_params)
|
474 |
+
|
475 |
+
|
476 |
+
def resnet101(num_classes=20, pretrained="./models/resnet101-imagenet.pth"):
|
477 |
+
model = ResNet(Bottleneck, [3, 4, 23, 3], num_classes)
|
478 |
+
settings = pretrained_settings["resnet101"]["imagenet"]
|
479 |
+
initialize_pretrained_model(model, settings, pretrained)
|
480 |
+
return model
|
3rdparty/SCHP/networks/__init__.py
ADDED
@@ -0,0 +1,13 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
from __future__ import absolute_import
|
2 |
+
|
3 |
+
from SCHP.networks.AugmentCE2P import resnet101
|
4 |
+
|
5 |
+
__factory = {
|
6 |
+
"resnet101": resnet101,
|
7 |
+
}
|
8 |
+
|
9 |
+
|
10 |
+
def init_model(name, *args, **kwargs):
|
11 |
+
if name not in __factory.keys():
|
12 |
+
raise KeyError("Unknown model arch: {}".format(name))
|
13 |
+
return __factory[name](*args, **kwargs)
|
3rdparty/SCHP/utils/transforms.py
ADDED
@@ -0,0 +1,174 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# ------------------------------------------------------------------------------
|
2 |
+
# Copyright (c) Microsoft
|
3 |
+
# Licensed under the MIT License.
|
4 |
+
# Written by Bin Xiao ([email protected])
|
5 |
+
# ------------------------------------------------------------------------------
|
6 |
+
|
7 |
+
from __future__ import absolute_import, division, print_function
|
8 |
+
|
9 |
+
import cv2
|
10 |
+
|
11 |
+
import numpy as np
|
12 |
+
import torch
|
13 |
+
|
14 |
+
|
15 |
+
class BRG2Tensor_transform(object):
|
16 |
+
def __call__(self, pic):
|
17 |
+
img = torch.from_numpy(pic.transpose((2, 0, 1)))
|
18 |
+
if isinstance(img, torch.ByteTensor):
|
19 |
+
return img.float()
|
20 |
+
else:
|
21 |
+
return img
|
22 |
+
|
23 |
+
|
24 |
+
class BGR2RGB_transform(object):
|
25 |
+
def __call__(self, tensor):
|
26 |
+
return tensor[[2, 1, 0], :, :]
|
27 |
+
|
28 |
+
|
29 |
+
def flip_back(output_flipped, matched_parts):
|
30 |
+
"""
|
31 |
+
ouput_flipped: numpy.ndarray(batch_size, num_joints, height, width)
|
32 |
+
"""
|
33 |
+
assert (
|
34 |
+
output_flipped.ndim == 4
|
35 |
+
), "output_flipped should be [batch_size, num_joints, height, width]"
|
36 |
+
|
37 |
+
output_flipped = output_flipped[:, :, :, ::-1]
|
38 |
+
|
39 |
+
for pair in matched_parts:
|
40 |
+
tmp = output_flipped[:, pair[0], :, :].copy()
|
41 |
+
output_flipped[:, pair[0], :, :] = output_flipped[:, pair[1], :, :]
|
42 |
+
output_flipped[:, pair[1], :, :] = tmp
|
43 |
+
|
44 |
+
return output_flipped
|
45 |
+
|
46 |
+
|
47 |
+
def fliplr_joints(joints, joints_vis, width, matched_parts):
|
48 |
+
"""
|
49 |
+
flip coords
|
50 |
+
"""
|
51 |
+
# Flip horizontal
|
52 |
+
joints[:, 0] = width - joints[:, 0] - 1
|
53 |
+
|
54 |
+
# Change left-right parts
|
55 |
+
for pair in matched_parts:
|
56 |
+
joints[pair[0], :], joints[pair[1], :] = (
|
57 |
+
joints[pair[1], :],
|
58 |
+
joints[pair[0], :].copy(),
|
59 |
+
)
|
60 |
+
joints_vis[pair[0], :], joints_vis[pair[1], :] = (
|
61 |
+
joints_vis[pair[1], :],
|
62 |
+
joints_vis[pair[0], :].copy(),
|
63 |
+
)
|
64 |
+
|
65 |
+
return joints * joints_vis, joints_vis
|
66 |
+
|
67 |
+
|
68 |
+
def transform_preds(coords, center, scale, input_size):
|
69 |
+
target_coords = np.zeros(coords.shape)
|
70 |
+
trans = get_affine_transform(center, scale, 0, input_size, inv=1)
|
71 |
+
for p in range(coords.shape[0]):
|
72 |
+
target_coords[p, 0:2] = affine_transform(coords[p, 0:2], trans)
|
73 |
+
return target_coords
|
74 |
+
|
75 |
+
|
76 |
+
def transform_parsing(pred, center, scale, width, height, input_size):
|
77 |
+
|
78 |
+
trans = get_affine_transform(center, scale, 0, input_size, inv=1)
|
79 |
+
target_pred = cv2.warpAffine(
|
80 |
+
pred,
|
81 |
+
trans,
|
82 |
+
(int(width), int(height)), # (int(width), int(height)),
|
83 |
+
flags=cv2.INTER_NEAREST,
|
84 |
+
borderMode=cv2.BORDER_CONSTANT,
|
85 |
+
borderValue=(0),
|
86 |
+
)
|
87 |
+
|
88 |
+
return target_pred
|
89 |
+
|
90 |
+
|
91 |
+
def transform_logits(logits, center, scale, width, height, input_size):
|
92 |
+
|
93 |
+
trans = get_affine_transform(center, scale, 0, input_size, inv=1)
|
94 |
+
channel = logits.shape[2]
|
95 |
+
target_logits = []
|
96 |
+
for i in range(channel):
|
97 |
+
target_logit = cv2.warpAffine(
|
98 |
+
logits[:, :, i],
|
99 |
+
trans,
|
100 |
+
(int(width), int(height)), # (int(width), int(height)),
|
101 |
+
flags=cv2.INTER_LINEAR,
|
102 |
+
borderMode=cv2.BORDER_CONSTANT,
|
103 |
+
borderValue=(0),
|
104 |
+
)
|
105 |
+
target_logits.append(target_logit)
|
106 |
+
target_logits = np.stack(target_logits, axis=2)
|
107 |
+
|
108 |
+
return target_logits
|
109 |
+
|
110 |
+
|
111 |
+
def get_affine_transform(
|
112 |
+
center, scale, rot, output_size, shift=np.array([0, 0], dtype=np.float32), inv=0
|
113 |
+
):
|
114 |
+
if not isinstance(scale, np.ndarray) and not isinstance(scale, list):
|
115 |
+
print(scale)
|
116 |
+
scale = np.array([scale, scale])
|
117 |
+
|
118 |
+
scale_tmp = scale
|
119 |
+
|
120 |
+
src_w = scale_tmp[0]
|
121 |
+
dst_w = output_size[1]
|
122 |
+
dst_h = output_size[0]
|
123 |
+
|
124 |
+
rot_rad = np.pi * rot / 180
|
125 |
+
src_dir = get_dir([0, src_w * -0.5], rot_rad)
|
126 |
+
dst_dir = np.array([0, (dst_w - 1) * -0.5], np.float32)
|
127 |
+
|
128 |
+
src = np.zeros((3, 2), dtype=np.float32)
|
129 |
+
dst = np.zeros((3, 2), dtype=np.float32)
|
130 |
+
src[0, :] = center + scale_tmp * shift
|
131 |
+
src[1, :] = center + src_dir + scale_tmp * shift
|
132 |
+
dst[0, :] = [(dst_w - 1) * 0.5, (dst_h - 1) * 0.5]
|
133 |
+
dst[1, :] = np.array([(dst_w - 1) * 0.5, (dst_h - 1) * 0.5]) + dst_dir
|
134 |
+
|
135 |
+
src[2:, :] = get_3rd_point(src[0, :], src[1, :])
|
136 |
+
dst[2:, :] = get_3rd_point(dst[0, :], dst[1, :])
|
137 |
+
|
138 |
+
if inv:
|
139 |
+
trans = cv2.getAffineTransform(np.float32(dst), np.float32(src))
|
140 |
+
else:
|
141 |
+
trans = cv2.getAffineTransform(np.float32(src), np.float32(dst))
|
142 |
+
|
143 |
+
return trans
|
144 |
+
|
145 |
+
|
146 |
+
def affine_transform(pt, t):
|
147 |
+
new_pt = np.array([pt[0], pt[1], 1.0]).T
|
148 |
+
new_pt = np.dot(t, new_pt)
|
149 |
+
return new_pt[:2]
|
150 |
+
|
151 |
+
|
152 |
+
def get_3rd_point(a, b):
|
153 |
+
direct = a - b
|
154 |
+
return b + np.array([-direct[1], direct[0]], dtype=np.float32)
|
155 |
+
|
156 |
+
|
157 |
+
def get_dir(src_point, rot_rad):
|
158 |
+
sn, cs = np.sin(rot_rad), np.cos(rot_rad)
|
159 |
+
|
160 |
+
src_result = [0, 0]
|
161 |
+
src_result[0] = src_point[0] * cs - src_point[1] * sn
|
162 |
+
src_result[1] = src_point[0] * sn + src_point[1] * cs
|
163 |
+
|
164 |
+
return src_result
|
165 |
+
|
166 |
+
|
167 |
+
def crop(img, center, scale, output_size, rot=0):
|
168 |
+
trans = get_affine_transform(center, scale, rot, output_size)
|
169 |
+
|
170 |
+
dst_img = cv2.warpAffine(
|
171 |
+
img, trans, (int(output_size[1]), int(output_size[0])), flags=cv2.INTER_LINEAR
|
172 |
+
)
|
173 |
+
|
174 |
+
return dst_img
|
3rdparty/densepose/__init__.py
ADDED
@@ -0,0 +1,20 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
from .data.datasets import builtin # just to register data
|
3 |
+
from .converters import builtin as builtin_converters # register converters
|
4 |
+
from .config import (
|
5 |
+
add_densepose_config,
|
6 |
+
add_densepose_head_config,
|
7 |
+
add_hrnet_config,
|
8 |
+
add_dataset_category_config,
|
9 |
+
add_bootstrap_config,
|
10 |
+
load_bootstrap_config,
|
11 |
+
)
|
12 |
+
from .structures import DensePoseDataRelative, DensePoseList, DensePoseTransformData
|
13 |
+
from .evaluation import DensePoseCOCOEvaluator
|
14 |
+
from .modeling.roi_heads import DensePoseROIHeads
|
15 |
+
from .modeling.test_time_augmentation import (
|
16 |
+
DensePoseGeneralizedRCNNWithTTA,
|
17 |
+
DensePoseDatasetMapperTTA,
|
18 |
+
)
|
19 |
+
from .utils.transform import load_from_cfg
|
20 |
+
from .modeling.hrfpn import build_hrfpn_backbone
|
3rdparty/densepose/config.py
ADDED
@@ -0,0 +1,277 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# -*- coding = utf-8 -*-
|
2 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
3 |
+
# pyre-ignore-all-errors
|
4 |
+
|
5 |
+
from detectron2.config import CfgNode as CN
|
6 |
+
|
7 |
+
|
8 |
+
def add_dataset_category_config(cfg: CN) -> None:
|
9 |
+
"""
|
10 |
+
Add config for additional category-related dataset options
|
11 |
+
- category whitelisting
|
12 |
+
- category mapping
|
13 |
+
"""
|
14 |
+
_C = cfg
|
15 |
+
_C.DATASETS.CATEGORY_MAPS = CN(new_allowed=True)
|
16 |
+
_C.DATASETS.WHITELISTED_CATEGORIES = CN(new_allowed=True)
|
17 |
+
# class to mesh mapping
|
18 |
+
_C.DATASETS.CLASS_TO_MESH_NAME_MAPPING = CN(new_allowed=True)
|
19 |
+
|
20 |
+
|
21 |
+
def add_evaluation_config(cfg: CN) -> None:
|
22 |
+
_C = cfg
|
23 |
+
_C.DENSEPOSE_EVALUATION = CN()
|
24 |
+
# evaluator type, possible values:
|
25 |
+
# - "iou": evaluator for models that produce iou data
|
26 |
+
# - "cse": evaluator for models that produce cse data
|
27 |
+
_C.DENSEPOSE_EVALUATION.TYPE = "iou"
|
28 |
+
# storage for DensePose results, possible values:
|
29 |
+
# - "none": no explicit storage, all the results are stored in the
|
30 |
+
# dictionary with predictions, memory intensive;
|
31 |
+
# historically the default storage type
|
32 |
+
# - "ram": RAM storage, uses per-process RAM storage, which is
|
33 |
+
# reduced to a single process storage on later stages,
|
34 |
+
# less memory intensive
|
35 |
+
# - "file": file storage, uses per-process file-based storage,
|
36 |
+
# the least memory intensive, but may create bottlenecks
|
37 |
+
# on file system accesses
|
38 |
+
_C.DENSEPOSE_EVALUATION.STORAGE = "none"
|
39 |
+
# minimum threshold for IOU values: the lower its values is,
|
40 |
+
# the more matches are produced (and the higher the AP score)
|
41 |
+
_C.DENSEPOSE_EVALUATION.MIN_IOU_THRESHOLD = 0.5
|
42 |
+
# Non-distributed inference is slower (at inference time) but can avoid RAM OOM
|
43 |
+
_C.DENSEPOSE_EVALUATION.DISTRIBUTED_INFERENCE = True
|
44 |
+
# evaluate mesh alignment based on vertex embeddings, only makes sense in CSE context
|
45 |
+
_C.DENSEPOSE_EVALUATION.EVALUATE_MESH_ALIGNMENT = False
|
46 |
+
# meshes to compute mesh alignment for
|
47 |
+
_C.DENSEPOSE_EVALUATION.MESH_ALIGNMENT_MESH_NAMES = []
|
48 |
+
|
49 |
+
|
50 |
+
def add_bootstrap_config(cfg: CN) -> None:
|
51 |
+
""" """
|
52 |
+
_C = cfg
|
53 |
+
_C.BOOTSTRAP_DATASETS = []
|
54 |
+
_C.BOOTSTRAP_MODEL = CN()
|
55 |
+
_C.BOOTSTRAP_MODEL.WEIGHTS = ""
|
56 |
+
_C.BOOTSTRAP_MODEL.DEVICE = "cuda"
|
57 |
+
|
58 |
+
|
59 |
+
def get_bootstrap_dataset_config() -> CN:
|
60 |
+
_C = CN()
|
61 |
+
_C.DATASET = ""
|
62 |
+
# ratio used to mix data loaders
|
63 |
+
_C.RATIO = 0.1
|
64 |
+
# image loader
|
65 |
+
_C.IMAGE_LOADER = CN(new_allowed=True)
|
66 |
+
_C.IMAGE_LOADER.TYPE = ""
|
67 |
+
_C.IMAGE_LOADER.BATCH_SIZE = 4
|
68 |
+
_C.IMAGE_LOADER.NUM_WORKERS = 4
|
69 |
+
_C.IMAGE_LOADER.CATEGORIES = []
|
70 |
+
_C.IMAGE_LOADER.MAX_COUNT_PER_CATEGORY = 1_000_000
|
71 |
+
_C.IMAGE_LOADER.CATEGORY_TO_CLASS_MAPPING = CN(new_allowed=True)
|
72 |
+
# inference
|
73 |
+
_C.INFERENCE = CN()
|
74 |
+
# batch size for model inputs
|
75 |
+
_C.INFERENCE.INPUT_BATCH_SIZE = 4
|
76 |
+
# batch size to group model outputs
|
77 |
+
_C.INFERENCE.OUTPUT_BATCH_SIZE = 2
|
78 |
+
# sampled data
|
79 |
+
_C.DATA_SAMPLER = CN(new_allowed=True)
|
80 |
+
_C.DATA_SAMPLER.TYPE = ""
|
81 |
+
_C.DATA_SAMPLER.USE_GROUND_TRUTH_CATEGORIES = False
|
82 |
+
# filter
|
83 |
+
_C.FILTER = CN(new_allowed=True)
|
84 |
+
_C.FILTER.TYPE = ""
|
85 |
+
return _C
|
86 |
+
|
87 |
+
|
88 |
+
def load_bootstrap_config(cfg: CN) -> None:
|
89 |
+
"""
|
90 |
+
Bootstrap datasets are given as a list of `dict` that are not automatically
|
91 |
+
converted into CfgNode. This method processes all bootstrap dataset entries
|
92 |
+
and ensures that they are in CfgNode format and comply with the specification
|
93 |
+
"""
|
94 |
+
if not cfg.BOOTSTRAP_DATASETS:
|
95 |
+
return
|
96 |
+
|
97 |
+
bootstrap_datasets_cfgnodes = []
|
98 |
+
for dataset_cfg in cfg.BOOTSTRAP_DATASETS:
|
99 |
+
_C = get_bootstrap_dataset_config().clone()
|
100 |
+
_C.merge_from_other_cfg(CN(dataset_cfg))
|
101 |
+
bootstrap_datasets_cfgnodes.append(_C)
|
102 |
+
cfg.BOOTSTRAP_DATASETS = bootstrap_datasets_cfgnodes
|
103 |
+
|
104 |
+
|
105 |
+
def add_densepose_head_cse_config(cfg: CN) -> None:
|
106 |
+
"""
|
107 |
+
Add configuration options for Continuous Surface Embeddings (CSE)
|
108 |
+
"""
|
109 |
+
_C = cfg
|
110 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.CSE = CN()
|
111 |
+
# Dimensionality D of the embedding space
|
112 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBED_SIZE = 16
|
113 |
+
# Embedder specifications for various mesh IDs
|
114 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBEDDERS = CN(new_allowed=True)
|
115 |
+
# normalization coefficient for embedding distances
|
116 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBEDDING_DIST_GAUSS_SIGMA = 0.01
|
117 |
+
# normalization coefficient for geodesic distances
|
118 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.CSE.GEODESIC_DIST_GAUSS_SIGMA = 0.01
|
119 |
+
# embedding loss weight
|
120 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBED_LOSS_WEIGHT = 0.6
|
121 |
+
# embedding loss name, currently the following options are supported:
|
122 |
+
# - EmbeddingLoss: cross-entropy on vertex labels
|
123 |
+
# - SoftEmbeddingLoss: cross-entropy on vertex label combined with
|
124 |
+
# Gaussian penalty on distance between vertices
|
125 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBED_LOSS_NAME = "EmbeddingLoss"
|
126 |
+
# optimizer hyperparameters
|
127 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.CSE.FEATURES_LR_FACTOR = 1.0
|
128 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBEDDING_LR_FACTOR = 1.0
|
129 |
+
# Shape to shape cycle consistency loss parameters:
|
130 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.CSE.SHAPE_TO_SHAPE_CYCLE_LOSS = CN({"ENABLED": False})
|
131 |
+
# shape to shape cycle consistency loss weight
|
132 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.CSE.SHAPE_TO_SHAPE_CYCLE_LOSS.WEIGHT = 0.025
|
133 |
+
# norm type used for loss computation
|
134 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.CSE.SHAPE_TO_SHAPE_CYCLE_LOSS.NORM_P = 2
|
135 |
+
# normalization term for embedding similarity matrices
|
136 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.CSE.SHAPE_TO_SHAPE_CYCLE_LOSS.TEMPERATURE = 0.05
|
137 |
+
# maximum number of vertices to include into shape to shape cycle loss
|
138 |
+
# if negative or zero, all vertices are considered
|
139 |
+
# if positive, random subset of vertices of given size is considered
|
140 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.CSE.SHAPE_TO_SHAPE_CYCLE_LOSS.MAX_NUM_VERTICES = 4936
|
141 |
+
# Pixel to shape cycle consistency loss parameters:
|
142 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS = CN({"ENABLED": False})
|
143 |
+
# pixel to shape cycle consistency loss weight
|
144 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.WEIGHT = 0.0001
|
145 |
+
# norm type used for loss computation
|
146 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.NORM_P = 2
|
147 |
+
# map images to all meshes and back (if false, use only gt meshes from the batch)
|
148 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.USE_ALL_MESHES_NOT_GT_ONLY = False
|
149 |
+
# Randomly select at most this number of pixels from every instance
|
150 |
+
# if negative or zero, all vertices are considered
|
151 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.NUM_PIXELS_TO_SAMPLE = 100
|
152 |
+
# normalization factor for pixel to pixel distances (higher value = smoother distribution)
|
153 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.PIXEL_SIGMA = 5.0
|
154 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.TEMPERATURE_PIXEL_TO_VERTEX = 0.05
|
155 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.CSE.PIX_TO_SHAPE_CYCLE_LOSS.TEMPERATURE_VERTEX_TO_PIXEL = 0.05
|
156 |
+
|
157 |
+
|
158 |
+
def add_densepose_head_config(cfg: CN) -> None:
|
159 |
+
"""
|
160 |
+
Add config for densepose head.
|
161 |
+
"""
|
162 |
+
_C = cfg
|
163 |
+
|
164 |
+
_C.MODEL.DENSEPOSE_ON = True
|
165 |
+
|
166 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD = CN()
|
167 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.NAME = ""
|
168 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.NUM_STACKED_CONVS = 8
|
169 |
+
# Number of parts used for point labels
|
170 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.NUM_PATCHES = 24
|
171 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.DECONV_KERNEL = 4
|
172 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.CONV_HEAD_DIM = 512
|
173 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.CONV_HEAD_KERNEL = 3
|
174 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.UP_SCALE = 2
|
175 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.HEATMAP_SIZE = 112
|
176 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.POOLER_TYPE = "ROIAlignV2"
|
177 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.POOLER_RESOLUTION = 28
|
178 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.POOLER_SAMPLING_RATIO = 2
|
179 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.NUM_COARSE_SEGM_CHANNELS = 2 # 15 or 2
|
180 |
+
# Overlap threshold for an RoI to be considered foreground (if >= FG_IOU_THRESHOLD)
|
181 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.FG_IOU_THRESHOLD = 0.7
|
182 |
+
# Loss weights for annotation masks.(14 Parts)
|
183 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.INDEX_WEIGHTS = 5.0
|
184 |
+
# Loss weights for surface parts. (24 Parts)
|
185 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.PART_WEIGHTS = 1.0
|
186 |
+
# Loss weights for UV regression.
|
187 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.POINT_REGRESSION_WEIGHTS = 0.01
|
188 |
+
# Coarse segmentation is trained using instance segmentation task data
|
189 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.COARSE_SEGM_TRAINED_BY_MASKS = False
|
190 |
+
# For Decoder
|
191 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.DECODER_ON = True
|
192 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.DECODER_NUM_CLASSES = 256
|
193 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.DECODER_CONV_DIMS = 256
|
194 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.DECODER_NORM = ""
|
195 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.DECODER_COMMON_STRIDE = 4
|
196 |
+
# For DeepLab head
|
197 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.DEEPLAB = CN()
|
198 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.DEEPLAB.NORM = "GN"
|
199 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.DEEPLAB.NONLOCAL_ON = 0
|
200 |
+
# Predictor class name, must be registered in DENSEPOSE_PREDICTOR_REGISTRY
|
201 |
+
# Some registered predictors:
|
202 |
+
# "DensePoseChartPredictor": predicts segmentation and UV coordinates for predefined charts
|
203 |
+
# "DensePoseChartWithConfidencePredictor": predicts segmentation, UV coordinates
|
204 |
+
# and associated confidences for predefined charts (default)
|
205 |
+
# "DensePoseEmbeddingWithConfidencePredictor": predicts segmentation, embeddings
|
206 |
+
# and associated confidences for CSE
|
207 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.PREDICTOR_NAME = "DensePoseChartWithConfidencePredictor"
|
208 |
+
# Loss class name, must be registered in DENSEPOSE_LOSS_REGISTRY
|
209 |
+
# Some registered losses:
|
210 |
+
# "DensePoseChartLoss": loss for chart-based models that estimate
|
211 |
+
# segmentation and UV coordinates
|
212 |
+
# "DensePoseChartWithConfidenceLoss": loss for chart-based models that estimate
|
213 |
+
# segmentation, UV coordinates and the corresponding confidences (default)
|
214 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.LOSS_NAME = "DensePoseChartWithConfidenceLoss"
|
215 |
+
# Confidences
|
216 |
+
# Enable learning UV confidences (variances) along with the actual values
|
217 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.UV_CONFIDENCE = CN({"ENABLED": False})
|
218 |
+
# UV confidence lower bound
|
219 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.UV_CONFIDENCE.EPSILON = 0.01
|
220 |
+
# Enable learning segmentation confidences (variances) along with the actual values
|
221 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.SEGM_CONFIDENCE = CN({"ENABLED": False})
|
222 |
+
# Segmentation confidence lower bound
|
223 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.SEGM_CONFIDENCE.EPSILON = 0.01
|
224 |
+
# Statistical model type for confidence learning, possible values:
|
225 |
+
# - "iid_iso": statistically independent identically distributed residuals
|
226 |
+
# with isotropic covariance
|
227 |
+
# - "indep_aniso": statistically independent residuals with anisotropic
|
228 |
+
# covariances
|
229 |
+
_C.MODEL.ROI_DENSEPOSE_HEAD.UV_CONFIDENCE.TYPE = "iid_iso"
|
230 |
+
# List of angles for rotation in data augmentation during training
|
231 |
+
_C.INPUT.ROTATION_ANGLES = [0]
|
232 |
+
_C.TEST.AUG.ROTATION_ANGLES = () # Rotation TTA
|
233 |
+
|
234 |
+
add_densepose_head_cse_config(cfg)
|
235 |
+
|
236 |
+
|
237 |
+
def add_hrnet_config(cfg: CN) -> None:
|
238 |
+
"""
|
239 |
+
Add config for HRNet backbone.
|
240 |
+
"""
|
241 |
+
_C = cfg
|
242 |
+
|
243 |
+
# For HigherHRNet w32
|
244 |
+
_C.MODEL.HRNET = CN()
|
245 |
+
_C.MODEL.HRNET.STEM_INPLANES = 64
|
246 |
+
_C.MODEL.HRNET.STAGE2 = CN()
|
247 |
+
_C.MODEL.HRNET.STAGE2.NUM_MODULES = 1
|
248 |
+
_C.MODEL.HRNET.STAGE2.NUM_BRANCHES = 2
|
249 |
+
_C.MODEL.HRNET.STAGE2.BLOCK = "BASIC"
|
250 |
+
_C.MODEL.HRNET.STAGE2.NUM_BLOCKS = [4, 4]
|
251 |
+
_C.MODEL.HRNET.STAGE2.NUM_CHANNELS = [32, 64]
|
252 |
+
_C.MODEL.HRNET.STAGE2.FUSE_METHOD = "SUM"
|
253 |
+
_C.MODEL.HRNET.STAGE3 = CN()
|
254 |
+
_C.MODEL.HRNET.STAGE3.NUM_MODULES = 4
|
255 |
+
_C.MODEL.HRNET.STAGE3.NUM_BRANCHES = 3
|
256 |
+
_C.MODEL.HRNET.STAGE3.BLOCK = "BASIC"
|
257 |
+
_C.MODEL.HRNET.STAGE3.NUM_BLOCKS = [4, 4, 4]
|
258 |
+
_C.MODEL.HRNET.STAGE3.NUM_CHANNELS = [32, 64, 128]
|
259 |
+
_C.MODEL.HRNET.STAGE3.FUSE_METHOD = "SUM"
|
260 |
+
_C.MODEL.HRNET.STAGE4 = CN()
|
261 |
+
_C.MODEL.HRNET.STAGE4.NUM_MODULES = 3
|
262 |
+
_C.MODEL.HRNET.STAGE4.NUM_BRANCHES = 4
|
263 |
+
_C.MODEL.HRNET.STAGE4.BLOCK = "BASIC"
|
264 |
+
_C.MODEL.HRNET.STAGE4.NUM_BLOCKS = [4, 4, 4, 4]
|
265 |
+
_C.MODEL.HRNET.STAGE4.NUM_CHANNELS = [32, 64, 128, 256]
|
266 |
+
_C.MODEL.HRNET.STAGE4.FUSE_METHOD = "SUM"
|
267 |
+
|
268 |
+
_C.MODEL.HRNET.HRFPN = CN()
|
269 |
+
_C.MODEL.HRNET.HRFPN.OUT_CHANNELS = 256
|
270 |
+
|
271 |
+
|
272 |
+
def add_densepose_config(cfg: CN) -> None:
|
273 |
+
add_densepose_head_config(cfg)
|
274 |
+
add_hrnet_config(cfg)
|
275 |
+
add_bootstrap_config(cfg)
|
276 |
+
add_dataset_category_config(cfg)
|
277 |
+
add_evaluation_config(cfg)
|
3rdparty/densepose/converters/__init__.py
ADDED
@@ -0,0 +1,15 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
|
3 |
+
from .hflip import HFlipConverter
|
4 |
+
from .to_mask import ToMaskConverter
|
5 |
+
from .to_chart_result import ToChartResultConverter, ToChartResultConverterWithConfidences
|
6 |
+
from .segm_to_mask import (
|
7 |
+
predictor_output_with_fine_and_coarse_segm_to_mask,
|
8 |
+
predictor_output_with_coarse_segm_to_mask,
|
9 |
+
resample_fine_and_coarse_segm_to_bbox,
|
10 |
+
)
|
11 |
+
from .chart_output_to_chart_result import (
|
12 |
+
densepose_chart_predictor_output_to_result,
|
13 |
+
densepose_chart_predictor_output_to_result_with_confidences,
|
14 |
+
)
|
15 |
+
from .chart_output_hflip import densepose_chart_predictor_output_hflip
|
3rdparty/densepose/converters/base.py
ADDED
@@ -0,0 +1,93 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
|
3 |
+
from typing import Any, Tuple, Type
|
4 |
+
import torch
|
5 |
+
|
6 |
+
|
7 |
+
class BaseConverter:
|
8 |
+
"""
|
9 |
+
Converter base class to be reused by various converters.
|
10 |
+
Converter allows one to convert data from various source types to a particular
|
11 |
+
destination type. Each source type needs to register its converter. The
|
12 |
+
registration for each source type is valid for all descendants of that type.
|
13 |
+
"""
|
14 |
+
|
15 |
+
@classmethod
|
16 |
+
def register(cls, from_type: Type, converter: Any = None):
|
17 |
+
"""
|
18 |
+
Registers a converter for the specified type.
|
19 |
+
Can be used as a decorator (if converter is None), or called as a method.
|
20 |
+
|
21 |
+
Args:
|
22 |
+
from_type (type): type to register the converter for;
|
23 |
+
all instances of this type will use the same converter
|
24 |
+
converter (callable): converter to be registered for the given
|
25 |
+
type; if None, this method is assumed to be a decorator for the converter
|
26 |
+
"""
|
27 |
+
|
28 |
+
if converter is not None:
|
29 |
+
cls._do_register(from_type, converter)
|
30 |
+
|
31 |
+
def wrapper(converter: Any) -> Any:
|
32 |
+
cls._do_register(from_type, converter)
|
33 |
+
return converter
|
34 |
+
|
35 |
+
return wrapper
|
36 |
+
|
37 |
+
@classmethod
|
38 |
+
def _do_register(cls, from_type: Type, converter: Any):
|
39 |
+
cls.registry[from_type] = converter # pyre-ignore[16]
|
40 |
+
|
41 |
+
@classmethod
|
42 |
+
def _lookup_converter(cls, from_type: Type) -> Any:
|
43 |
+
"""
|
44 |
+
Perform recursive lookup for the given type
|
45 |
+
to find registered converter. If a converter was found for some base
|
46 |
+
class, it gets registered for this class to save on further lookups.
|
47 |
+
|
48 |
+
Args:
|
49 |
+
from_type: type for which to find a converter
|
50 |
+
Return:
|
51 |
+
callable or None - registered converter or None
|
52 |
+
if no suitable entry was found in the registry
|
53 |
+
"""
|
54 |
+
if from_type in cls.registry: # pyre-ignore[16]
|
55 |
+
return cls.registry[from_type]
|
56 |
+
for base in from_type.__bases__:
|
57 |
+
converter = cls._lookup_converter(base)
|
58 |
+
if converter is not None:
|
59 |
+
cls._do_register(from_type, converter)
|
60 |
+
return converter
|
61 |
+
return None
|
62 |
+
|
63 |
+
@classmethod
|
64 |
+
def convert(cls, instance: Any, *args, **kwargs):
|
65 |
+
"""
|
66 |
+
Convert an instance to the destination type using some registered
|
67 |
+
converter. Does recursive lookup for base classes, so there's no need
|
68 |
+
for explicit registration for derived classes.
|
69 |
+
|
70 |
+
Args:
|
71 |
+
instance: source instance to convert to the destination type
|
72 |
+
Return:
|
73 |
+
An instance of the destination type obtained from the source instance
|
74 |
+
Raises KeyError, if no suitable converter found
|
75 |
+
"""
|
76 |
+
instance_type = type(instance)
|
77 |
+
converter = cls._lookup_converter(instance_type)
|
78 |
+
if converter is None:
|
79 |
+
if cls.dst_type is None: # pyre-ignore[16]
|
80 |
+
output_type_str = "itself"
|
81 |
+
else:
|
82 |
+
output_type_str = cls.dst_type
|
83 |
+
raise KeyError(f"Could not find converter from {instance_type} to {output_type_str}")
|
84 |
+
return converter(instance, *args, **kwargs)
|
85 |
+
|
86 |
+
|
87 |
+
IntTupleBox = Tuple[int, int, int, int]
|
88 |
+
|
89 |
+
|
90 |
+
def make_int_box(box: torch.Tensor) -> IntTupleBox:
|
91 |
+
int_box = [0, 0, 0, 0]
|
92 |
+
int_box[0], int_box[1], int_box[2], int_box[3] = tuple(box.long().tolist())
|
93 |
+
return int_box[0], int_box[1], int_box[2], int_box[3]
|
3rdparty/densepose/converters/builtin.py
ADDED
@@ -0,0 +1,31 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
|
3 |
+
from ..structures import DensePoseChartPredictorOutput, DensePoseEmbeddingPredictorOutput
|
4 |
+
from . import (
|
5 |
+
HFlipConverter,
|
6 |
+
ToChartResultConverter,
|
7 |
+
ToChartResultConverterWithConfidences,
|
8 |
+
ToMaskConverter,
|
9 |
+
densepose_chart_predictor_output_hflip,
|
10 |
+
densepose_chart_predictor_output_to_result,
|
11 |
+
densepose_chart_predictor_output_to_result_with_confidences,
|
12 |
+
predictor_output_with_coarse_segm_to_mask,
|
13 |
+
predictor_output_with_fine_and_coarse_segm_to_mask,
|
14 |
+
)
|
15 |
+
|
16 |
+
ToMaskConverter.register(
|
17 |
+
DensePoseChartPredictorOutput, predictor_output_with_fine_and_coarse_segm_to_mask
|
18 |
+
)
|
19 |
+
ToMaskConverter.register(
|
20 |
+
DensePoseEmbeddingPredictorOutput, predictor_output_with_coarse_segm_to_mask
|
21 |
+
)
|
22 |
+
|
23 |
+
ToChartResultConverter.register(
|
24 |
+
DensePoseChartPredictorOutput, densepose_chart_predictor_output_to_result
|
25 |
+
)
|
26 |
+
|
27 |
+
ToChartResultConverterWithConfidences.register(
|
28 |
+
DensePoseChartPredictorOutput, densepose_chart_predictor_output_to_result_with_confidences
|
29 |
+
)
|
30 |
+
|
31 |
+
HFlipConverter.register(DensePoseChartPredictorOutput, densepose_chart_predictor_output_hflip)
|
3rdparty/densepose/converters/chart_output_hflip.py
ADDED
@@ -0,0 +1,71 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
from dataclasses import fields
|
3 |
+
import torch
|
4 |
+
|
5 |
+
from densepose.structures import DensePoseChartPredictorOutput, DensePoseTransformData
|
6 |
+
|
7 |
+
|
8 |
+
def densepose_chart_predictor_output_hflip(
|
9 |
+
densepose_predictor_output: DensePoseChartPredictorOutput,
|
10 |
+
transform_data: DensePoseTransformData,
|
11 |
+
) -> DensePoseChartPredictorOutput:
|
12 |
+
"""
|
13 |
+
Change to take into account a Horizontal flip.
|
14 |
+
"""
|
15 |
+
if len(densepose_predictor_output) > 0:
|
16 |
+
|
17 |
+
PredictorOutput = type(densepose_predictor_output)
|
18 |
+
output_dict = {}
|
19 |
+
|
20 |
+
for field in fields(densepose_predictor_output):
|
21 |
+
field_value = getattr(densepose_predictor_output, field.name)
|
22 |
+
# flip tensors
|
23 |
+
if isinstance(field_value, torch.Tensor):
|
24 |
+
setattr(densepose_predictor_output, field.name, torch.flip(field_value, [3]))
|
25 |
+
|
26 |
+
densepose_predictor_output = _flip_iuv_semantics_tensor(
|
27 |
+
densepose_predictor_output, transform_data
|
28 |
+
)
|
29 |
+
densepose_predictor_output = _flip_segm_semantics_tensor(
|
30 |
+
densepose_predictor_output, transform_data
|
31 |
+
)
|
32 |
+
|
33 |
+
for field in fields(densepose_predictor_output):
|
34 |
+
output_dict[field.name] = getattr(densepose_predictor_output, field.name)
|
35 |
+
|
36 |
+
return PredictorOutput(**output_dict)
|
37 |
+
else:
|
38 |
+
return densepose_predictor_output
|
39 |
+
|
40 |
+
|
41 |
+
def _flip_iuv_semantics_tensor(
|
42 |
+
densepose_predictor_output: DensePoseChartPredictorOutput,
|
43 |
+
dp_transform_data: DensePoseTransformData,
|
44 |
+
) -> DensePoseChartPredictorOutput:
|
45 |
+
point_label_symmetries = dp_transform_data.point_label_symmetries
|
46 |
+
uv_symmetries = dp_transform_data.uv_symmetries
|
47 |
+
|
48 |
+
N, C, H, W = densepose_predictor_output.u.shape
|
49 |
+
u_loc = (densepose_predictor_output.u[:, 1:, :, :].clamp(0, 1) * 255).long()
|
50 |
+
v_loc = (densepose_predictor_output.v[:, 1:, :, :].clamp(0, 1) * 255).long()
|
51 |
+
Iindex = torch.arange(C - 1, device=densepose_predictor_output.u.device)[
|
52 |
+
None, :, None, None
|
53 |
+
].expand(N, C - 1, H, W)
|
54 |
+
densepose_predictor_output.u[:, 1:, :, :] = uv_symmetries["U_transforms"][Iindex, v_loc, u_loc]
|
55 |
+
densepose_predictor_output.v[:, 1:, :, :] = uv_symmetries["V_transforms"][Iindex, v_loc, u_loc]
|
56 |
+
|
57 |
+
for el in ["fine_segm", "u", "v"]:
|
58 |
+
densepose_predictor_output.__dict__[el] = densepose_predictor_output.__dict__[el][
|
59 |
+
:, point_label_symmetries, :, :
|
60 |
+
]
|
61 |
+
return densepose_predictor_output
|
62 |
+
|
63 |
+
|
64 |
+
def _flip_segm_semantics_tensor(
|
65 |
+
densepose_predictor_output: DensePoseChartPredictorOutput, dp_transform_data
|
66 |
+
):
|
67 |
+
if densepose_predictor_output.coarse_segm.shape[1] > 2:
|
68 |
+
densepose_predictor_output.coarse_segm = densepose_predictor_output.coarse_segm[
|
69 |
+
:, dp_transform_data.mask_label_symmetries, :, :
|
70 |
+
]
|
71 |
+
return densepose_predictor_output
|
3rdparty/densepose/converters/chart_output_to_chart_result.py
ADDED
@@ -0,0 +1,188 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
|
3 |
+
from typing import Dict
|
4 |
+
import torch
|
5 |
+
from torch.nn import functional as F
|
6 |
+
|
7 |
+
from detectron2.structures.boxes import Boxes, BoxMode
|
8 |
+
|
9 |
+
from ..structures import (
|
10 |
+
DensePoseChartPredictorOutput,
|
11 |
+
DensePoseChartResult,
|
12 |
+
DensePoseChartResultWithConfidences,
|
13 |
+
)
|
14 |
+
from . import resample_fine_and_coarse_segm_to_bbox
|
15 |
+
from .base import IntTupleBox, make_int_box
|
16 |
+
|
17 |
+
|
18 |
+
def resample_uv_tensors_to_bbox(
|
19 |
+
u: torch.Tensor,
|
20 |
+
v: torch.Tensor,
|
21 |
+
labels: torch.Tensor,
|
22 |
+
box_xywh_abs: IntTupleBox,
|
23 |
+
) -> torch.Tensor:
|
24 |
+
"""
|
25 |
+
Resamples U and V coordinate estimates for the given bounding box
|
26 |
+
|
27 |
+
Args:
|
28 |
+
u (tensor [1, C, H, W] of float): U coordinates
|
29 |
+
v (tensor [1, C, H, W] of float): V coordinates
|
30 |
+
labels (tensor [H, W] of long): labels obtained by resampling segmentation
|
31 |
+
outputs for the given bounding box
|
32 |
+
box_xywh_abs (tuple of 4 int): bounding box that corresponds to predictor outputs
|
33 |
+
Return:
|
34 |
+
Resampled U and V coordinates - a tensor [2, H, W] of float
|
35 |
+
"""
|
36 |
+
x, y, w, h = box_xywh_abs
|
37 |
+
w = max(int(w), 1)
|
38 |
+
h = max(int(h), 1)
|
39 |
+
u_bbox = F.interpolate(u, (h, w), mode="bilinear", align_corners=False)
|
40 |
+
v_bbox = F.interpolate(v, (h, w), mode="bilinear", align_corners=False)
|
41 |
+
uv = torch.zeros([2, h, w], dtype=torch.float32, device=u.device)
|
42 |
+
for part_id in range(1, u_bbox.size(1)):
|
43 |
+
uv[0][labels == part_id] = u_bbox[0, part_id][labels == part_id]
|
44 |
+
uv[1][labels == part_id] = v_bbox[0, part_id][labels == part_id]
|
45 |
+
return uv
|
46 |
+
|
47 |
+
|
48 |
+
def resample_uv_to_bbox(
|
49 |
+
predictor_output: DensePoseChartPredictorOutput,
|
50 |
+
labels: torch.Tensor,
|
51 |
+
box_xywh_abs: IntTupleBox,
|
52 |
+
) -> torch.Tensor:
|
53 |
+
"""
|
54 |
+
Resamples U and V coordinate estimates for the given bounding box
|
55 |
+
|
56 |
+
Args:
|
57 |
+
predictor_output (DensePoseChartPredictorOutput): DensePose predictor
|
58 |
+
output to be resampled
|
59 |
+
labels (tensor [H, W] of long): labels obtained by resampling segmentation
|
60 |
+
outputs for the given bounding box
|
61 |
+
box_xywh_abs (tuple of 4 int): bounding box that corresponds to predictor outputs
|
62 |
+
Return:
|
63 |
+
Resampled U and V coordinates - a tensor [2, H, W] of float
|
64 |
+
"""
|
65 |
+
return resample_uv_tensors_to_bbox(
|
66 |
+
predictor_output.u,
|
67 |
+
predictor_output.v,
|
68 |
+
labels,
|
69 |
+
box_xywh_abs,
|
70 |
+
)
|
71 |
+
|
72 |
+
|
73 |
+
def densepose_chart_predictor_output_to_result(
|
74 |
+
predictor_output: DensePoseChartPredictorOutput, boxes: Boxes
|
75 |
+
) -> DensePoseChartResult:
|
76 |
+
"""
|
77 |
+
Convert densepose chart predictor outputs to results
|
78 |
+
|
79 |
+
Args:
|
80 |
+
predictor_output (DensePoseChartPredictorOutput): DensePose predictor
|
81 |
+
output to be converted to results, must contain only 1 output
|
82 |
+
boxes (Boxes): bounding box that corresponds to the predictor output,
|
83 |
+
must contain only 1 bounding box
|
84 |
+
Return:
|
85 |
+
DensePose chart-based result (DensePoseChartResult)
|
86 |
+
"""
|
87 |
+
assert len(predictor_output) == 1 and len(boxes) == 1, (
|
88 |
+
f"Predictor output to result conversion can operate only single outputs"
|
89 |
+
f", got {len(predictor_output)} predictor outputs and {len(boxes)} boxes"
|
90 |
+
)
|
91 |
+
|
92 |
+
boxes_xyxy_abs = boxes.tensor.clone()
|
93 |
+
boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
|
94 |
+
box_xywh = make_int_box(boxes_xywh_abs[0])
|
95 |
+
|
96 |
+
labels = resample_fine_and_coarse_segm_to_bbox(predictor_output, box_xywh).squeeze(0)
|
97 |
+
uv = resample_uv_to_bbox(predictor_output, labels, box_xywh)
|
98 |
+
return DensePoseChartResult(labels=labels, uv=uv)
|
99 |
+
|
100 |
+
|
101 |
+
def resample_confidences_to_bbox(
|
102 |
+
predictor_output: DensePoseChartPredictorOutput,
|
103 |
+
labels: torch.Tensor,
|
104 |
+
box_xywh_abs: IntTupleBox,
|
105 |
+
) -> Dict[str, torch.Tensor]:
|
106 |
+
"""
|
107 |
+
Resamples confidences for the given bounding box
|
108 |
+
|
109 |
+
Args:
|
110 |
+
predictor_output (DensePoseChartPredictorOutput): DensePose predictor
|
111 |
+
output to be resampled
|
112 |
+
labels (tensor [H, W] of long): labels obtained by resampling segmentation
|
113 |
+
outputs for the given bounding box
|
114 |
+
box_xywh_abs (tuple of 4 int): bounding box that corresponds to predictor outputs
|
115 |
+
Return:
|
116 |
+
Resampled confidences - a dict of [H, W] tensors of float
|
117 |
+
"""
|
118 |
+
|
119 |
+
x, y, w, h = box_xywh_abs
|
120 |
+
w = max(int(w), 1)
|
121 |
+
h = max(int(h), 1)
|
122 |
+
|
123 |
+
confidence_names = [
|
124 |
+
"sigma_1",
|
125 |
+
"sigma_2",
|
126 |
+
"kappa_u",
|
127 |
+
"kappa_v",
|
128 |
+
"fine_segm_confidence",
|
129 |
+
"coarse_segm_confidence",
|
130 |
+
]
|
131 |
+
confidence_results = {key: None for key in confidence_names}
|
132 |
+
confidence_names = [
|
133 |
+
key for key in confidence_names if getattr(predictor_output, key) is not None
|
134 |
+
]
|
135 |
+
confidence_base = torch.zeros([h, w], dtype=torch.float32, device=predictor_output.u.device)
|
136 |
+
|
137 |
+
# assign data from channels that correspond to the labels
|
138 |
+
for key in confidence_names:
|
139 |
+
resampled_confidence = F.interpolate(
|
140 |
+
getattr(predictor_output, key),
|
141 |
+
(h, w),
|
142 |
+
mode="bilinear",
|
143 |
+
align_corners=False,
|
144 |
+
)
|
145 |
+
result = confidence_base.clone()
|
146 |
+
for part_id in range(1, predictor_output.u.size(1)):
|
147 |
+
if resampled_confidence.size(1) != predictor_output.u.size(1):
|
148 |
+
# confidence is not part-based, don't try to fill it part by part
|
149 |
+
continue
|
150 |
+
result[labels == part_id] = resampled_confidence[0, part_id][labels == part_id]
|
151 |
+
|
152 |
+
if resampled_confidence.size(1) != predictor_output.u.size(1):
|
153 |
+
# confidence is not part-based, fill the data with the first channel
|
154 |
+
# (targeted for segmentation confidences that have only 1 channel)
|
155 |
+
result = resampled_confidence[0, 0]
|
156 |
+
|
157 |
+
confidence_results[key] = result
|
158 |
+
|
159 |
+
return confidence_results # pyre-ignore[7]
|
160 |
+
|
161 |
+
|
162 |
+
def densepose_chart_predictor_output_to_result_with_confidences(
|
163 |
+
predictor_output: DensePoseChartPredictorOutput, boxes: Boxes
|
164 |
+
) -> DensePoseChartResultWithConfidences:
|
165 |
+
"""
|
166 |
+
Convert densepose chart predictor outputs to results
|
167 |
+
|
168 |
+
Args:
|
169 |
+
predictor_output (DensePoseChartPredictorOutput): DensePose predictor
|
170 |
+
output with confidences to be converted to results, must contain only 1 output
|
171 |
+
boxes (Boxes): bounding box that corresponds to the predictor output,
|
172 |
+
must contain only 1 bounding box
|
173 |
+
Return:
|
174 |
+
DensePose chart-based result with confidences (DensePoseChartResultWithConfidences)
|
175 |
+
"""
|
176 |
+
assert len(predictor_output) == 1 and len(boxes) == 1, (
|
177 |
+
f"Predictor output to result conversion can operate only single outputs"
|
178 |
+
f", got {len(predictor_output)} predictor outputs and {len(boxes)} boxes"
|
179 |
+
)
|
180 |
+
|
181 |
+
boxes_xyxy_abs = boxes.tensor.clone()
|
182 |
+
boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
|
183 |
+
box_xywh = make_int_box(boxes_xywh_abs[0])
|
184 |
+
|
185 |
+
labels = resample_fine_and_coarse_segm_to_bbox(predictor_output, box_xywh).squeeze(0)
|
186 |
+
uv = resample_uv_to_bbox(predictor_output, labels, box_xywh)
|
187 |
+
confidences = resample_confidences_to_bbox(predictor_output, labels, box_xywh)
|
188 |
+
return DensePoseChartResultWithConfidences(labels=labels, uv=uv, **confidences)
|
3rdparty/densepose/converters/hflip.py
ADDED
@@ -0,0 +1,34 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
|
3 |
+
from typing import Any
|
4 |
+
|
5 |
+
from .base import BaseConverter
|
6 |
+
|
7 |
+
|
8 |
+
class HFlipConverter(BaseConverter):
|
9 |
+
"""
|
10 |
+
Converts various DensePose predictor outputs to DensePose results.
|
11 |
+
Each DensePose predictor output type has to register its convertion strategy.
|
12 |
+
"""
|
13 |
+
|
14 |
+
registry = {}
|
15 |
+
dst_type = None
|
16 |
+
|
17 |
+
@classmethod
|
18 |
+
# pyre-fixme[14]: `convert` overrides method defined in `BaseConverter`
|
19 |
+
# inconsistently.
|
20 |
+
def convert(cls, predictor_outputs: Any, transform_data: Any, *args, **kwargs):
|
21 |
+
"""
|
22 |
+
Performs an horizontal flip on DensePose predictor outputs.
|
23 |
+
Does recursive lookup for base classes, so there's no need
|
24 |
+
for explicit registration for derived classes.
|
25 |
+
|
26 |
+
Args:
|
27 |
+
predictor_outputs: DensePose predictor output to be converted to BitMasks
|
28 |
+
transform_data: Anything useful for the flip
|
29 |
+
Return:
|
30 |
+
An instance of the same type as predictor_outputs
|
31 |
+
"""
|
32 |
+
return super(HFlipConverter, cls).convert(
|
33 |
+
predictor_outputs, transform_data, *args, **kwargs
|
34 |
+
)
|
3rdparty/densepose/converters/segm_to_mask.py
ADDED
@@ -0,0 +1,150 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
|
3 |
+
from typing import Any
|
4 |
+
import torch
|
5 |
+
from torch.nn import functional as F
|
6 |
+
|
7 |
+
from detectron2.structures import BitMasks, Boxes, BoxMode
|
8 |
+
|
9 |
+
from .base import IntTupleBox, make_int_box
|
10 |
+
from .to_mask import ImageSizeType
|
11 |
+
|
12 |
+
|
13 |
+
def resample_coarse_segm_tensor_to_bbox(coarse_segm: torch.Tensor, box_xywh_abs: IntTupleBox):
|
14 |
+
"""
|
15 |
+
Resample coarse segmentation tensor to the given
|
16 |
+
bounding box and derive labels for each pixel of the bounding box
|
17 |
+
|
18 |
+
Args:
|
19 |
+
coarse_segm: float tensor of shape [1, K, Hout, Wout]
|
20 |
+
box_xywh_abs (tuple of 4 int): bounding box given by its upper-left
|
21 |
+
corner coordinates, width (W) and height (H)
|
22 |
+
Return:
|
23 |
+
Labels for each pixel of the bounding box, a long tensor of size [1, H, W]
|
24 |
+
"""
|
25 |
+
x, y, w, h = box_xywh_abs
|
26 |
+
w = max(int(w), 1)
|
27 |
+
h = max(int(h), 1)
|
28 |
+
labels = F.interpolate(coarse_segm, (h, w), mode="bilinear", align_corners=False).argmax(dim=1)
|
29 |
+
return labels
|
30 |
+
|
31 |
+
|
32 |
+
def resample_fine_and_coarse_segm_tensors_to_bbox(
|
33 |
+
fine_segm: torch.Tensor, coarse_segm: torch.Tensor, box_xywh_abs: IntTupleBox
|
34 |
+
):
|
35 |
+
"""
|
36 |
+
Resample fine and coarse segmentation tensors to the given
|
37 |
+
bounding box and derive labels for each pixel of the bounding box
|
38 |
+
|
39 |
+
Args:
|
40 |
+
fine_segm: float tensor of shape [1, C, Hout, Wout]
|
41 |
+
coarse_segm: float tensor of shape [1, K, Hout, Wout]
|
42 |
+
box_xywh_abs (tuple of 4 int): bounding box given by its upper-left
|
43 |
+
corner coordinates, width (W) and height (H)
|
44 |
+
Return:
|
45 |
+
Labels for each pixel of the bounding box, a long tensor of size [1, H, W]
|
46 |
+
"""
|
47 |
+
x, y, w, h = box_xywh_abs
|
48 |
+
w = max(int(w), 1)
|
49 |
+
h = max(int(h), 1)
|
50 |
+
# coarse segmentation
|
51 |
+
coarse_segm_bbox = F.interpolate(
|
52 |
+
coarse_segm,
|
53 |
+
(h, w),
|
54 |
+
mode="bilinear",
|
55 |
+
align_corners=False,
|
56 |
+
).argmax(dim=1)
|
57 |
+
# combined coarse and fine segmentation
|
58 |
+
labels = (
|
59 |
+
F.interpolate(fine_segm, (h, w), mode="bilinear", align_corners=False).argmax(dim=1)
|
60 |
+
* (coarse_segm_bbox > 0).long()
|
61 |
+
)
|
62 |
+
return labels
|
63 |
+
|
64 |
+
|
65 |
+
def resample_fine_and_coarse_segm_to_bbox(predictor_output: Any, box_xywh_abs: IntTupleBox):
|
66 |
+
"""
|
67 |
+
Resample fine and coarse segmentation outputs from a predictor to the given
|
68 |
+
bounding box and derive labels for each pixel of the bounding box
|
69 |
+
|
70 |
+
Args:
|
71 |
+
predictor_output: DensePose predictor output that contains segmentation
|
72 |
+
results to be resampled
|
73 |
+
box_xywh_abs (tuple of 4 int): bounding box given by its upper-left
|
74 |
+
corner coordinates, width (W) and height (H)
|
75 |
+
Return:
|
76 |
+
Labels for each pixel of the bounding box, a long tensor of size [1, H, W]
|
77 |
+
"""
|
78 |
+
return resample_fine_and_coarse_segm_tensors_to_bbox(
|
79 |
+
predictor_output.fine_segm,
|
80 |
+
predictor_output.coarse_segm,
|
81 |
+
box_xywh_abs,
|
82 |
+
)
|
83 |
+
|
84 |
+
|
85 |
+
def predictor_output_with_coarse_segm_to_mask(
|
86 |
+
predictor_output: Any, boxes: Boxes, image_size_hw: ImageSizeType
|
87 |
+
) -> BitMasks:
|
88 |
+
"""
|
89 |
+
Convert predictor output with coarse and fine segmentation to a mask.
|
90 |
+
Assumes that predictor output has the following attributes:
|
91 |
+
- coarse_segm (tensor of size [N, D, H, W]): coarse segmentation
|
92 |
+
unnormalized scores for N instances; D is the number of coarse
|
93 |
+
segmentation labels, H and W is the resolution of the estimate
|
94 |
+
|
95 |
+
Args:
|
96 |
+
predictor_output: DensePose predictor output to be converted to mask
|
97 |
+
boxes (Boxes): bounding boxes that correspond to the DensePose
|
98 |
+
predictor outputs
|
99 |
+
image_size_hw (tuple [int, int]): image height Himg and width Wimg
|
100 |
+
Return:
|
101 |
+
BitMasks that contain a bool tensor of size [N, Himg, Wimg] with
|
102 |
+
a mask of the size of the image for each instance
|
103 |
+
"""
|
104 |
+
H, W = image_size_hw
|
105 |
+
boxes_xyxy_abs = boxes.tensor.clone()
|
106 |
+
boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
|
107 |
+
N = len(boxes_xywh_abs)
|
108 |
+
masks = torch.zeros((N, H, W), dtype=torch.bool, device=boxes.tensor.device)
|
109 |
+
for i in range(len(boxes_xywh_abs)):
|
110 |
+
box_xywh = make_int_box(boxes_xywh_abs[i])
|
111 |
+
box_mask = resample_coarse_segm_tensor_to_bbox(predictor_output[i].coarse_segm, box_xywh)
|
112 |
+
x, y, w, h = box_xywh
|
113 |
+
masks[i, y : y + h, x : x + w] = box_mask
|
114 |
+
|
115 |
+
return BitMasks(masks)
|
116 |
+
|
117 |
+
|
118 |
+
def predictor_output_with_fine_and_coarse_segm_to_mask(
|
119 |
+
predictor_output: Any, boxes: Boxes, image_size_hw: ImageSizeType
|
120 |
+
) -> BitMasks:
|
121 |
+
"""
|
122 |
+
Convert predictor output with coarse and fine segmentation to a mask.
|
123 |
+
Assumes that predictor output has the following attributes:
|
124 |
+
- coarse_segm (tensor of size [N, D, H, W]): coarse segmentation
|
125 |
+
unnormalized scores for N instances; D is the number of coarse
|
126 |
+
segmentation labels, H and W is the resolution of the estimate
|
127 |
+
- fine_segm (tensor of size [N, C, H, W]): fine segmentation
|
128 |
+
unnormalized scores for N instances; C is the number of fine
|
129 |
+
segmentation labels, H and W is the resolution of the estimate
|
130 |
+
|
131 |
+
Args:
|
132 |
+
predictor_output: DensePose predictor output to be converted to mask
|
133 |
+
boxes (Boxes): bounding boxes that correspond to the DensePose
|
134 |
+
predictor outputs
|
135 |
+
image_size_hw (tuple [int, int]): image height Himg and width Wimg
|
136 |
+
Return:
|
137 |
+
BitMasks that contain a bool tensor of size [N, Himg, Wimg] with
|
138 |
+
a mask of the size of the image for each instance
|
139 |
+
"""
|
140 |
+
H, W = image_size_hw
|
141 |
+
boxes_xyxy_abs = boxes.tensor.clone()
|
142 |
+
boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
|
143 |
+
N = len(boxes_xywh_abs)
|
144 |
+
masks = torch.zeros((N, H, W), dtype=torch.bool, device=boxes.tensor.device)
|
145 |
+
for i in range(len(boxes_xywh_abs)):
|
146 |
+
box_xywh = make_int_box(boxes_xywh_abs[i])
|
147 |
+
labels_i = resample_fine_and_coarse_segm_to_bbox(predictor_output[i], box_xywh)
|
148 |
+
x, y, w, h = box_xywh
|
149 |
+
masks[i, y : y + h, x : x + w] = labels_i > 0
|
150 |
+
return BitMasks(masks)
|
3rdparty/densepose/converters/to_chart_result.py
ADDED
@@ -0,0 +1,70 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
|
3 |
+
from typing import Any
|
4 |
+
|
5 |
+
from detectron2.structures import Boxes
|
6 |
+
|
7 |
+
from ..structures import DensePoseChartResult, DensePoseChartResultWithConfidences
|
8 |
+
from .base import BaseConverter
|
9 |
+
|
10 |
+
|
11 |
+
class ToChartResultConverter(BaseConverter):
|
12 |
+
"""
|
13 |
+
Converts various DensePose predictor outputs to DensePose results.
|
14 |
+
Each DensePose predictor output type has to register its convertion strategy.
|
15 |
+
"""
|
16 |
+
|
17 |
+
registry = {}
|
18 |
+
dst_type = DensePoseChartResult
|
19 |
+
|
20 |
+
@classmethod
|
21 |
+
# pyre-fixme[14]: `convert` overrides method defined in `BaseConverter`
|
22 |
+
# inconsistently.
|
23 |
+
def convert(cls, predictor_outputs: Any, boxes: Boxes, *args, **kwargs) -> DensePoseChartResult:
|
24 |
+
"""
|
25 |
+
Convert DensePose predictor outputs to DensePoseResult using some registered
|
26 |
+
converter. Does recursive lookup for base classes, so there's no need
|
27 |
+
for explicit registration for derived classes.
|
28 |
+
|
29 |
+
Args:
|
30 |
+
densepose_predictor_outputs: DensePose predictor output to be
|
31 |
+
converted to BitMasks
|
32 |
+
boxes (Boxes): bounding boxes that correspond to the DensePose
|
33 |
+
predictor outputs
|
34 |
+
Return:
|
35 |
+
An instance of DensePoseResult. If no suitable converter was found, raises KeyError
|
36 |
+
"""
|
37 |
+
return super(ToChartResultConverter, cls).convert(predictor_outputs, boxes, *args, **kwargs)
|
38 |
+
|
39 |
+
|
40 |
+
class ToChartResultConverterWithConfidences(BaseConverter):
|
41 |
+
"""
|
42 |
+
Converts various DensePose predictor outputs to DensePose results.
|
43 |
+
Each DensePose predictor output type has to register its convertion strategy.
|
44 |
+
"""
|
45 |
+
|
46 |
+
registry = {}
|
47 |
+
dst_type = DensePoseChartResultWithConfidences
|
48 |
+
|
49 |
+
@classmethod
|
50 |
+
# pyre-fixme[14]: `convert` overrides method defined in `BaseConverter`
|
51 |
+
# inconsistently.
|
52 |
+
def convert(
|
53 |
+
cls, predictor_outputs: Any, boxes: Boxes, *args, **kwargs
|
54 |
+
) -> DensePoseChartResultWithConfidences:
|
55 |
+
"""
|
56 |
+
Convert DensePose predictor outputs to DensePoseResult with confidences
|
57 |
+
using some registered converter. Does recursive lookup for base classes,
|
58 |
+
so there's no need for explicit registration for derived classes.
|
59 |
+
|
60 |
+
Args:
|
61 |
+
densepose_predictor_outputs: DensePose predictor output with confidences
|
62 |
+
to be converted to BitMasks
|
63 |
+
boxes (Boxes): bounding boxes that correspond to the DensePose
|
64 |
+
predictor outputs
|
65 |
+
Return:
|
66 |
+
An instance of DensePoseResult. If no suitable converter was found, raises KeyError
|
67 |
+
"""
|
68 |
+
return super(ToChartResultConverterWithConfidences, cls).convert(
|
69 |
+
predictor_outputs, boxes, *args, **kwargs
|
70 |
+
)
|
3rdparty/densepose/converters/to_mask.py
ADDED
@@ -0,0 +1,49 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
|
3 |
+
from typing import Any, Tuple
|
4 |
+
|
5 |
+
from detectron2.structures import BitMasks, Boxes
|
6 |
+
|
7 |
+
from .base import BaseConverter
|
8 |
+
|
9 |
+
ImageSizeType = Tuple[int, int]
|
10 |
+
|
11 |
+
|
12 |
+
class ToMaskConverter(BaseConverter):
|
13 |
+
"""
|
14 |
+
Converts various DensePose predictor outputs to masks
|
15 |
+
in bit mask format (see `BitMasks`). Each DensePose predictor output type
|
16 |
+
has to register its convertion strategy.
|
17 |
+
"""
|
18 |
+
|
19 |
+
registry = {}
|
20 |
+
dst_type = BitMasks
|
21 |
+
|
22 |
+
@classmethod
|
23 |
+
# pyre-fixme[14]: `convert` overrides method defined in `BaseConverter`
|
24 |
+
# inconsistently.
|
25 |
+
def convert(
|
26 |
+
cls,
|
27 |
+
densepose_predictor_outputs: Any,
|
28 |
+
boxes: Boxes,
|
29 |
+
image_size_hw: ImageSizeType,
|
30 |
+
*args,
|
31 |
+
**kwargs
|
32 |
+
) -> BitMasks:
|
33 |
+
"""
|
34 |
+
Convert DensePose predictor outputs to BitMasks using some registered
|
35 |
+
converter. Does recursive lookup for base classes, so there's no need
|
36 |
+
for explicit registration for derived classes.
|
37 |
+
|
38 |
+
Args:
|
39 |
+
densepose_predictor_outputs: DensePose predictor output to be
|
40 |
+
converted to BitMasks
|
41 |
+
boxes (Boxes): bounding boxes that correspond to the DensePose
|
42 |
+
predictor outputs
|
43 |
+
image_size_hw (tuple [int, int]): image height and width
|
44 |
+
Return:
|
45 |
+
An instance of `BitMasks`. If no suitable converter was found, raises KeyError
|
46 |
+
"""
|
47 |
+
return super(ToMaskConverter, cls).convert(
|
48 |
+
densepose_predictor_outputs, boxes, image_size_hw, *args, **kwargs
|
49 |
+
)
|
3rdparty/densepose/data/__init__.py
ADDED
@@ -0,0 +1,25 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
|
3 |
+
from .meshes import builtin
|
4 |
+
from .build import (
|
5 |
+
build_detection_test_loader,
|
6 |
+
build_detection_train_loader,
|
7 |
+
build_combined_loader,
|
8 |
+
build_frame_selector,
|
9 |
+
build_inference_based_loaders,
|
10 |
+
has_inference_based_loaders,
|
11 |
+
BootstrapDatasetFactoryCatalog,
|
12 |
+
)
|
13 |
+
from .combined_loader import CombinedDataLoader
|
14 |
+
from .dataset_mapper import DatasetMapper
|
15 |
+
from .inference_based_loader import InferenceBasedLoader, ScoreBasedFilter
|
16 |
+
from .image_list_dataset import ImageListDataset
|
17 |
+
from .utils import is_relative_local_path, maybe_prepend_base_path
|
18 |
+
|
19 |
+
# ensure the builtin datasets are registered
|
20 |
+
from . import datasets
|
21 |
+
|
22 |
+
# ensure the bootstrap datasets builders are registered
|
23 |
+
from . import build
|
24 |
+
|
25 |
+
__all__ = [k for k in globals().keys() if not k.startswith("_")]
|
3rdparty/densepose/data/build.py
ADDED
@@ -0,0 +1,736 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
|
3 |
+
import itertools
|
4 |
+
import logging
|
5 |
+
import numpy as np
|
6 |
+
from collections import UserDict, defaultdict
|
7 |
+
from dataclasses import dataclass
|
8 |
+
from typing import Any, Callable, Collection, Dict, Iterable, List, Optional, Sequence, Tuple
|
9 |
+
import torch
|
10 |
+
from torch.utils.data.dataset import Dataset
|
11 |
+
|
12 |
+
from detectron2.config import CfgNode
|
13 |
+
from detectron2.data.build import build_detection_test_loader as d2_build_detection_test_loader
|
14 |
+
from detectron2.data.build import build_detection_train_loader as d2_build_detection_train_loader
|
15 |
+
from detectron2.data.build import (
|
16 |
+
load_proposals_into_dataset,
|
17 |
+
print_instances_class_histogram,
|
18 |
+
trivial_batch_collator,
|
19 |
+
worker_init_reset_seed,
|
20 |
+
)
|
21 |
+
from detectron2.data.catalog import DatasetCatalog, Metadata, MetadataCatalog
|
22 |
+
from detectron2.data.samplers import TrainingSampler
|
23 |
+
from detectron2.utils.comm import get_world_size
|
24 |
+
|
25 |
+
from densepose.config import get_bootstrap_dataset_config
|
26 |
+
from densepose.modeling import build_densepose_embedder
|
27 |
+
|
28 |
+
from .combined_loader import CombinedDataLoader, Loader
|
29 |
+
from .dataset_mapper import DatasetMapper
|
30 |
+
from .datasets.coco import DENSEPOSE_CSE_KEYS_WITHOUT_MASK, DENSEPOSE_IUV_KEYS_WITHOUT_MASK
|
31 |
+
from .datasets.dataset_type import DatasetType
|
32 |
+
from .inference_based_loader import InferenceBasedLoader, ScoreBasedFilter
|
33 |
+
from .samplers import (
|
34 |
+
DensePoseConfidenceBasedSampler,
|
35 |
+
DensePoseCSEConfidenceBasedSampler,
|
36 |
+
DensePoseCSEUniformSampler,
|
37 |
+
DensePoseUniformSampler,
|
38 |
+
MaskFromDensePoseSampler,
|
39 |
+
PredictionToGroundTruthSampler,
|
40 |
+
)
|
41 |
+
from .transform import ImageResizeTransform
|
42 |
+
from .utils import get_category_to_class_mapping, get_class_to_mesh_name_mapping
|
43 |
+
from .video import (
|
44 |
+
FirstKFramesSelector,
|
45 |
+
FrameSelectionStrategy,
|
46 |
+
LastKFramesSelector,
|
47 |
+
RandomKFramesSelector,
|
48 |
+
VideoKeyframeDataset,
|
49 |
+
video_list_from_file,
|
50 |
+
)
|
51 |
+
|
52 |
+
__all__ = ["build_detection_train_loader", "build_detection_test_loader"]
|
53 |
+
|
54 |
+
|
55 |
+
Instance = Dict[str, Any]
|
56 |
+
InstancePredicate = Callable[[Instance], bool]
|
57 |
+
|
58 |
+
|
59 |
+
def _compute_num_images_per_worker(cfg: CfgNode) -> int:
|
60 |
+
num_workers = get_world_size()
|
61 |
+
images_per_batch = cfg.SOLVER.IMS_PER_BATCH
|
62 |
+
assert (
|
63 |
+
images_per_batch % num_workers == 0
|
64 |
+
), "SOLVER.IMS_PER_BATCH ({}) must be divisible by the number of workers ({}).".format(
|
65 |
+
images_per_batch, num_workers
|
66 |
+
)
|
67 |
+
assert (
|
68 |
+
images_per_batch >= num_workers
|
69 |
+
), "SOLVER.IMS_PER_BATCH ({}) must be larger than the number of workers ({}).".format(
|
70 |
+
images_per_batch, num_workers
|
71 |
+
)
|
72 |
+
images_per_worker = images_per_batch // num_workers
|
73 |
+
return images_per_worker
|
74 |
+
|
75 |
+
|
76 |
+
def _map_category_id_to_contiguous_id(dataset_name: str, dataset_dicts: Iterable[Instance]) -> None:
|
77 |
+
meta = MetadataCatalog.get(dataset_name)
|
78 |
+
for dataset_dict in dataset_dicts:
|
79 |
+
for ann in dataset_dict["annotations"]:
|
80 |
+
ann["category_id"] = meta.thing_dataset_id_to_contiguous_id[ann["category_id"]]
|
81 |
+
|
82 |
+
|
83 |
+
@dataclass
|
84 |
+
class _DatasetCategory:
|
85 |
+
"""
|
86 |
+
Class representing category data in a dataset:
|
87 |
+
- id: category ID, as specified in the dataset annotations file
|
88 |
+
- name: category name, as specified in the dataset annotations file
|
89 |
+
- mapped_id: category ID after applying category maps (DATASETS.CATEGORY_MAPS config option)
|
90 |
+
- mapped_name: category name after applying category maps
|
91 |
+
- dataset_name: dataset in which the category is defined
|
92 |
+
|
93 |
+
For example, when training models in a class-agnostic manner, one could take LVIS 1.0
|
94 |
+
dataset and map the animal categories to the same category as human data from COCO:
|
95 |
+
id = 225
|
96 |
+
name = "cat"
|
97 |
+
mapped_id = 1
|
98 |
+
mapped_name = "person"
|
99 |
+
dataset_name = "lvis_v1_animals_dp_train"
|
100 |
+
"""
|
101 |
+
|
102 |
+
id: int
|
103 |
+
name: str
|
104 |
+
mapped_id: int
|
105 |
+
mapped_name: str
|
106 |
+
dataset_name: str
|
107 |
+
|
108 |
+
|
109 |
+
_MergedCategoriesT = Dict[int, List[_DatasetCategory]]
|
110 |
+
|
111 |
+
|
112 |
+
def _add_category_id_to_contiguous_id_maps_to_metadata(
|
113 |
+
merged_categories: _MergedCategoriesT,
|
114 |
+
) -> None:
|
115 |
+
merged_categories_per_dataset = {}
|
116 |
+
for contiguous_cat_id, cat_id in enumerate(sorted(merged_categories.keys())):
|
117 |
+
for cat in merged_categories[cat_id]:
|
118 |
+
if cat.dataset_name not in merged_categories_per_dataset:
|
119 |
+
merged_categories_per_dataset[cat.dataset_name] = defaultdict(list)
|
120 |
+
merged_categories_per_dataset[cat.dataset_name][cat_id].append(
|
121 |
+
(
|
122 |
+
contiguous_cat_id,
|
123 |
+
cat,
|
124 |
+
)
|
125 |
+
)
|
126 |
+
|
127 |
+
logger = logging.getLogger(__name__)
|
128 |
+
for dataset_name, merged_categories in merged_categories_per_dataset.items():
|
129 |
+
meta = MetadataCatalog.get(dataset_name)
|
130 |
+
if not hasattr(meta, "thing_classes"):
|
131 |
+
meta.thing_classes = []
|
132 |
+
meta.thing_dataset_id_to_contiguous_id = {}
|
133 |
+
meta.thing_dataset_id_to_merged_id = {}
|
134 |
+
else:
|
135 |
+
meta.thing_classes.clear()
|
136 |
+
meta.thing_dataset_id_to_contiguous_id.clear()
|
137 |
+
meta.thing_dataset_id_to_merged_id.clear()
|
138 |
+
logger.info(f"Dataset {dataset_name}: category ID to contiguous ID mapping:")
|
139 |
+
for _cat_id, categories in sorted(merged_categories.items()):
|
140 |
+
added_to_thing_classes = False
|
141 |
+
for contiguous_cat_id, cat in categories:
|
142 |
+
if not added_to_thing_classes:
|
143 |
+
meta.thing_classes.append(cat.mapped_name)
|
144 |
+
added_to_thing_classes = True
|
145 |
+
meta.thing_dataset_id_to_contiguous_id[cat.id] = contiguous_cat_id
|
146 |
+
meta.thing_dataset_id_to_merged_id[cat.id] = cat.mapped_id
|
147 |
+
logger.info(f"{cat.id} ({cat.name}) -> {contiguous_cat_id}")
|
148 |
+
|
149 |
+
|
150 |
+
def _maybe_create_general_keep_instance_predicate(cfg: CfgNode) -> Optional[InstancePredicate]:
|
151 |
+
def has_annotations(instance: Instance) -> bool:
|
152 |
+
return "annotations" in instance
|
153 |
+
|
154 |
+
def has_only_crowd_anotations(instance: Instance) -> bool:
|
155 |
+
for ann in instance["annotations"]:
|
156 |
+
if ann.get("is_crowd", 0) == 0:
|
157 |
+
return False
|
158 |
+
return True
|
159 |
+
|
160 |
+
def general_keep_instance_predicate(instance: Instance) -> bool:
|
161 |
+
return has_annotations(instance) and not has_only_crowd_anotations(instance)
|
162 |
+
|
163 |
+
if not cfg.DATALOADER.FILTER_EMPTY_ANNOTATIONS:
|
164 |
+
return None
|
165 |
+
return general_keep_instance_predicate
|
166 |
+
|
167 |
+
|
168 |
+
def _maybe_create_keypoints_keep_instance_predicate(cfg: CfgNode) -> Optional[InstancePredicate]:
|
169 |
+
|
170 |
+
min_num_keypoints = cfg.MODEL.ROI_KEYPOINT_HEAD.MIN_KEYPOINTS_PER_IMAGE
|
171 |
+
|
172 |
+
def has_sufficient_num_keypoints(instance: Instance) -> bool:
|
173 |
+
num_kpts = sum(
|
174 |
+
(np.array(ann["keypoints"][2::3]) > 0).sum()
|
175 |
+
for ann in instance["annotations"]
|
176 |
+
if "keypoints" in ann
|
177 |
+
)
|
178 |
+
return num_kpts >= min_num_keypoints
|
179 |
+
|
180 |
+
if cfg.MODEL.KEYPOINT_ON and (min_num_keypoints > 0):
|
181 |
+
return has_sufficient_num_keypoints
|
182 |
+
return None
|
183 |
+
|
184 |
+
|
185 |
+
def _maybe_create_mask_keep_instance_predicate(cfg: CfgNode) -> Optional[InstancePredicate]:
|
186 |
+
if not cfg.MODEL.MASK_ON:
|
187 |
+
return None
|
188 |
+
|
189 |
+
def has_mask_annotations(instance: Instance) -> bool:
|
190 |
+
return any("segmentation" in ann for ann in instance["annotations"])
|
191 |
+
|
192 |
+
return has_mask_annotations
|
193 |
+
|
194 |
+
|
195 |
+
def _maybe_create_densepose_keep_instance_predicate(cfg: CfgNode) -> Optional[InstancePredicate]:
|
196 |
+
if not cfg.MODEL.DENSEPOSE_ON:
|
197 |
+
return None
|
198 |
+
|
199 |
+
use_masks = cfg.MODEL.ROI_DENSEPOSE_HEAD.COARSE_SEGM_TRAINED_BY_MASKS
|
200 |
+
|
201 |
+
def has_densepose_annotations(instance: Instance) -> bool:
|
202 |
+
for ann in instance["annotations"]:
|
203 |
+
if all(key in ann for key in DENSEPOSE_IUV_KEYS_WITHOUT_MASK) or all(
|
204 |
+
key in ann for key in DENSEPOSE_CSE_KEYS_WITHOUT_MASK
|
205 |
+
):
|
206 |
+
return True
|
207 |
+
if use_masks and "segmentation" in ann:
|
208 |
+
return True
|
209 |
+
return False
|
210 |
+
|
211 |
+
return has_densepose_annotations
|
212 |
+
|
213 |
+
|
214 |
+
def _maybe_create_specific_keep_instance_predicate(cfg: CfgNode) -> Optional[InstancePredicate]:
|
215 |
+
specific_predicate_creators = [
|
216 |
+
_maybe_create_keypoints_keep_instance_predicate,
|
217 |
+
_maybe_create_mask_keep_instance_predicate,
|
218 |
+
_maybe_create_densepose_keep_instance_predicate,
|
219 |
+
]
|
220 |
+
predicates = [creator(cfg) for creator in specific_predicate_creators]
|
221 |
+
predicates = [p for p in predicates if p is not None]
|
222 |
+
if not predicates:
|
223 |
+
return None
|
224 |
+
|
225 |
+
def combined_predicate(instance: Instance) -> bool:
|
226 |
+
return any(p(instance) for p in predicates)
|
227 |
+
|
228 |
+
return combined_predicate
|
229 |
+
|
230 |
+
|
231 |
+
def _get_train_keep_instance_predicate(cfg: CfgNode):
|
232 |
+
general_keep_predicate = _maybe_create_general_keep_instance_predicate(cfg)
|
233 |
+
combined_specific_keep_predicate = _maybe_create_specific_keep_instance_predicate(cfg)
|
234 |
+
|
235 |
+
def combined_general_specific_keep_predicate(instance: Instance) -> bool:
|
236 |
+
return general_keep_predicate(instance) and combined_specific_keep_predicate(instance)
|
237 |
+
|
238 |
+
if (general_keep_predicate is None) and (combined_specific_keep_predicate is None):
|
239 |
+
return None
|
240 |
+
if general_keep_predicate is None:
|
241 |
+
return combined_specific_keep_predicate
|
242 |
+
if combined_specific_keep_predicate is None:
|
243 |
+
return general_keep_predicate
|
244 |
+
return combined_general_specific_keep_predicate
|
245 |
+
|
246 |
+
|
247 |
+
def _get_test_keep_instance_predicate(cfg: CfgNode):
|
248 |
+
general_keep_predicate = _maybe_create_general_keep_instance_predicate(cfg)
|
249 |
+
return general_keep_predicate
|
250 |
+
|
251 |
+
|
252 |
+
def _maybe_filter_and_map_categories(
|
253 |
+
dataset_name: str, dataset_dicts: List[Instance]
|
254 |
+
) -> List[Instance]:
|
255 |
+
meta = MetadataCatalog.get(dataset_name)
|
256 |
+
category_id_map = meta.thing_dataset_id_to_contiguous_id
|
257 |
+
filtered_dataset_dicts = []
|
258 |
+
for dataset_dict in dataset_dicts:
|
259 |
+
anns = []
|
260 |
+
for ann in dataset_dict["annotations"]:
|
261 |
+
cat_id = ann["category_id"]
|
262 |
+
if cat_id not in category_id_map:
|
263 |
+
continue
|
264 |
+
ann["category_id"] = category_id_map[cat_id]
|
265 |
+
anns.append(ann)
|
266 |
+
dataset_dict["annotations"] = anns
|
267 |
+
filtered_dataset_dicts.append(dataset_dict)
|
268 |
+
return filtered_dataset_dicts
|
269 |
+
|
270 |
+
|
271 |
+
def _add_category_whitelists_to_metadata(cfg: CfgNode) -> None:
|
272 |
+
for dataset_name, whitelisted_cat_ids in cfg.DATASETS.WHITELISTED_CATEGORIES.items():
|
273 |
+
meta = MetadataCatalog.get(dataset_name)
|
274 |
+
meta.whitelisted_categories = whitelisted_cat_ids
|
275 |
+
logger = logging.getLogger(__name__)
|
276 |
+
logger.info(
|
277 |
+
"Whitelisted categories for dataset {}: {}".format(
|
278 |
+
dataset_name, meta.whitelisted_categories
|
279 |
+
)
|
280 |
+
)
|
281 |
+
|
282 |
+
|
283 |
+
def _add_category_maps_to_metadata(cfg: CfgNode) -> None:
|
284 |
+
for dataset_name, category_map in cfg.DATASETS.CATEGORY_MAPS.items():
|
285 |
+
category_map = {
|
286 |
+
int(cat_id_src): int(cat_id_dst) for cat_id_src, cat_id_dst in category_map.items()
|
287 |
+
}
|
288 |
+
meta = MetadataCatalog.get(dataset_name)
|
289 |
+
meta.category_map = category_map
|
290 |
+
logger = logging.getLogger(__name__)
|
291 |
+
logger.info("Category maps for dataset {}: {}".format(dataset_name, meta.category_map))
|
292 |
+
|
293 |
+
|
294 |
+
def _add_category_info_to_bootstrapping_metadata(dataset_name: str, dataset_cfg: CfgNode) -> None:
|
295 |
+
meta = MetadataCatalog.get(dataset_name)
|
296 |
+
meta.category_to_class_mapping = get_category_to_class_mapping(dataset_cfg)
|
297 |
+
meta.categories = dataset_cfg.CATEGORIES
|
298 |
+
meta.max_count_per_category = dataset_cfg.MAX_COUNT_PER_CATEGORY
|
299 |
+
logger = logging.getLogger(__name__)
|
300 |
+
logger.info(
|
301 |
+
"Category to class mapping for dataset {}: {}".format(
|
302 |
+
dataset_name, meta.category_to_class_mapping
|
303 |
+
)
|
304 |
+
)
|
305 |
+
|
306 |
+
|
307 |
+
def _maybe_add_class_to_mesh_name_map_to_metadata(dataset_names: List[str], cfg: CfgNode) -> None:
|
308 |
+
for dataset_name in dataset_names:
|
309 |
+
meta = MetadataCatalog.get(dataset_name)
|
310 |
+
if not hasattr(meta, "class_to_mesh_name"):
|
311 |
+
meta.class_to_mesh_name = get_class_to_mesh_name_mapping(cfg)
|
312 |
+
|
313 |
+
|
314 |
+
def _merge_categories(dataset_names: Collection[str]) -> _MergedCategoriesT:
|
315 |
+
merged_categories = defaultdict(list)
|
316 |
+
category_names = {}
|
317 |
+
for dataset_name in dataset_names:
|
318 |
+
meta = MetadataCatalog.get(dataset_name)
|
319 |
+
whitelisted_categories = meta.get("whitelisted_categories")
|
320 |
+
category_map = meta.get("category_map", {})
|
321 |
+
cat_ids = (
|
322 |
+
whitelisted_categories if whitelisted_categories is not None else meta.categories.keys()
|
323 |
+
)
|
324 |
+
for cat_id in cat_ids:
|
325 |
+
cat_name = meta.categories[cat_id]
|
326 |
+
cat_id_mapped = category_map.get(cat_id, cat_id)
|
327 |
+
if cat_id_mapped == cat_id or cat_id_mapped in cat_ids:
|
328 |
+
category_names[cat_id] = cat_name
|
329 |
+
else:
|
330 |
+
category_names[cat_id] = str(cat_id_mapped)
|
331 |
+
# assign temporary mapped category name, this name can be changed
|
332 |
+
# during the second pass, since mapped ID can correspond to a category
|
333 |
+
# from a different dataset
|
334 |
+
cat_name_mapped = meta.categories[cat_id_mapped]
|
335 |
+
merged_categories[cat_id_mapped].append(
|
336 |
+
_DatasetCategory(
|
337 |
+
id=cat_id,
|
338 |
+
name=cat_name,
|
339 |
+
mapped_id=cat_id_mapped,
|
340 |
+
mapped_name=cat_name_mapped,
|
341 |
+
dataset_name=dataset_name,
|
342 |
+
)
|
343 |
+
)
|
344 |
+
# second pass to assign proper mapped category names
|
345 |
+
for cat_id, categories in merged_categories.items():
|
346 |
+
for cat in categories:
|
347 |
+
if cat_id in category_names and cat.mapped_name != category_names[cat_id]:
|
348 |
+
cat.mapped_name = category_names[cat_id]
|
349 |
+
|
350 |
+
return merged_categories
|
351 |
+
|
352 |
+
|
353 |
+
def _warn_if_merged_different_categories(merged_categories: _MergedCategoriesT) -> None:
|
354 |
+
logger = logging.getLogger(__name__)
|
355 |
+
for cat_id in merged_categories:
|
356 |
+
merged_categories_i = merged_categories[cat_id]
|
357 |
+
first_cat_name = merged_categories_i[0].name
|
358 |
+
if len(merged_categories_i) > 1 and not all(
|
359 |
+
cat.name == first_cat_name for cat in merged_categories_i[1:]
|
360 |
+
):
|
361 |
+
cat_summary_str = ", ".join(
|
362 |
+
[f"{cat.id} ({cat.name}) from {cat.dataset_name}" for cat in merged_categories_i]
|
363 |
+
)
|
364 |
+
logger.warning(
|
365 |
+
f"Merged category {cat_id} corresponds to the following categories: "
|
366 |
+
f"{cat_summary_str}"
|
367 |
+
)
|
368 |
+
|
369 |
+
|
370 |
+
def combine_detection_dataset_dicts(
|
371 |
+
dataset_names: Collection[str],
|
372 |
+
keep_instance_predicate: Optional[InstancePredicate] = None,
|
373 |
+
proposal_files: Optional[Collection[str]] = None,
|
374 |
+
) -> List[Instance]:
|
375 |
+
"""
|
376 |
+
Load and prepare dataset dicts for training / testing
|
377 |
+
|
378 |
+
Args:
|
379 |
+
dataset_names (Collection[str]): a list of dataset names
|
380 |
+
keep_instance_predicate (Callable: Dict[str, Any] -> bool): predicate
|
381 |
+
applied to instance dicts which defines whether to keep the instance
|
382 |
+
proposal_files (Collection[str]): if given, a list of object proposal files
|
383 |
+
that match each dataset in `dataset_names`.
|
384 |
+
"""
|
385 |
+
assert len(dataset_names)
|
386 |
+
if proposal_files is None:
|
387 |
+
proposal_files = [None] * len(dataset_names)
|
388 |
+
assert len(dataset_names) == len(proposal_files)
|
389 |
+
# load datasets and metadata
|
390 |
+
dataset_name_to_dicts = {}
|
391 |
+
for dataset_name in dataset_names:
|
392 |
+
dataset_name_to_dicts[dataset_name] = DatasetCatalog.get(dataset_name)
|
393 |
+
assert len(dataset_name_to_dicts), f"Dataset '{dataset_name}' is empty!"
|
394 |
+
# merge categories, requires category metadata to be loaded
|
395 |
+
# cat_id -> [(orig_cat_id, cat_name, dataset_name)]
|
396 |
+
merged_categories = _merge_categories(dataset_names)
|
397 |
+
_warn_if_merged_different_categories(merged_categories)
|
398 |
+
merged_category_names = [
|
399 |
+
merged_categories[cat_id][0].mapped_name for cat_id in sorted(merged_categories)
|
400 |
+
]
|
401 |
+
# map to contiguous category IDs
|
402 |
+
_add_category_id_to_contiguous_id_maps_to_metadata(merged_categories)
|
403 |
+
# load annotations and dataset metadata
|
404 |
+
for dataset_name, proposal_file in zip(dataset_names, proposal_files):
|
405 |
+
dataset_dicts = dataset_name_to_dicts[dataset_name]
|
406 |
+
assert len(dataset_dicts), f"Dataset '{dataset_name}' is empty!"
|
407 |
+
if proposal_file is not None:
|
408 |
+
dataset_dicts = load_proposals_into_dataset(dataset_dicts, proposal_file)
|
409 |
+
dataset_dicts = _maybe_filter_and_map_categories(dataset_name, dataset_dicts)
|
410 |
+
print_instances_class_histogram(dataset_dicts, merged_category_names)
|
411 |
+
dataset_name_to_dicts[dataset_name] = dataset_dicts
|
412 |
+
|
413 |
+
if keep_instance_predicate is not None:
|
414 |
+
all_datasets_dicts_plain = [
|
415 |
+
d
|
416 |
+
for d in itertools.chain.from_iterable(dataset_name_to_dicts.values())
|
417 |
+
if keep_instance_predicate(d)
|
418 |
+
]
|
419 |
+
else:
|
420 |
+
all_datasets_dicts_plain = list(
|
421 |
+
itertools.chain.from_iterable(dataset_name_to_dicts.values())
|
422 |
+
)
|
423 |
+
return all_datasets_dicts_plain
|
424 |
+
|
425 |
+
|
426 |
+
def build_detection_train_loader(cfg: CfgNode, mapper=None):
|
427 |
+
"""
|
428 |
+
A data loader is created in a way similar to that of Detectron2.
|
429 |
+
The main differences are:
|
430 |
+
- it allows to combine datasets with different but compatible object category sets
|
431 |
+
|
432 |
+
The data loader is created by the following steps:
|
433 |
+
1. Use the dataset names in config to query :class:`DatasetCatalog`, and obtain a list of dicts.
|
434 |
+
2. Start workers to work on the dicts. Each worker will:
|
435 |
+
* Map each metadata dict into another format to be consumed by the model.
|
436 |
+
* Batch them by simply putting dicts into a list.
|
437 |
+
The batched ``list[mapped_dict]`` is what this dataloader will return.
|
438 |
+
|
439 |
+
Args:
|
440 |
+
cfg (CfgNode): the config
|
441 |
+
mapper (callable): a callable which takes a sample (dict) from dataset and
|
442 |
+
returns the format to be consumed by the model.
|
443 |
+
By default it will be `DatasetMapper(cfg, True)`.
|
444 |
+
|
445 |
+
Returns:
|
446 |
+
an infinite iterator of training data
|
447 |
+
"""
|
448 |
+
|
449 |
+
_add_category_whitelists_to_metadata(cfg)
|
450 |
+
_add_category_maps_to_metadata(cfg)
|
451 |
+
_maybe_add_class_to_mesh_name_map_to_metadata(cfg.DATASETS.TRAIN, cfg)
|
452 |
+
dataset_dicts = combine_detection_dataset_dicts(
|
453 |
+
cfg.DATASETS.TRAIN,
|
454 |
+
keep_instance_predicate=_get_train_keep_instance_predicate(cfg),
|
455 |
+
proposal_files=cfg.DATASETS.PROPOSAL_FILES_TRAIN if cfg.MODEL.LOAD_PROPOSALS else None,
|
456 |
+
)
|
457 |
+
if mapper is None:
|
458 |
+
mapper = DatasetMapper(cfg, True)
|
459 |
+
return d2_build_detection_train_loader(cfg, dataset=dataset_dicts, mapper=mapper)
|
460 |
+
|
461 |
+
|
462 |
+
def build_detection_test_loader(cfg, dataset_name, mapper=None):
|
463 |
+
"""
|
464 |
+
Similar to `build_detection_train_loader`.
|
465 |
+
But this function uses the given `dataset_name` argument (instead of the names in cfg),
|
466 |
+
and uses batch size 1.
|
467 |
+
|
468 |
+
Args:
|
469 |
+
cfg: a detectron2 CfgNode
|
470 |
+
dataset_name (str): a name of the dataset that's available in the DatasetCatalog
|
471 |
+
mapper (callable): a callable which takes a sample (dict) from dataset
|
472 |
+
and returns the format to be consumed by the model.
|
473 |
+
By default it will be `DatasetMapper(cfg, False)`.
|
474 |
+
|
475 |
+
Returns:
|
476 |
+
DataLoader: a torch DataLoader, that loads the given detection
|
477 |
+
dataset, with test-time transformation and batching.
|
478 |
+
"""
|
479 |
+
_add_category_whitelists_to_metadata(cfg)
|
480 |
+
_add_category_maps_to_metadata(cfg)
|
481 |
+
_maybe_add_class_to_mesh_name_map_to_metadata([dataset_name], cfg)
|
482 |
+
dataset_dicts = combine_detection_dataset_dicts(
|
483 |
+
[dataset_name],
|
484 |
+
keep_instance_predicate=_get_test_keep_instance_predicate(cfg),
|
485 |
+
proposal_files=[
|
486 |
+
cfg.DATASETS.PROPOSAL_FILES_TEST[list(cfg.DATASETS.TEST).index(dataset_name)]
|
487 |
+
]
|
488 |
+
if cfg.MODEL.LOAD_PROPOSALS
|
489 |
+
else None,
|
490 |
+
)
|
491 |
+
sampler = None
|
492 |
+
if not cfg.DENSEPOSE_EVALUATION.DISTRIBUTED_INFERENCE:
|
493 |
+
sampler = torch.utils.data.SequentialSampler(dataset_dicts)
|
494 |
+
if mapper is None:
|
495 |
+
mapper = DatasetMapper(cfg, False)
|
496 |
+
return d2_build_detection_test_loader(
|
497 |
+
dataset_dicts, mapper=mapper, num_workers=cfg.DATALOADER.NUM_WORKERS, sampler=sampler
|
498 |
+
)
|
499 |
+
|
500 |
+
|
501 |
+
def build_frame_selector(cfg: CfgNode):
|
502 |
+
strategy = FrameSelectionStrategy(cfg.STRATEGY)
|
503 |
+
if strategy == FrameSelectionStrategy.RANDOM_K:
|
504 |
+
frame_selector = RandomKFramesSelector(cfg.NUM_IMAGES)
|
505 |
+
elif strategy == FrameSelectionStrategy.FIRST_K:
|
506 |
+
frame_selector = FirstKFramesSelector(cfg.NUM_IMAGES)
|
507 |
+
elif strategy == FrameSelectionStrategy.LAST_K:
|
508 |
+
frame_selector = LastKFramesSelector(cfg.NUM_IMAGES)
|
509 |
+
elif strategy == FrameSelectionStrategy.ALL:
|
510 |
+
frame_selector = None
|
511 |
+
# pyre-fixme[61]: `frame_selector` may not be initialized here.
|
512 |
+
return frame_selector
|
513 |
+
|
514 |
+
|
515 |
+
def build_transform(cfg: CfgNode, data_type: str):
|
516 |
+
if cfg.TYPE == "resize":
|
517 |
+
if data_type == "image":
|
518 |
+
return ImageResizeTransform(cfg.MIN_SIZE, cfg.MAX_SIZE)
|
519 |
+
raise ValueError(f"Unknown transform {cfg.TYPE} for data type {data_type}")
|
520 |
+
|
521 |
+
|
522 |
+
def build_combined_loader(cfg: CfgNode, loaders: Collection[Loader], ratios: Sequence[float]):
|
523 |
+
images_per_worker = _compute_num_images_per_worker(cfg)
|
524 |
+
return CombinedDataLoader(loaders, images_per_worker, ratios)
|
525 |
+
|
526 |
+
|
527 |
+
def build_bootstrap_dataset(dataset_name: str, cfg: CfgNode) -> Sequence[torch.Tensor]:
|
528 |
+
"""
|
529 |
+
Build dataset that provides data to bootstrap on
|
530 |
+
|
531 |
+
Args:
|
532 |
+
dataset_name (str): Name of the dataset, needs to have associated metadata
|
533 |
+
to load the data
|
534 |
+
cfg (CfgNode): bootstrapping config
|
535 |
+
Returns:
|
536 |
+
Sequence[Tensor] - dataset that provides image batches, Tensors of size
|
537 |
+
[N, C, H, W] of type float32
|
538 |
+
"""
|
539 |
+
logger = logging.getLogger(__name__)
|
540 |
+
_add_category_info_to_bootstrapping_metadata(dataset_name, cfg)
|
541 |
+
meta = MetadataCatalog.get(dataset_name)
|
542 |
+
factory = BootstrapDatasetFactoryCatalog.get(meta.dataset_type)
|
543 |
+
dataset = None
|
544 |
+
if factory is not None:
|
545 |
+
dataset = factory(meta, cfg)
|
546 |
+
if dataset is None:
|
547 |
+
logger.warning(f"Failed to create dataset {dataset_name} of type {meta.dataset_type}")
|
548 |
+
return dataset
|
549 |
+
|
550 |
+
|
551 |
+
def build_data_sampler(cfg: CfgNode, sampler_cfg: CfgNode, embedder: Optional[torch.nn.Module]):
|
552 |
+
if sampler_cfg.TYPE == "densepose_uniform":
|
553 |
+
data_sampler = PredictionToGroundTruthSampler()
|
554 |
+
# transform densepose pred -> gt
|
555 |
+
data_sampler.register_sampler(
|
556 |
+
"pred_densepose",
|
557 |
+
"gt_densepose",
|
558 |
+
DensePoseUniformSampler(count_per_class=sampler_cfg.COUNT_PER_CLASS),
|
559 |
+
)
|
560 |
+
data_sampler.register_sampler("pred_densepose", "gt_masks", MaskFromDensePoseSampler())
|
561 |
+
return data_sampler
|
562 |
+
elif sampler_cfg.TYPE == "densepose_UV_confidence":
|
563 |
+
data_sampler = PredictionToGroundTruthSampler()
|
564 |
+
# transform densepose pred -> gt
|
565 |
+
data_sampler.register_sampler(
|
566 |
+
"pred_densepose",
|
567 |
+
"gt_densepose",
|
568 |
+
DensePoseConfidenceBasedSampler(
|
569 |
+
confidence_channel="sigma_2",
|
570 |
+
count_per_class=sampler_cfg.COUNT_PER_CLASS,
|
571 |
+
search_proportion=0.5,
|
572 |
+
),
|
573 |
+
)
|
574 |
+
data_sampler.register_sampler("pred_densepose", "gt_masks", MaskFromDensePoseSampler())
|
575 |
+
return data_sampler
|
576 |
+
elif sampler_cfg.TYPE == "densepose_fine_segm_confidence":
|
577 |
+
data_sampler = PredictionToGroundTruthSampler()
|
578 |
+
# transform densepose pred -> gt
|
579 |
+
data_sampler.register_sampler(
|
580 |
+
"pred_densepose",
|
581 |
+
"gt_densepose",
|
582 |
+
DensePoseConfidenceBasedSampler(
|
583 |
+
confidence_channel="fine_segm_confidence",
|
584 |
+
count_per_class=sampler_cfg.COUNT_PER_CLASS,
|
585 |
+
search_proportion=0.5,
|
586 |
+
),
|
587 |
+
)
|
588 |
+
data_sampler.register_sampler("pred_densepose", "gt_masks", MaskFromDensePoseSampler())
|
589 |
+
return data_sampler
|
590 |
+
elif sampler_cfg.TYPE == "densepose_coarse_segm_confidence":
|
591 |
+
data_sampler = PredictionToGroundTruthSampler()
|
592 |
+
# transform densepose pred -> gt
|
593 |
+
data_sampler.register_sampler(
|
594 |
+
"pred_densepose",
|
595 |
+
"gt_densepose",
|
596 |
+
DensePoseConfidenceBasedSampler(
|
597 |
+
confidence_channel="coarse_segm_confidence",
|
598 |
+
count_per_class=sampler_cfg.COUNT_PER_CLASS,
|
599 |
+
search_proportion=0.5,
|
600 |
+
),
|
601 |
+
)
|
602 |
+
data_sampler.register_sampler("pred_densepose", "gt_masks", MaskFromDensePoseSampler())
|
603 |
+
return data_sampler
|
604 |
+
elif sampler_cfg.TYPE == "densepose_cse_uniform":
|
605 |
+
assert embedder is not None
|
606 |
+
data_sampler = PredictionToGroundTruthSampler()
|
607 |
+
# transform densepose pred -> gt
|
608 |
+
data_sampler.register_sampler(
|
609 |
+
"pred_densepose",
|
610 |
+
"gt_densepose",
|
611 |
+
DensePoseCSEUniformSampler(
|
612 |
+
cfg=cfg,
|
613 |
+
use_gt_categories=sampler_cfg.USE_GROUND_TRUTH_CATEGORIES,
|
614 |
+
embedder=embedder,
|
615 |
+
count_per_class=sampler_cfg.COUNT_PER_CLASS,
|
616 |
+
),
|
617 |
+
)
|
618 |
+
data_sampler.register_sampler("pred_densepose", "gt_masks", MaskFromDensePoseSampler())
|
619 |
+
return data_sampler
|
620 |
+
elif sampler_cfg.TYPE == "densepose_cse_coarse_segm_confidence":
|
621 |
+
assert embedder is not None
|
622 |
+
data_sampler = PredictionToGroundTruthSampler()
|
623 |
+
# transform densepose pred -> gt
|
624 |
+
data_sampler.register_sampler(
|
625 |
+
"pred_densepose",
|
626 |
+
"gt_densepose",
|
627 |
+
DensePoseCSEConfidenceBasedSampler(
|
628 |
+
cfg=cfg,
|
629 |
+
use_gt_categories=sampler_cfg.USE_GROUND_TRUTH_CATEGORIES,
|
630 |
+
embedder=embedder,
|
631 |
+
confidence_channel="coarse_segm_confidence",
|
632 |
+
count_per_class=sampler_cfg.COUNT_PER_CLASS,
|
633 |
+
search_proportion=0.5,
|
634 |
+
),
|
635 |
+
)
|
636 |
+
data_sampler.register_sampler("pred_densepose", "gt_masks", MaskFromDensePoseSampler())
|
637 |
+
return data_sampler
|
638 |
+
|
639 |
+
raise ValueError(f"Unknown data sampler type {sampler_cfg.TYPE}")
|
640 |
+
|
641 |
+
|
642 |
+
def build_data_filter(cfg: CfgNode):
|
643 |
+
if cfg.TYPE == "detection_score":
|
644 |
+
min_score = cfg.MIN_VALUE
|
645 |
+
return ScoreBasedFilter(min_score=min_score)
|
646 |
+
raise ValueError(f"Unknown data filter type {cfg.TYPE}")
|
647 |
+
|
648 |
+
|
649 |
+
def build_inference_based_loader(
|
650 |
+
cfg: CfgNode,
|
651 |
+
dataset_cfg: CfgNode,
|
652 |
+
model: torch.nn.Module,
|
653 |
+
embedder: Optional[torch.nn.Module] = None,
|
654 |
+
) -> InferenceBasedLoader:
|
655 |
+
"""
|
656 |
+
Constructs data loader based on inference results of a model.
|
657 |
+
"""
|
658 |
+
dataset = build_bootstrap_dataset(dataset_cfg.DATASET, dataset_cfg.IMAGE_LOADER)
|
659 |
+
meta = MetadataCatalog.get(dataset_cfg.DATASET)
|
660 |
+
training_sampler = TrainingSampler(len(dataset))
|
661 |
+
data_loader = torch.utils.data.DataLoader(
|
662 |
+
dataset, # pyre-ignore[6]
|
663 |
+
batch_size=dataset_cfg.IMAGE_LOADER.BATCH_SIZE,
|
664 |
+
sampler=training_sampler,
|
665 |
+
num_workers=dataset_cfg.IMAGE_LOADER.NUM_WORKERS,
|
666 |
+
collate_fn=trivial_batch_collator,
|
667 |
+
worker_init_fn=worker_init_reset_seed,
|
668 |
+
)
|
669 |
+
return InferenceBasedLoader(
|
670 |
+
model,
|
671 |
+
data_loader=data_loader,
|
672 |
+
data_sampler=build_data_sampler(cfg, dataset_cfg.DATA_SAMPLER, embedder),
|
673 |
+
data_filter=build_data_filter(dataset_cfg.FILTER),
|
674 |
+
shuffle=True,
|
675 |
+
batch_size=dataset_cfg.INFERENCE.OUTPUT_BATCH_SIZE,
|
676 |
+
inference_batch_size=dataset_cfg.INFERENCE.INPUT_BATCH_SIZE,
|
677 |
+
category_to_class_mapping=meta.category_to_class_mapping,
|
678 |
+
)
|
679 |
+
|
680 |
+
|
681 |
+
def has_inference_based_loaders(cfg: CfgNode) -> bool:
|
682 |
+
"""
|
683 |
+
Returns True, if at least one inferense-based loader must
|
684 |
+
be instantiated for training
|
685 |
+
"""
|
686 |
+
return len(cfg.BOOTSTRAP_DATASETS) > 0
|
687 |
+
|
688 |
+
|
689 |
+
def build_inference_based_loaders(
|
690 |
+
cfg: CfgNode, model: torch.nn.Module
|
691 |
+
) -> Tuple[List[InferenceBasedLoader], List[float]]:
|
692 |
+
loaders = []
|
693 |
+
ratios = []
|
694 |
+
embedder = build_densepose_embedder(cfg).to(device=model.device) # pyre-ignore[16]
|
695 |
+
for dataset_spec in cfg.BOOTSTRAP_DATASETS:
|
696 |
+
dataset_cfg = get_bootstrap_dataset_config().clone()
|
697 |
+
dataset_cfg.merge_from_other_cfg(CfgNode(dataset_spec))
|
698 |
+
loader = build_inference_based_loader(cfg, dataset_cfg, model, embedder)
|
699 |
+
loaders.append(loader)
|
700 |
+
ratios.append(dataset_cfg.RATIO)
|
701 |
+
return loaders, ratios
|
702 |
+
|
703 |
+
|
704 |
+
def build_video_list_dataset(meta: Metadata, cfg: CfgNode):
|
705 |
+
video_list_fpath = meta.video_list_fpath
|
706 |
+
video_base_path = meta.video_base_path
|
707 |
+
category = meta.category
|
708 |
+
if cfg.TYPE == "video_keyframe":
|
709 |
+
frame_selector = build_frame_selector(cfg.SELECT)
|
710 |
+
transform = build_transform(cfg.TRANSFORM, data_type="image")
|
711 |
+
video_list = video_list_from_file(video_list_fpath, video_base_path)
|
712 |
+
keyframe_helper_fpath = getattr(cfg, "KEYFRAME_HELPER", None)
|
713 |
+
return VideoKeyframeDataset(
|
714 |
+
video_list, category, frame_selector, transform, keyframe_helper_fpath
|
715 |
+
)
|
716 |
+
|
717 |
+
|
718 |
+
class _BootstrapDatasetFactoryCatalog(UserDict):
|
719 |
+
"""
|
720 |
+
A global dictionary that stores information about bootstrapped datasets creation functions
|
721 |
+
from metadata and config, for diverse DatasetType
|
722 |
+
"""
|
723 |
+
|
724 |
+
def register(self, dataset_type: DatasetType, factory: Callable[[Metadata, CfgNode], Dataset]):
|
725 |
+
"""
|
726 |
+
Args:
|
727 |
+
dataset_type (DatasetType): a DatasetType e.g. DatasetType.VIDEO_LIST
|
728 |
+
factory (Callable[Metadata, CfgNode]): a callable which takes Metadata and cfg
|
729 |
+
arguments and returns a dataset object.
|
730 |
+
"""
|
731 |
+
assert dataset_type not in self, "Dataset '{}' is already registered!".format(dataset_type)
|
732 |
+
self[dataset_type] = factory
|
733 |
+
|
734 |
+
|
735 |
+
BootstrapDatasetFactoryCatalog = _BootstrapDatasetFactoryCatalog()
|
736 |
+
BootstrapDatasetFactoryCatalog.register(DatasetType.VIDEO_LIST, build_video_list_dataset)
|
3rdparty/densepose/data/combined_loader.py
ADDED
@@ -0,0 +1,44 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
|
3 |
+
import random
|
4 |
+
from collections import deque
|
5 |
+
from typing import Any, Collection, Deque, Iterable, Iterator, List, Sequence
|
6 |
+
|
7 |
+
Loader = Iterable[Any]
|
8 |
+
|
9 |
+
|
10 |
+
def _pooled_next(iterator: Iterator[Any], pool: Deque[Any]):
|
11 |
+
if not pool:
|
12 |
+
pool.extend(next(iterator))
|
13 |
+
return pool.popleft()
|
14 |
+
|
15 |
+
|
16 |
+
class CombinedDataLoader:
|
17 |
+
"""
|
18 |
+
Combines data loaders using the provided sampling ratios
|
19 |
+
"""
|
20 |
+
|
21 |
+
BATCH_COUNT = 100
|
22 |
+
|
23 |
+
def __init__(self, loaders: Collection[Loader], batch_size: int, ratios: Sequence[float]):
|
24 |
+
self.loaders = loaders
|
25 |
+
self.batch_size = batch_size
|
26 |
+
self.ratios = ratios
|
27 |
+
|
28 |
+
def __iter__(self) -> Iterator[List[Any]]:
|
29 |
+
iters = [iter(loader) for loader in self.loaders]
|
30 |
+
indices = []
|
31 |
+
pool = [deque()] * len(iters)
|
32 |
+
# infinite iterator, as in D2
|
33 |
+
while True:
|
34 |
+
if not indices:
|
35 |
+
# just a buffer of indices, its size doesn't matter
|
36 |
+
# as long as it's a multiple of batch_size
|
37 |
+
k = self.batch_size * self.BATCH_COUNT
|
38 |
+
indices = random.choices(range(len(self.loaders)), self.ratios, k=k)
|
39 |
+
try:
|
40 |
+
batch = [_pooled_next(iters[i], pool[i]) for i in indices[: self.batch_size]]
|
41 |
+
except StopIteration:
|
42 |
+
break
|
43 |
+
indices = indices[self.batch_size :]
|
44 |
+
yield batch
|
3rdparty/densepose/data/dataset_mapper.py
ADDED
@@ -0,0 +1,168 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# -*- coding: utf-8 -*-
|
2 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
3 |
+
|
4 |
+
import copy
|
5 |
+
import logging
|
6 |
+
from typing import Any, Dict, List, Tuple
|
7 |
+
import torch
|
8 |
+
|
9 |
+
from detectron2.data import MetadataCatalog
|
10 |
+
from detectron2.data import detection_utils as utils
|
11 |
+
from detectron2.data import transforms as T
|
12 |
+
from detectron2.layers import ROIAlign
|
13 |
+
from detectron2.structures import BoxMode
|
14 |
+
from detectron2.utils.file_io import PathManager
|
15 |
+
|
16 |
+
from densepose.structures import DensePoseDataRelative, DensePoseList, DensePoseTransformData
|
17 |
+
|
18 |
+
|
19 |
+
def build_augmentation(cfg, is_train):
|
20 |
+
logger = logging.getLogger(__name__)
|
21 |
+
result = utils.build_augmentation(cfg, is_train)
|
22 |
+
if is_train:
|
23 |
+
random_rotation = T.RandomRotation(
|
24 |
+
cfg.INPUT.ROTATION_ANGLES, expand=False, sample_style="choice"
|
25 |
+
)
|
26 |
+
result.append(random_rotation)
|
27 |
+
logger.info("DensePose-specific augmentation used in training: " + str(random_rotation))
|
28 |
+
return result
|
29 |
+
|
30 |
+
|
31 |
+
class DatasetMapper:
|
32 |
+
"""
|
33 |
+
A customized version of `detectron2.data.DatasetMapper`
|
34 |
+
"""
|
35 |
+
|
36 |
+
def __init__(self, cfg, is_train=True):
|
37 |
+
self.augmentation = build_augmentation(cfg, is_train)
|
38 |
+
|
39 |
+
# fmt: off
|
40 |
+
self.img_format = cfg.INPUT.FORMAT
|
41 |
+
self.mask_on = (
|
42 |
+
cfg.MODEL.MASK_ON or (
|
43 |
+
cfg.MODEL.DENSEPOSE_ON
|
44 |
+
and cfg.MODEL.ROI_DENSEPOSE_HEAD.COARSE_SEGM_TRAINED_BY_MASKS)
|
45 |
+
)
|
46 |
+
self.keypoint_on = cfg.MODEL.KEYPOINT_ON
|
47 |
+
self.densepose_on = cfg.MODEL.DENSEPOSE_ON
|
48 |
+
assert not cfg.MODEL.LOAD_PROPOSALS, "not supported yet"
|
49 |
+
# fmt: on
|
50 |
+
if self.keypoint_on and is_train:
|
51 |
+
# Flip only makes sense in training
|
52 |
+
self.keypoint_hflip_indices = utils.create_keypoint_hflip_indices(cfg.DATASETS.TRAIN)
|
53 |
+
else:
|
54 |
+
self.keypoint_hflip_indices = None
|
55 |
+
|
56 |
+
if self.densepose_on:
|
57 |
+
densepose_transform_srcs = [
|
58 |
+
MetadataCatalog.get(ds).densepose_transform_src
|
59 |
+
for ds in cfg.DATASETS.TRAIN + cfg.DATASETS.TEST
|
60 |
+
]
|
61 |
+
assert len(densepose_transform_srcs) > 0
|
62 |
+
# TODO: check that DensePose transformation data is the same for
|
63 |
+
# all the datasets. Otherwise one would have to pass DB ID with
|
64 |
+
# each entry to select proper transformation data. For now, since
|
65 |
+
# all DensePose annotated data uses the same data semantics, we
|
66 |
+
# omit this check.
|
67 |
+
densepose_transform_data_fpath = PathManager.get_local_path(densepose_transform_srcs[0])
|
68 |
+
self.densepose_transform_data = DensePoseTransformData.load(
|
69 |
+
densepose_transform_data_fpath
|
70 |
+
)
|
71 |
+
|
72 |
+
self.is_train = is_train
|
73 |
+
|
74 |
+
def __call__(self, dataset_dict):
|
75 |
+
"""
|
76 |
+
Args:
|
77 |
+
dataset_dict (dict): Metadata of one image, in Detectron2 Dataset format.
|
78 |
+
|
79 |
+
Returns:
|
80 |
+
dict: a format that builtin models in detectron2 accept
|
81 |
+
"""
|
82 |
+
dataset_dict = copy.deepcopy(dataset_dict) # it will be modified by code below
|
83 |
+
image = utils.read_image(dataset_dict["file_name"], format=self.img_format)
|
84 |
+
utils.check_image_size(dataset_dict, image)
|
85 |
+
|
86 |
+
image, transforms = T.apply_transform_gens(self.augmentation, image)
|
87 |
+
image_shape = image.shape[:2] # h, w
|
88 |
+
dataset_dict["image"] = torch.as_tensor(image.transpose(2, 0, 1).astype("float32"))
|
89 |
+
|
90 |
+
if not self.is_train:
|
91 |
+
dataset_dict.pop("annotations", None)
|
92 |
+
return dataset_dict
|
93 |
+
|
94 |
+
for anno in dataset_dict["annotations"]:
|
95 |
+
if not self.mask_on:
|
96 |
+
anno.pop("segmentation", None)
|
97 |
+
if not self.keypoint_on:
|
98 |
+
anno.pop("keypoints", None)
|
99 |
+
|
100 |
+
# USER: Implement additional transformations if you have other types of data
|
101 |
+
# USER: Don't call transpose_densepose if you don't need
|
102 |
+
annos = [
|
103 |
+
self._transform_densepose(
|
104 |
+
utils.transform_instance_annotations(
|
105 |
+
obj, transforms, image_shape, keypoint_hflip_indices=self.keypoint_hflip_indices
|
106 |
+
),
|
107 |
+
transforms,
|
108 |
+
)
|
109 |
+
for obj in dataset_dict.pop("annotations")
|
110 |
+
if obj.get("iscrowd", 0) == 0
|
111 |
+
]
|
112 |
+
|
113 |
+
if self.mask_on:
|
114 |
+
self._add_densepose_masks_as_segmentation(annos, image_shape)
|
115 |
+
|
116 |
+
instances = utils.annotations_to_instances(annos, image_shape, mask_format="bitmask")
|
117 |
+
densepose_annotations = [obj.get("densepose") for obj in annos]
|
118 |
+
if densepose_annotations and not all(v is None for v in densepose_annotations):
|
119 |
+
instances.gt_densepose = DensePoseList(
|
120 |
+
densepose_annotations, instances.gt_boxes, image_shape
|
121 |
+
)
|
122 |
+
|
123 |
+
dataset_dict["instances"] = instances[instances.gt_boxes.nonempty()]
|
124 |
+
return dataset_dict
|
125 |
+
|
126 |
+
def _transform_densepose(self, annotation, transforms):
|
127 |
+
if not self.densepose_on:
|
128 |
+
return annotation
|
129 |
+
|
130 |
+
# Handle densepose annotations
|
131 |
+
is_valid, reason_not_valid = DensePoseDataRelative.validate_annotation(annotation)
|
132 |
+
if is_valid:
|
133 |
+
densepose_data = DensePoseDataRelative(annotation, cleanup=True)
|
134 |
+
densepose_data.apply_transform(transforms, self.densepose_transform_data)
|
135 |
+
annotation["densepose"] = densepose_data
|
136 |
+
else:
|
137 |
+
# logger = logging.getLogger(__name__)
|
138 |
+
# logger.debug("Could not load DensePose annotation: {}".format(reason_not_valid))
|
139 |
+
DensePoseDataRelative.cleanup_annotation(annotation)
|
140 |
+
# NOTE: annotations for certain instances may be unavailable.
|
141 |
+
# 'None' is accepted by the DensePostList data structure.
|
142 |
+
annotation["densepose"] = None
|
143 |
+
return annotation
|
144 |
+
|
145 |
+
def _add_densepose_masks_as_segmentation(
|
146 |
+
self, annotations: List[Dict[str, Any]], image_shape_hw: Tuple[int, int]
|
147 |
+
):
|
148 |
+
for obj in annotations:
|
149 |
+
if ("densepose" not in obj) or ("segmentation" in obj):
|
150 |
+
continue
|
151 |
+
# DP segmentation: torch.Tensor [S, S] of float32, S=256
|
152 |
+
segm_dp = torch.zeros_like(obj["densepose"].segm)
|
153 |
+
segm_dp[obj["densepose"].segm > 0] = 1
|
154 |
+
segm_h, segm_w = segm_dp.shape
|
155 |
+
bbox_segm_dp = torch.tensor((0, 0, segm_h - 1, segm_w - 1), dtype=torch.float32)
|
156 |
+
# image bbox
|
157 |
+
x0, y0, x1, y1 = (
|
158 |
+
v.item() for v in BoxMode.convert(obj["bbox"], obj["bbox_mode"], BoxMode.XYXY_ABS)
|
159 |
+
)
|
160 |
+
segm_aligned = (
|
161 |
+
ROIAlign((y1 - y0, x1 - x0), 1.0, 0, aligned=True)
|
162 |
+
.forward(segm_dp.view(1, 1, *segm_dp.shape), bbox_segm_dp)
|
163 |
+
.squeeze()
|
164 |
+
)
|
165 |
+
image_mask = torch.zeros(*image_shape_hw, dtype=torch.float32)
|
166 |
+
image_mask[y0:y1, x0:x1] = segm_aligned
|
167 |
+
# segmentation for BitMask: np.array [H, W] of bool
|
168 |
+
obj["segmentation"] = image_mask >= 0.5
|
3rdparty/densepose/data/datasets/__init__.py
ADDED
@@ -0,0 +1,5 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
|
3 |
+
from . import builtin # ensure the builtin datasets are registered
|
4 |
+
|
5 |
+
__all__ = [k for k in globals().keys() if "builtin" not in k and not k.startswith("_")]
|
3rdparty/densepose/data/datasets/builtin.py
ADDED
@@ -0,0 +1,16 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
from .chimpnsee import register_dataset as register_chimpnsee_dataset
|
3 |
+
from .coco import BASE_DATASETS as BASE_COCO_DATASETS
|
4 |
+
from .coco import DATASETS as COCO_DATASETS
|
5 |
+
from .coco import register_datasets as register_coco_datasets
|
6 |
+
from .lvis import DATASETS as LVIS_DATASETS
|
7 |
+
from .lvis import register_datasets as register_lvis_datasets
|
8 |
+
|
9 |
+
DEFAULT_DATASETS_ROOT = "datasets"
|
10 |
+
|
11 |
+
|
12 |
+
register_coco_datasets(COCO_DATASETS, DEFAULT_DATASETS_ROOT)
|
13 |
+
register_coco_datasets(BASE_COCO_DATASETS, DEFAULT_DATASETS_ROOT)
|
14 |
+
register_lvis_datasets(LVIS_DATASETS, DEFAULT_DATASETS_ROOT)
|
15 |
+
|
16 |
+
register_chimpnsee_dataset(DEFAULT_DATASETS_ROOT) # pyre-ignore[19]
|
3rdparty/densepose/data/datasets/chimpnsee.py
ADDED
@@ -0,0 +1,29 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
|
3 |
+
from typing import Optional
|
4 |
+
|
5 |
+
from detectron2.data import DatasetCatalog, MetadataCatalog
|
6 |
+
|
7 |
+
from ..utils import maybe_prepend_base_path
|
8 |
+
from .dataset_type import DatasetType
|
9 |
+
|
10 |
+
CHIMPNSEE_DATASET_NAME = "chimpnsee"
|
11 |
+
|
12 |
+
|
13 |
+
def register_dataset(datasets_root: Optional[str] = None) -> None:
|
14 |
+
def empty_load_callback():
|
15 |
+
pass
|
16 |
+
|
17 |
+
video_list_fpath = maybe_prepend_base_path(
|
18 |
+
datasets_root,
|
19 |
+
"chimpnsee/cdna.eva.mpg.de/video_list.txt",
|
20 |
+
)
|
21 |
+
video_base_path = maybe_prepend_base_path(datasets_root, "chimpnsee/cdna.eva.mpg.de")
|
22 |
+
|
23 |
+
DatasetCatalog.register(CHIMPNSEE_DATASET_NAME, empty_load_callback)
|
24 |
+
MetadataCatalog.get(CHIMPNSEE_DATASET_NAME).set(
|
25 |
+
dataset_type=DatasetType.VIDEO_LIST,
|
26 |
+
video_list_fpath=video_list_fpath,
|
27 |
+
video_base_path=video_base_path,
|
28 |
+
category="chimpanzee",
|
29 |
+
)
|
3rdparty/densepose/data/datasets/coco.py
ADDED
@@ -0,0 +1,432 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
import contextlib
|
3 |
+
import io
|
4 |
+
import logging
|
5 |
+
import os
|
6 |
+
from collections import defaultdict
|
7 |
+
from dataclasses import dataclass
|
8 |
+
from typing import Any, Dict, Iterable, List, Optional
|
9 |
+
from fvcore.common.timer import Timer
|
10 |
+
|
11 |
+
from detectron2.data import DatasetCatalog, MetadataCatalog
|
12 |
+
from detectron2.structures import BoxMode
|
13 |
+
from detectron2.utils.file_io import PathManager
|
14 |
+
|
15 |
+
from ..utils import maybe_prepend_base_path
|
16 |
+
|
17 |
+
DENSEPOSE_MASK_KEY = "dp_masks"
|
18 |
+
DENSEPOSE_IUV_KEYS_WITHOUT_MASK = ["dp_x", "dp_y", "dp_I", "dp_U", "dp_V"]
|
19 |
+
DENSEPOSE_CSE_KEYS_WITHOUT_MASK = ["dp_x", "dp_y", "dp_vertex", "ref_model"]
|
20 |
+
DENSEPOSE_ALL_POSSIBLE_KEYS = set(
|
21 |
+
DENSEPOSE_IUV_KEYS_WITHOUT_MASK + DENSEPOSE_CSE_KEYS_WITHOUT_MASK + [DENSEPOSE_MASK_KEY]
|
22 |
+
)
|
23 |
+
DENSEPOSE_METADATA_URL_PREFIX = "https://dl.fbaipublicfiles.com/densepose/data/"
|
24 |
+
|
25 |
+
|
26 |
+
@dataclass
|
27 |
+
class CocoDatasetInfo:
|
28 |
+
name: str
|
29 |
+
images_root: str
|
30 |
+
annotations_fpath: str
|
31 |
+
|
32 |
+
|
33 |
+
DATASETS = [
|
34 |
+
CocoDatasetInfo(
|
35 |
+
name="densepose_coco_2014_train",
|
36 |
+
images_root="coco/train2014",
|
37 |
+
annotations_fpath="coco/annotations/densepose_train2014.json",
|
38 |
+
),
|
39 |
+
CocoDatasetInfo(
|
40 |
+
name="densepose_coco_2014_minival",
|
41 |
+
images_root="coco/val2014",
|
42 |
+
annotations_fpath="coco/annotations/densepose_minival2014.json",
|
43 |
+
),
|
44 |
+
CocoDatasetInfo(
|
45 |
+
name="densepose_coco_2014_minival_100",
|
46 |
+
images_root="coco/val2014",
|
47 |
+
annotations_fpath="coco/annotations/densepose_minival2014_100.json",
|
48 |
+
),
|
49 |
+
CocoDatasetInfo(
|
50 |
+
name="densepose_coco_2014_valminusminival",
|
51 |
+
images_root="coco/val2014",
|
52 |
+
annotations_fpath="coco/annotations/densepose_valminusminival2014.json",
|
53 |
+
),
|
54 |
+
CocoDatasetInfo(
|
55 |
+
name="densepose_coco_2014_train_cse",
|
56 |
+
images_root="coco/train2014",
|
57 |
+
annotations_fpath="coco_cse/densepose_train2014_cse.json",
|
58 |
+
),
|
59 |
+
CocoDatasetInfo(
|
60 |
+
name="densepose_coco_2014_minival_cse",
|
61 |
+
images_root="coco/val2014",
|
62 |
+
annotations_fpath="coco_cse/densepose_minival2014_cse.json",
|
63 |
+
),
|
64 |
+
CocoDatasetInfo(
|
65 |
+
name="densepose_coco_2014_minival_100_cse",
|
66 |
+
images_root="coco/val2014",
|
67 |
+
annotations_fpath="coco_cse/densepose_minival2014_100_cse.json",
|
68 |
+
),
|
69 |
+
CocoDatasetInfo(
|
70 |
+
name="densepose_coco_2014_valminusminival_cse",
|
71 |
+
images_root="coco/val2014",
|
72 |
+
annotations_fpath="coco_cse/densepose_valminusminival2014_cse.json",
|
73 |
+
),
|
74 |
+
CocoDatasetInfo(
|
75 |
+
name="densepose_chimps",
|
76 |
+
images_root="densepose_chimps/images",
|
77 |
+
annotations_fpath="densepose_chimps/densepose_chimps_densepose.json",
|
78 |
+
),
|
79 |
+
CocoDatasetInfo(
|
80 |
+
name="densepose_chimps_cse_train",
|
81 |
+
images_root="densepose_chimps/images",
|
82 |
+
annotations_fpath="densepose_chimps/densepose_chimps_cse_train.json",
|
83 |
+
),
|
84 |
+
CocoDatasetInfo(
|
85 |
+
name="densepose_chimps_cse_val",
|
86 |
+
images_root="densepose_chimps/images",
|
87 |
+
annotations_fpath="densepose_chimps/densepose_chimps_cse_val.json",
|
88 |
+
),
|
89 |
+
CocoDatasetInfo(
|
90 |
+
name="posetrack2017_train",
|
91 |
+
images_root="posetrack2017/posetrack_data_2017",
|
92 |
+
annotations_fpath="posetrack2017/densepose_posetrack_train2017.json",
|
93 |
+
),
|
94 |
+
CocoDatasetInfo(
|
95 |
+
name="posetrack2017_val",
|
96 |
+
images_root="posetrack2017/posetrack_data_2017",
|
97 |
+
annotations_fpath="posetrack2017/densepose_posetrack_val2017.json",
|
98 |
+
),
|
99 |
+
CocoDatasetInfo(
|
100 |
+
name="lvis_v05_train",
|
101 |
+
images_root="coco/train2017",
|
102 |
+
annotations_fpath="lvis/lvis_v0.5_plus_dp_train.json",
|
103 |
+
),
|
104 |
+
CocoDatasetInfo(
|
105 |
+
name="lvis_v05_val",
|
106 |
+
images_root="coco/val2017",
|
107 |
+
annotations_fpath="lvis/lvis_v0.5_plus_dp_val.json",
|
108 |
+
),
|
109 |
+
]
|
110 |
+
|
111 |
+
|
112 |
+
BASE_DATASETS = [
|
113 |
+
CocoDatasetInfo(
|
114 |
+
name="base_coco_2017_train",
|
115 |
+
images_root="coco/train2017",
|
116 |
+
annotations_fpath="coco/annotations/instances_train2017.json",
|
117 |
+
),
|
118 |
+
CocoDatasetInfo(
|
119 |
+
name="base_coco_2017_val",
|
120 |
+
images_root="coco/val2017",
|
121 |
+
annotations_fpath="coco/annotations/instances_val2017.json",
|
122 |
+
),
|
123 |
+
CocoDatasetInfo(
|
124 |
+
name="base_coco_2017_val_100",
|
125 |
+
images_root="coco/val2017",
|
126 |
+
annotations_fpath="coco/annotations/instances_val2017_100.json",
|
127 |
+
),
|
128 |
+
]
|
129 |
+
|
130 |
+
|
131 |
+
def get_metadata(base_path: Optional[str]) -> Dict[str, Any]:
|
132 |
+
"""
|
133 |
+
Returns metadata associated with COCO DensePose datasets
|
134 |
+
|
135 |
+
Args:
|
136 |
+
base_path: Optional[str]
|
137 |
+
Base path used to load metadata from
|
138 |
+
|
139 |
+
Returns:
|
140 |
+
Dict[str, Any]
|
141 |
+
Metadata in the form of a dictionary
|
142 |
+
"""
|
143 |
+
meta = {
|
144 |
+
"densepose_transform_src": maybe_prepend_base_path(base_path, "UV_symmetry_transforms.mat"),
|
145 |
+
"densepose_smpl_subdiv": maybe_prepend_base_path(base_path, "SMPL_subdiv.mat"),
|
146 |
+
"densepose_smpl_subdiv_transform": maybe_prepend_base_path(
|
147 |
+
base_path,
|
148 |
+
"SMPL_SUBDIV_TRANSFORM.mat",
|
149 |
+
),
|
150 |
+
}
|
151 |
+
return meta
|
152 |
+
|
153 |
+
|
154 |
+
def _load_coco_annotations(json_file: str):
|
155 |
+
"""
|
156 |
+
Load COCO annotations from a JSON file
|
157 |
+
|
158 |
+
Args:
|
159 |
+
json_file: str
|
160 |
+
Path to the file to load annotations from
|
161 |
+
Returns:
|
162 |
+
Instance of `pycocotools.coco.COCO` that provides access to annotations
|
163 |
+
data
|
164 |
+
"""
|
165 |
+
from pycocotools.coco import COCO
|
166 |
+
|
167 |
+
logger = logging.getLogger(__name__)
|
168 |
+
timer = Timer()
|
169 |
+
with contextlib.redirect_stdout(io.StringIO()):
|
170 |
+
coco_api = COCO(json_file)
|
171 |
+
if timer.seconds() > 1:
|
172 |
+
logger.info("Loading {} takes {:.2f} seconds.".format(json_file, timer.seconds()))
|
173 |
+
return coco_api
|
174 |
+
|
175 |
+
|
176 |
+
def _add_categories_metadata(dataset_name: str, categories: List[Dict[str, Any]]):
|
177 |
+
meta = MetadataCatalog.get(dataset_name)
|
178 |
+
meta.categories = {c["id"]: c["name"] for c in categories}
|
179 |
+
logger = logging.getLogger(__name__)
|
180 |
+
logger.info("Dataset {} categories: {}".format(dataset_name, meta.categories))
|
181 |
+
|
182 |
+
|
183 |
+
def _verify_annotations_have_unique_ids(json_file: str, anns: List[List[Dict[str, Any]]]):
|
184 |
+
if "minival" in json_file:
|
185 |
+
# Skip validation on COCO2014 valminusminival and minival annotations
|
186 |
+
# The ratio of buggy annotations there is tiny and does not affect accuracy
|
187 |
+
# Therefore we explicitly white-list them
|
188 |
+
return
|
189 |
+
ann_ids = [ann["id"] for anns_per_image in anns for ann in anns_per_image]
|
190 |
+
assert len(set(ann_ids)) == len(ann_ids), "Annotation ids in '{}' are not unique!".format(
|
191 |
+
json_file
|
192 |
+
)
|
193 |
+
|
194 |
+
|
195 |
+
def _maybe_add_bbox(obj: Dict[str, Any], ann_dict: Dict[str, Any]):
|
196 |
+
if "bbox" not in ann_dict:
|
197 |
+
return
|
198 |
+
obj["bbox"] = ann_dict["bbox"]
|
199 |
+
obj["bbox_mode"] = BoxMode.XYWH_ABS
|
200 |
+
|
201 |
+
|
202 |
+
def _maybe_add_segm(obj: Dict[str, Any], ann_dict: Dict[str, Any]):
|
203 |
+
if "segmentation" not in ann_dict:
|
204 |
+
return
|
205 |
+
segm = ann_dict["segmentation"]
|
206 |
+
if not isinstance(segm, dict):
|
207 |
+
# filter out invalid polygons (< 3 points)
|
208 |
+
segm = [poly for poly in segm if len(poly) % 2 == 0 and len(poly) >= 6]
|
209 |
+
if len(segm) == 0:
|
210 |
+
return
|
211 |
+
obj["segmentation"] = segm
|
212 |
+
|
213 |
+
|
214 |
+
def _maybe_add_keypoints(obj: Dict[str, Any], ann_dict: Dict[str, Any]):
|
215 |
+
if "keypoints" not in ann_dict:
|
216 |
+
return
|
217 |
+
keypts = ann_dict["keypoints"] # list[int]
|
218 |
+
for idx, v in enumerate(keypts):
|
219 |
+
if idx % 3 != 2:
|
220 |
+
# COCO's segmentation coordinates are floating points in [0, H or W],
|
221 |
+
# but keypoint coordinates are integers in [0, H-1 or W-1]
|
222 |
+
# Therefore we assume the coordinates are "pixel indices" and
|
223 |
+
# add 0.5 to convert to floating point coordinates.
|
224 |
+
keypts[idx] = v + 0.5
|
225 |
+
obj["keypoints"] = keypts
|
226 |
+
|
227 |
+
|
228 |
+
def _maybe_add_densepose(obj: Dict[str, Any], ann_dict: Dict[str, Any]):
|
229 |
+
for key in DENSEPOSE_ALL_POSSIBLE_KEYS:
|
230 |
+
if key in ann_dict:
|
231 |
+
obj[key] = ann_dict[key]
|
232 |
+
|
233 |
+
|
234 |
+
def _combine_images_with_annotations(
|
235 |
+
dataset_name: str,
|
236 |
+
image_root: str,
|
237 |
+
img_datas: Iterable[Dict[str, Any]],
|
238 |
+
ann_datas: Iterable[Iterable[Dict[str, Any]]],
|
239 |
+
):
|
240 |
+
|
241 |
+
ann_keys = ["iscrowd", "category_id"]
|
242 |
+
dataset_dicts = []
|
243 |
+
contains_video_frame_info = False
|
244 |
+
|
245 |
+
for img_dict, ann_dicts in zip(img_datas, ann_datas):
|
246 |
+
record = {}
|
247 |
+
record["file_name"] = os.path.join(image_root, img_dict["file_name"])
|
248 |
+
record["height"] = img_dict["height"]
|
249 |
+
record["width"] = img_dict["width"]
|
250 |
+
record["image_id"] = img_dict["id"]
|
251 |
+
record["dataset"] = dataset_name
|
252 |
+
if "frame_id" in img_dict:
|
253 |
+
record["frame_id"] = img_dict["frame_id"]
|
254 |
+
record["video_id"] = img_dict.get("vid_id", None)
|
255 |
+
contains_video_frame_info = True
|
256 |
+
objs = []
|
257 |
+
for ann_dict in ann_dicts:
|
258 |
+
assert ann_dict["image_id"] == record["image_id"]
|
259 |
+
assert ann_dict.get("ignore", 0) == 0
|
260 |
+
obj = {key: ann_dict[key] for key in ann_keys if key in ann_dict}
|
261 |
+
_maybe_add_bbox(obj, ann_dict)
|
262 |
+
_maybe_add_segm(obj, ann_dict)
|
263 |
+
_maybe_add_keypoints(obj, ann_dict)
|
264 |
+
_maybe_add_densepose(obj, ann_dict)
|
265 |
+
objs.append(obj)
|
266 |
+
record["annotations"] = objs
|
267 |
+
dataset_dicts.append(record)
|
268 |
+
if contains_video_frame_info:
|
269 |
+
create_video_frame_mapping(dataset_name, dataset_dicts)
|
270 |
+
return dataset_dicts
|
271 |
+
|
272 |
+
|
273 |
+
def get_contiguous_id_to_category_id_map(metadata):
|
274 |
+
cat_id_2_cont_id = metadata.thing_dataset_id_to_contiguous_id
|
275 |
+
cont_id_2_cat_id = {}
|
276 |
+
for cat_id, cont_id in cat_id_2_cont_id.items():
|
277 |
+
if cont_id in cont_id_2_cat_id:
|
278 |
+
continue
|
279 |
+
cont_id_2_cat_id[cont_id] = cat_id
|
280 |
+
return cont_id_2_cat_id
|
281 |
+
|
282 |
+
|
283 |
+
def maybe_filter_categories_cocoapi(dataset_name, coco_api):
|
284 |
+
meta = MetadataCatalog.get(dataset_name)
|
285 |
+
cont_id_2_cat_id = get_contiguous_id_to_category_id_map(meta)
|
286 |
+
cat_id_2_cont_id = meta.thing_dataset_id_to_contiguous_id
|
287 |
+
# filter categories
|
288 |
+
cats = []
|
289 |
+
for cat in coco_api.dataset["categories"]:
|
290 |
+
cat_id = cat["id"]
|
291 |
+
if cat_id not in cat_id_2_cont_id:
|
292 |
+
continue
|
293 |
+
cont_id = cat_id_2_cont_id[cat_id]
|
294 |
+
if (cont_id in cont_id_2_cat_id) and (cont_id_2_cat_id[cont_id] == cat_id):
|
295 |
+
cats.append(cat)
|
296 |
+
coco_api.dataset["categories"] = cats
|
297 |
+
# filter annotations, if multiple categories are mapped to a single
|
298 |
+
# contiguous ID, use only one category ID and map all annotations to that category ID
|
299 |
+
anns = []
|
300 |
+
for ann in coco_api.dataset["annotations"]:
|
301 |
+
cat_id = ann["category_id"]
|
302 |
+
if cat_id not in cat_id_2_cont_id:
|
303 |
+
continue
|
304 |
+
cont_id = cat_id_2_cont_id[cat_id]
|
305 |
+
ann["category_id"] = cont_id_2_cat_id[cont_id]
|
306 |
+
anns.append(ann)
|
307 |
+
coco_api.dataset["annotations"] = anns
|
308 |
+
# recreate index
|
309 |
+
coco_api.createIndex()
|
310 |
+
|
311 |
+
|
312 |
+
def maybe_filter_and_map_categories_cocoapi(dataset_name, coco_api):
|
313 |
+
meta = MetadataCatalog.get(dataset_name)
|
314 |
+
category_id_map = meta.thing_dataset_id_to_contiguous_id
|
315 |
+
# map categories
|
316 |
+
cats = []
|
317 |
+
for cat in coco_api.dataset["categories"]:
|
318 |
+
cat_id = cat["id"]
|
319 |
+
if cat_id not in category_id_map:
|
320 |
+
continue
|
321 |
+
cat["id"] = category_id_map[cat_id]
|
322 |
+
cats.append(cat)
|
323 |
+
coco_api.dataset["categories"] = cats
|
324 |
+
# map annotation categories
|
325 |
+
anns = []
|
326 |
+
for ann in coco_api.dataset["annotations"]:
|
327 |
+
cat_id = ann["category_id"]
|
328 |
+
if cat_id not in category_id_map:
|
329 |
+
continue
|
330 |
+
ann["category_id"] = category_id_map[cat_id]
|
331 |
+
anns.append(ann)
|
332 |
+
coco_api.dataset["annotations"] = anns
|
333 |
+
# recreate index
|
334 |
+
coco_api.createIndex()
|
335 |
+
|
336 |
+
|
337 |
+
def create_video_frame_mapping(dataset_name, dataset_dicts):
|
338 |
+
mapping = defaultdict(dict)
|
339 |
+
for d in dataset_dicts:
|
340 |
+
video_id = d.get("video_id")
|
341 |
+
if video_id is None:
|
342 |
+
continue
|
343 |
+
mapping[video_id].update({d["frame_id"]: d["file_name"]})
|
344 |
+
MetadataCatalog.get(dataset_name).set(video_frame_mapping=mapping)
|
345 |
+
|
346 |
+
|
347 |
+
def load_coco_json(annotations_json_file: str, image_root: str, dataset_name: str):
|
348 |
+
"""
|
349 |
+
Loads a JSON file with annotations in COCO instances format.
|
350 |
+
Replaces `detectron2.data.datasets.coco.load_coco_json` to handle metadata
|
351 |
+
in a more flexible way. Postpones category mapping to a later stage to be
|
352 |
+
able to combine several datasets with different (but coherent) sets of
|
353 |
+
categories.
|
354 |
+
|
355 |
+
Args:
|
356 |
+
|
357 |
+
annotations_json_file: str
|
358 |
+
Path to the JSON file with annotations in COCO instances format.
|
359 |
+
image_root: str
|
360 |
+
directory that contains all the images
|
361 |
+
dataset_name: str
|
362 |
+
the name that identifies a dataset, e.g. "densepose_coco_2014_train"
|
363 |
+
extra_annotation_keys: Optional[List[str]]
|
364 |
+
If provided, these keys are used to extract additional data from
|
365 |
+
the annotations.
|
366 |
+
"""
|
367 |
+
coco_api = _load_coco_annotations(PathManager.get_local_path(annotations_json_file))
|
368 |
+
_add_categories_metadata(dataset_name, coco_api.loadCats(coco_api.getCatIds()))
|
369 |
+
# sort indices for reproducible results
|
370 |
+
img_ids = sorted(coco_api.imgs.keys())
|
371 |
+
# imgs is a list of dicts, each looks something like:
|
372 |
+
# {'license': 4,
|
373 |
+
# 'url': 'http://farm6.staticflickr.com/5454/9413846304_881d5e5c3b_z.jpg',
|
374 |
+
# 'file_name': 'COCO_val2014_000000001268.jpg',
|
375 |
+
# 'height': 427,
|
376 |
+
# 'width': 640,
|
377 |
+
# 'date_captured': '2013-11-17 05:57:24',
|
378 |
+
# 'id': 1268}
|
379 |
+
imgs = coco_api.loadImgs(img_ids)
|
380 |
+
logger = logging.getLogger(__name__)
|
381 |
+
logger.info("Loaded {} images in COCO format from {}".format(len(imgs), annotations_json_file))
|
382 |
+
# anns is a list[list[dict]], where each dict is an annotation
|
383 |
+
# record for an object. The inner list enumerates the objects in an image
|
384 |
+
# and the outer list enumerates over images.
|
385 |
+
anns = [coco_api.imgToAnns[img_id] for img_id in img_ids]
|
386 |
+
_verify_annotations_have_unique_ids(annotations_json_file, anns)
|
387 |
+
dataset_records = _combine_images_with_annotations(dataset_name, image_root, imgs, anns)
|
388 |
+
return dataset_records
|
389 |
+
|
390 |
+
|
391 |
+
def register_dataset(dataset_data: CocoDatasetInfo, datasets_root: Optional[str] = None):
|
392 |
+
"""
|
393 |
+
Registers provided COCO DensePose dataset
|
394 |
+
|
395 |
+
Args:
|
396 |
+
dataset_data: CocoDatasetInfo
|
397 |
+
Dataset data
|
398 |
+
datasets_root: Optional[str]
|
399 |
+
Datasets root folder (default: None)
|
400 |
+
"""
|
401 |
+
annotations_fpath = maybe_prepend_base_path(datasets_root, dataset_data.annotations_fpath)
|
402 |
+
images_root = maybe_prepend_base_path(datasets_root, dataset_data.images_root)
|
403 |
+
|
404 |
+
def load_annotations():
|
405 |
+
return load_coco_json(
|
406 |
+
annotations_json_file=annotations_fpath,
|
407 |
+
image_root=images_root,
|
408 |
+
dataset_name=dataset_data.name,
|
409 |
+
)
|
410 |
+
|
411 |
+
DatasetCatalog.register(dataset_data.name, load_annotations)
|
412 |
+
MetadataCatalog.get(dataset_data.name).set(
|
413 |
+
json_file=annotations_fpath,
|
414 |
+
image_root=images_root,
|
415 |
+
**get_metadata(DENSEPOSE_METADATA_URL_PREFIX)
|
416 |
+
)
|
417 |
+
|
418 |
+
|
419 |
+
def register_datasets(
|
420 |
+
datasets_data: Iterable[CocoDatasetInfo], datasets_root: Optional[str] = None
|
421 |
+
):
|
422 |
+
"""
|
423 |
+
Registers provided COCO DensePose datasets
|
424 |
+
|
425 |
+
Args:
|
426 |
+
datasets_data: Iterable[CocoDatasetInfo]
|
427 |
+
An iterable of dataset datas
|
428 |
+
datasets_root: Optional[str]
|
429 |
+
Datasets root folder (default: None)
|
430 |
+
"""
|
431 |
+
for dataset_data in datasets_data:
|
432 |
+
register_dataset(dataset_data, datasets_root)
|
3rdparty/densepose/data/datasets/dataset_type.py
ADDED
@@ -0,0 +1,11 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
|
3 |
+
from enum import Enum
|
4 |
+
|
5 |
+
|
6 |
+
class DatasetType(Enum):
|
7 |
+
"""
|
8 |
+
Dataset type, mostly used for datasets that contain data to bootstrap models on
|
9 |
+
"""
|
10 |
+
|
11 |
+
VIDEO_LIST = "video_list"
|
3rdparty/densepose/data/datasets/lvis.py
ADDED
@@ -0,0 +1,257 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
import logging
|
3 |
+
import os
|
4 |
+
from typing import Any, Dict, Iterable, List, Optional
|
5 |
+
from fvcore.common.timer import Timer
|
6 |
+
|
7 |
+
from detectron2.data import DatasetCatalog, MetadataCatalog
|
8 |
+
from detectron2.data.datasets.lvis import get_lvis_instances_meta
|
9 |
+
from detectron2.structures import BoxMode
|
10 |
+
from detectron2.utils.file_io import PathManager
|
11 |
+
|
12 |
+
from ..utils import maybe_prepend_base_path
|
13 |
+
from .coco import (
|
14 |
+
DENSEPOSE_ALL_POSSIBLE_KEYS,
|
15 |
+
DENSEPOSE_METADATA_URL_PREFIX,
|
16 |
+
CocoDatasetInfo,
|
17 |
+
get_metadata,
|
18 |
+
)
|
19 |
+
|
20 |
+
DATASETS = [
|
21 |
+
CocoDatasetInfo(
|
22 |
+
name="densepose_lvis_v1_ds1_train_v1",
|
23 |
+
images_root="coco_",
|
24 |
+
annotations_fpath="lvis/densepose_lvis_v1_ds1_train_v1.json",
|
25 |
+
),
|
26 |
+
CocoDatasetInfo(
|
27 |
+
name="densepose_lvis_v1_ds1_val_v1",
|
28 |
+
images_root="coco_",
|
29 |
+
annotations_fpath="lvis/densepose_lvis_v1_ds1_val_v1.json",
|
30 |
+
),
|
31 |
+
CocoDatasetInfo(
|
32 |
+
name="densepose_lvis_v1_ds2_train_v1",
|
33 |
+
images_root="coco_",
|
34 |
+
annotations_fpath="lvis/densepose_lvis_v1_ds2_train_v1.json",
|
35 |
+
),
|
36 |
+
CocoDatasetInfo(
|
37 |
+
name="densepose_lvis_v1_ds2_val_v1",
|
38 |
+
images_root="coco_",
|
39 |
+
annotations_fpath="lvis/densepose_lvis_v1_ds2_val_v1.json",
|
40 |
+
),
|
41 |
+
CocoDatasetInfo(
|
42 |
+
name="densepose_lvis_v1_ds1_val_animals_100",
|
43 |
+
images_root="coco_",
|
44 |
+
annotations_fpath="lvis/densepose_lvis_v1_val_animals_100_v2.json",
|
45 |
+
),
|
46 |
+
]
|
47 |
+
|
48 |
+
|
49 |
+
def _load_lvis_annotations(json_file: str):
|
50 |
+
"""
|
51 |
+
Load COCO annotations from a JSON file
|
52 |
+
|
53 |
+
Args:
|
54 |
+
json_file: str
|
55 |
+
Path to the file to load annotations from
|
56 |
+
Returns:
|
57 |
+
Instance of `pycocotools.coco.COCO` that provides access to annotations
|
58 |
+
data
|
59 |
+
"""
|
60 |
+
from lvis import LVIS
|
61 |
+
|
62 |
+
json_file = PathManager.get_local_path(json_file)
|
63 |
+
logger = logging.getLogger(__name__)
|
64 |
+
timer = Timer()
|
65 |
+
lvis_api = LVIS(json_file)
|
66 |
+
if timer.seconds() > 1:
|
67 |
+
logger.info("Loading {} takes {:.2f} seconds.".format(json_file, timer.seconds()))
|
68 |
+
return lvis_api
|
69 |
+
|
70 |
+
|
71 |
+
def _add_categories_metadata(dataset_name: str) -> None:
|
72 |
+
metadict = get_lvis_instances_meta(dataset_name)
|
73 |
+
categories = metadict["thing_classes"]
|
74 |
+
metadata = MetadataCatalog.get(dataset_name)
|
75 |
+
metadata.categories = {i + 1: categories[i] for i in range(len(categories))}
|
76 |
+
logger = logging.getLogger(__name__)
|
77 |
+
logger.info(f"Dataset {dataset_name} has {len(categories)} categories")
|
78 |
+
|
79 |
+
|
80 |
+
def _verify_annotations_have_unique_ids(json_file: str, anns: List[List[Dict[str, Any]]]) -> None:
|
81 |
+
ann_ids = [ann["id"] for anns_per_image in anns for ann in anns_per_image]
|
82 |
+
assert len(set(ann_ids)) == len(ann_ids), "Annotation ids in '{}' are not unique!".format(
|
83 |
+
json_file
|
84 |
+
)
|
85 |
+
|
86 |
+
|
87 |
+
def _maybe_add_bbox(obj: Dict[str, Any], ann_dict: Dict[str, Any]) -> None:
|
88 |
+
if "bbox" not in ann_dict:
|
89 |
+
return
|
90 |
+
obj["bbox"] = ann_dict["bbox"]
|
91 |
+
obj["bbox_mode"] = BoxMode.XYWH_ABS
|
92 |
+
|
93 |
+
|
94 |
+
def _maybe_add_segm(obj: Dict[str, Any], ann_dict: Dict[str, Any]) -> None:
|
95 |
+
if "segmentation" not in ann_dict:
|
96 |
+
return
|
97 |
+
segm = ann_dict["segmentation"]
|
98 |
+
if not isinstance(segm, dict):
|
99 |
+
# filter out invalid polygons (< 3 points)
|
100 |
+
segm = [poly for poly in segm if len(poly) % 2 == 0 and len(poly) >= 6]
|
101 |
+
if len(segm) == 0:
|
102 |
+
return
|
103 |
+
obj["segmentation"] = segm
|
104 |
+
|
105 |
+
|
106 |
+
def _maybe_add_keypoints(obj: Dict[str, Any], ann_dict: Dict[str, Any]) -> None:
|
107 |
+
if "keypoints" not in ann_dict:
|
108 |
+
return
|
109 |
+
keypts = ann_dict["keypoints"] # list[int]
|
110 |
+
for idx, v in enumerate(keypts):
|
111 |
+
if idx % 3 != 2:
|
112 |
+
# COCO's segmentation coordinates are floating points in [0, H or W],
|
113 |
+
# but keypoint coordinates are integers in [0, H-1 or W-1]
|
114 |
+
# Therefore we assume the coordinates are "pixel indices" and
|
115 |
+
# add 0.5 to convert to floating point coordinates.
|
116 |
+
keypts[idx] = v + 0.5
|
117 |
+
obj["keypoints"] = keypts
|
118 |
+
|
119 |
+
|
120 |
+
def _maybe_add_densepose(obj: Dict[str, Any], ann_dict: Dict[str, Any]) -> None:
|
121 |
+
for key in DENSEPOSE_ALL_POSSIBLE_KEYS:
|
122 |
+
if key in ann_dict:
|
123 |
+
obj[key] = ann_dict[key]
|
124 |
+
|
125 |
+
|
126 |
+
def _combine_images_with_annotations(
|
127 |
+
dataset_name: str,
|
128 |
+
image_root: str,
|
129 |
+
img_datas: Iterable[Dict[str, Any]],
|
130 |
+
ann_datas: Iterable[Iterable[Dict[str, Any]]],
|
131 |
+
):
|
132 |
+
|
133 |
+
dataset_dicts = []
|
134 |
+
|
135 |
+
def get_file_name(img_root, img_dict):
|
136 |
+
# Determine the path including the split folder ("train2017", "val2017", "test2017") from
|
137 |
+
# the coco_url field. Example:
|
138 |
+
# 'coco_url': 'http://images.cocodataset.org/train2017/000000155379.jpg'
|
139 |
+
split_folder, file_name = img_dict["coco_url"].split("/")[-2:]
|
140 |
+
return os.path.join(img_root + split_folder, file_name)
|
141 |
+
|
142 |
+
for img_dict, ann_dicts in zip(img_datas, ann_datas):
|
143 |
+
record = {}
|
144 |
+
record["file_name"] = get_file_name(image_root, img_dict)
|
145 |
+
record["height"] = img_dict["height"]
|
146 |
+
record["width"] = img_dict["width"]
|
147 |
+
record["not_exhaustive_category_ids"] = img_dict.get("not_exhaustive_category_ids", [])
|
148 |
+
record["neg_category_ids"] = img_dict.get("neg_category_ids", [])
|
149 |
+
record["image_id"] = img_dict["id"]
|
150 |
+
record["dataset"] = dataset_name
|
151 |
+
|
152 |
+
objs = []
|
153 |
+
for ann_dict in ann_dicts:
|
154 |
+
assert ann_dict["image_id"] == record["image_id"]
|
155 |
+
obj = {}
|
156 |
+
_maybe_add_bbox(obj, ann_dict)
|
157 |
+
obj["iscrowd"] = ann_dict.get("iscrowd", 0)
|
158 |
+
obj["category_id"] = ann_dict["category_id"]
|
159 |
+
_maybe_add_segm(obj, ann_dict)
|
160 |
+
_maybe_add_keypoints(obj, ann_dict)
|
161 |
+
_maybe_add_densepose(obj, ann_dict)
|
162 |
+
objs.append(obj)
|
163 |
+
record["annotations"] = objs
|
164 |
+
dataset_dicts.append(record)
|
165 |
+
return dataset_dicts
|
166 |
+
|
167 |
+
|
168 |
+
def load_lvis_json(annotations_json_file: str, image_root: str, dataset_name: str):
|
169 |
+
"""
|
170 |
+
Loads a JSON file with annotations in LVIS instances format.
|
171 |
+
Replaces `detectron2.data.datasets.coco.load_lvis_json` to handle metadata
|
172 |
+
in a more flexible way. Postpones category mapping to a later stage to be
|
173 |
+
able to combine several datasets with different (but coherent) sets of
|
174 |
+
categories.
|
175 |
+
|
176 |
+
Args:
|
177 |
+
|
178 |
+
annotations_json_file: str
|
179 |
+
Path to the JSON file with annotations in COCO instances format.
|
180 |
+
image_root: str
|
181 |
+
directory that contains all the images
|
182 |
+
dataset_name: str
|
183 |
+
the name that identifies a dataset, e.g. "densepose_coco_2014_train"
|
184 |
+
extra_annotation_keys: Optional[List[str]]
|
185 |
+
If provided, these keys are used to extract additional data from
|
186 |
+
the annotations.
|
187 |
+
"""
|
188 |
+
lvis_api = _load_lvis_annotations(PathManager.get_local_path(annotations_json_file))
|
189 |
+
|
190 |
+
_add_categories_metadata(dataset_name)
|
191 |
+
|
192 |
+
# sort indices for reproducible results
|
193 |
+
img_ids = sorted(lvis_api.imgs.keys())
|
194 |
+
# imgs is a list of dicts, each looks something like:
|
195 |
+
# {'license': 4,
|
196 |
+
# 'url': 'http://farm6.staticflickr.com/5454/9413846304_881d5e5c3b_z.jpg',
|
197 |
+
# 'file_name': 'COCO_val2014_000000001268.jpg',
|
198 |
+
# 'height': 427,
|
199 |
+
# 'width': 640,
|
200 |
+
# 'date_captured': '2013-11-17 05:57:24',
|
201 |
+
# 'id': 1268}
|
202 |
+
imgs = lvis_api.load_imgs(img_ids)
|
203 |
+
logger = logging.getLogger(__name__)
|
204 |
+
logger.info("Loaded {} images in LVIS format from {}".format(len(imgs), annotations_json_file))
|
205 |
+
# anns is a list[list[dict]], where each dict is an annotation
|
206 |
+
# record for an object. The inner list enumerates the objects in an image
|
207 |
+
# and the outer list enumerates over images.
|
208 |
+
anns = [lvis_api.img_ann_map[img_id] for img_id in img_ids]
|
209 |
+
|
210 |
+
_verify_annotations_have_unique_ids(annotations_json_file, anns)
|
211 |
+
dataset_records = _combine_images_with_annotations(dataset_name, image_root, imgs, anns)
|
212 |
+
return dataset_records
|
213 |
+
|
214 |
+
|
215 |
+
def register_dataset(dataset_data: CocoDatasetInfo, datasets_root: Optional[str] = None) -> None:
|
216 |
+
"""
|
217 |
+
Registers provided LVIS DensePose dataset
|
218 |
+
|
219 |
+
Args:
|
220 |
+
dataset_data: CocoDatasetInfo
|
221 |
+
Dataset data
|
222 |
+
datasets_root: Optional[str]
|
223 |
+
Datasets root folder (default: None)
|
224 |
+
"""
|
225 |
+
annotations_fpath = maybe_prepend_base_path(datasets_root, dataset_data.annotations_fpath)
|
226 |
+
images_root = maybe_prepend_base_path(datasets_root, dataset_data.images_root)
|
227 |
+
|
228 |
+
def load_annotations():
|
229 |
+
return load_lvis_json(
|
230 |
+
annotations_json_file=annotations_fpath,
|
231 |
+
image_root=images_root,
|
232 |
+
dataset_name=dataset_data.name,
|
233 |
+
)
|
234 |
+
|
235 |
+
DatasetCatalog.register(dataset_data.name, load_annotations)
|
236 |
+
MetadataCatalog.get(dataset_data.name).set(
|
237 |
+
json_file=annotations_fpath,
|
238 |
+
image_root=images_root,
|
239 |
+
evaluator_type="lvis",
|
240 |
+
**get_metadata(DENSEPOSE_METADATA_URL_PREFIX),
|
241 |
+
)
|
242 |
+
|
243 |
+
|
244 |
+
def register_datasets(
|
245 |
+
datasets_data: Iterable[CocoDatasetInfo], datasets_root: Optional[str] = None
|
246 |
+
) -> None:
|
247 |
+
"""
|
248 |
+
Registers provided LVIS DensePose datasets
|
249 |
+
|
250 |
+
Args:
|
251 |
+
datasets_data: Iterable[CocoDatasetInfo]
|
252 |
+
An iterable of dataset datas
|
253 |
+
datasets_root: Optional[str]
|
254 |
+
Datasets root folder (default: None)
|
255 |
+
"""
|
256 |
+
for dataset_data in datasets_data:
|
257 |
+
register_dataset(dataset_data, datasets_root)
|
3rdparty/densepose/data/image_list_dataset.py
ADDED
@@ -0,0 +1,72 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# -*- coding: utf-8 -*-
|
2 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
3 |
+
|
4 |
+
import logging
|
5 |
+
import numpy as np
|
6 |
+
from typing import Any, Callable, Dict, List, Optional, Union
|
7 |
+
import torch
|
8 |
+
from torch.utils.data.dataset import Dataset
|
9 |
+
|
10 |
+
from detectron2.data.detection_utils import read_image
|
11 |
+
|
12 |
+
ImageTransform = Callable[[torch.Tensor], torch.Tensor]
|
13 |
+
|
14 |
+
|
15 |
+
class ImageListDataset(Dataset):
|
16 |
+
"""
|
17 |
+
Dataset that provides images from a list.
|
18 |
+
"""
|
19 |
+
|
20 |
+
_EMPTY_IMAGE = torch.empty((0, 3, 1, 1))
|
21 |
+
|
22 |
+
def __init__(
|
23 |
+
self,
|
24 |
+
image_list: List[str],
|
25 |
+
category_list: Union[str, List[str], None] = None,
|
26 |
+
transform: Optional[ImageTransform] = None,
|
27 |
+
):
|
28 |
+
"""
|
29 |
+
Args:
|
30 |
+
image_list (List[str]): list of paths to image files
|
31 |
+
category_list (Union[str, List[str], None]): list of animal categories for
|
32 |
+
each image. If it is a string, or None, this applies to all images
|
33 |
+
"""
|
34 |
+
if type(category_list) == list:
|
35 |
+
self.category_list = category_list
|
36 |
+
else:
|
37 |
+
self.category_list = [category_list] * len(image_list)
|
38 |
+
assert len(image_list) == len(
|
39 |
+
self.category_list
|
40 |
+
), "length of image and category lists must be equal"
|
41 |
+
self.image_list = image_list
|
42 |
+
self.transform = transform
|
43 |
+
|
44 |
+
def __getitem__(self, idx: int) -> Dict[str, Any]:
|
45 |
+
"""
|
46 |
+
Gets selected images from the list
|
47 |
+
|
48 |
+
Args:
|
49 |
+
idx (int): video index in the video list file
|
50 |
+
Returns:
|
51 |
+
A dictionary containing two keys:
|
52 |
+
images (torch.Tensor): tensor of size [N, 3, H, W] (N = 1, or 0 for _EMPTY_IMAGE)
|
53 |
+
categories (List[str]): categories of the frames
|
54 |
+
"""
|
55 |
+
categories = [self.category_list[idx]]
|
56 |
+
fpath = self.image_list[idx]
|
57 |
+
transform = self.transform
|
58 |
+
|
59 |
+
try:
|
60 |
+
image = torch.from_numpy(np.ascontiguousarray(read_image(fpath, format="BGR")))
|
61 |
+
image = image.permute(2, 0, 1).unsqueeze(0).float() # HWC -> NCHW
|
62 |
+
if transform is not None:
|
63 |
+
image = transform(image)
|
64 |
+
return {"images": image, "categories": categories}
|
65 |
+
except (OSError, RuntimeError) as e:
|
66 |
+
logger = logging.getLogger(__name__)
|
67 |
+
logger.warning(f"Error opening image file container {fpath}: {e}")
|
68 |
+
|
69 |
+
return {"images": self._EMPTY_IMAGE, "categories": []}
|
70 |
+
|
71 |
+
def __len__(self):
|
72 |
+
return len(self.image_list)
|
3rdparty/densepose/data/inference_based_loader.py
ADDED
@@ -0,0 +1,172 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
|
3 |
+
import random
|
4 |
+
from typing import Any, Callable, Dict, Iterable, Iterator, List, Optional, Tuple
|
5 |
+
import torch
|
6 |
+
from torch import nn
|
7 |
+
|
8 |
+
SampledData = Any
|
9 |
+
ModelOutput = Any
|
10 |
+
|
11 |
+
|
12 |
+
def _grouper(iterable: Iterable[Any], n: int, fillvalue=None) -> Iterator[Tuple[Any]]:
|
13 |
+
"""
|
14 |
+
Group elements of an iterable by chunks of size `n`, e.g.
|
15 |
+
grouper(range(9), 4) ->
|
16 |
+
(0, 1, 2, 3), (4, 5, 6, 7), (8, None, None, None)
|
17 |
+
"""
|
18 |
+
it = iter(iterable)
|
19 |
+
while True:
|
20 |
+
values = []
|
21 |
+
for _ in range(n):
|
22 |
+
try:
|
23 |
+
value = next(it)
|
24 |
+
except StopIteration:
|
25 |
+
if values:
|
26 |
+
values.extend([fillvalue] * (n - len(values)))
|
27 |
+
yield tuple(values)
|
28 |
+
return
|
29 |
+
values.append(value)
|
30 |
+
yield tuple(values)
|
31 |
+
|
32 |
+
|
33 |
+
class ScoreBasedFilter:
|
34 |
+
"""
|
35 |
+
Filters entries in model output based on their scores
|
36 |
+
Discards all entries with score less than the specified minimum
|
37 |
+
"""
|
38 |
+
|
39 |
+
def __init__(self, min_score: float = 0.8):
|
40 |
+
self.min_score = min_score
|
41 |
+
|
42 |
+
def __call__(self, model_output: ModelOutput) -> ModelOutput:
|
43 |
+
for model_output_i in model_output:
|
44 |
+
instances = model_output_i["instances"]
|
45 |
+
if not instances.has("scores"):
|
46 |
+
continue
|
47 |
+
instances_filtered = instances[instances.scores >= self.min_score]
|
48 |
+
model_output_i["instances"] = instances_filtered
|
49 |
+
return model_output
|
50 |
+
|
51 |
+
|
52 |
+
class InferenceBasedLoader:
|
53 |
+
"""
|
54 |
+
Data loader based on results inferred by a model. Consists of:
|
55 |
+
- a data loader that provides batches of images
|
56 |
+
- a model that is used to infer the results
|
57 |
+
- a data sampler that converts inferred results to annotations
|
58 |
+
"""
|
59 |
+
|
60 |
+
def __init__(
|
61 |
+
self,
|
62 |
+
model: nn.Module,
|
63 |
+
data_loader: Iterable[List[Dict[str, Any]]],
|
64 |
+
data_sampler: Optional[Callable[[ModelOutput], List[SampledData]]] = None,
|
65 |
+
data_filter: Optional[Callable[[ModelOutput], ModelOutput]] = None,
|
66 |
+
shuffle: bool = True,
|
67 |
+
batch_size: int = 4,
|
68 |
+
inference_batch_size: int = 4,
|
69 |
+
drop_last: bool = False,
|
70 |
+
category_to_class_mapping: Optional[dict] = None,
|
71 |
+
):
|
72 |
+
"""
|
73 |
+
Constructor
|
74 |
+
|
75 |
+
Args:
|
76 |
+
model (torch.nn.Module): model used to produce data
|
77 |
+
data_loader (Iterable[List[Dict[str, Any]]]): iterable that provides
|
78 |
+
dictionaries with "images" and "categories" fields to perform inference on
|
79 |
+
data_sampler (Callable: ModelOutput -> SampledData): functor
|
80 |
+
that produces annotation data from inference results;
|
81 |
+
(optional, default: None)
|
82 |
+
data_filter (Callable: ModelOutput -> ModelOutput): filter
|
83 |
+
that selects model outputs for further processing
|
84 |
+
(optional, default: None)
|
85 |
+
shuffle (bool): if True, the input images get shuffled
|
86 |
+
batch_size (int): batch size for the produced annotation data
|
87 |
+
inference_batch_size (int): batch size for input images
|
88 |
+
drop_last (bool): if True, drop the last batch if it is undersized
|
89 |
+
category_to_class_mapping (dict): category to class mapping
|
90 |
+
"""
|
91 |
+
self.model = model
|
92 |
+
self.model.eval()
|
93 |
+
self.data_loader = data_loader
|
94 |
+
self.data_sampler = data_sampler
|
95 |
+
self.data_filter = data_filter
|
96 |
+
self.shuffle = shuffle
|
97 |
+
self.batch_size = batch_size
|
98 |
+
self.inference_batch_size = inference_batch_size
|
99 |
+
self.drop_last = drop_last
|
100 |
+
if category_to_class_mapping is not None:
|
101 |
+
self.category_to_class_mapping = category_to_class_mapping
|
102 |
+
else:
|
103 |
+
self.category_to_class_mapping = {}
|
104 |
+
|
105 |
+
def __iter__(self) -> Iterator[List[SampledData]]:
|
106 |
+
for batch in self.data_loader:
|
107 |
+
# batch : List[Dict[str: Tensor[N, C, H, W], str: Optional[str]]]
|
108 |
+
# images_batch : Tensor[N, C, H, W]
|
109 |
+
# image : Tensor[C, H, W]
|
110 |
+
images_and_categories = [
|
111 |
+
{"image": image, "category": category}
|
112 |
+
for element in batch
|
113 |
+
for image, category in zip(element["images"], element["categories"])
|
114 |
+
]
|
115 |
+
if not images_and_categories:
|
116 |
+
continue
|
117 |
+
if self.shuffle:
|
118 |
+
random.shuffle(images_and_categories)
|
119 |
+
yield from self._produce_data(images_and_categories) # pyre-ignore[6]
|
120 |
+
|
121 |
+
def _produce_data(
|
122 |
+
self, images_and_categories: List[Tuple[torch.Tensor, Optional[str]]]
|
123 |
+
) -> Iterator[List[SampledData]]:
|
124 |
+
"""
|
125 |
+
Produce batches of data from images
|
126 |
+
|
127 |
+
Args:
|
128 |
+
images_and_categories (List[Tuple[torch.Tensor, Optional[str]]]):
|
129 |
+
list of images and corresponding categories to process
|
130 |
+
|
131 |
+
Returns:
|
132 |
+
Iterator over batches of data sampled from model outputs
|
133 |
+
"""
|
134 |
+
data_batches: List[SampledData] = []
|
135 |
+
category_to_class_mapping = self.category_to_class_mapping
|
136 |
+
batched_images_and_categories = _grouper(images_and_categories, self.inference_batch_size)
|
137 |
+
for batch in batched_images_and_categories:
|
138 |
+
batch = [
|
139 |
+
{
|
140 |
+
"image": image_and_category["image"].to(self.model.device),
|
141 |
+
"category": image_and_category["category"],
|
142 |
+
}
|
143 |
+
for image_and_category in batch
|
144 |
+
if image_and_category is not None
|
145 |
+
]
|
146 |
+
if not batch:
|
147 |
+
continue
|
148 |
+
with torch.no_grad():
|
149 |
+
model_output = self.model(batch)
|
150 |
+
for model_output_i, batch_i in zip(model_output, batch):
|
151 |
+
assert len(batch_i["image"].shape) == 3
|
152 |
+
model_output_i["image"] = batch_i["image"]
|
153 |
+
instance_class = category_to_class_mapping.get(batch_i["category"], 0)
|
154 |
+
model_output_i["instances"].dataset_classes = torch.tensor(
|
155 |
+
[instance_class] * len(model_output_i["instances"])
|
156 |
+
)
|
157 |
+
model_output_filtered = (
|
158 |
+
model_output if self.data_filter is None else self.data_filter(model_output)
|
159 |
+
)
|
160 |
+
data = (
|
161 |
+
model_output_filtered
|
162 |
+
if self.data_sampler is None
|
163 |
+
else self.data_sampler(model_output_filtered)
|
164 |
+
)
|
165 |
+
for data_i in data:
|
166 |
+
if len(data_i["instances"]):
|
167 |
+
data_batches.append(data_i)
|
168 |
+
if len(data_batches) >= self.batch_size:
|
169 |
+
yield data_batches[: self.batch_size]
|
170 |
+
data_batches = data_batches[self.batch_size :]
|
171 |
+
if not self.drop_last and data_batches:
|
172 |
+
yield data_batches
|
3rdparty/densepose/data/meshes/__init__.py
ADDED
@@ -0,0 +1,5 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
|
2 |
+
|
3 |
+
from . import builtin
|
4 |
+
|
5 |
+
__all__ = [k for k in globals().keys() if "builtin" not in k and not k.startswith("_")]
|
3rdparty/densepose/data/meshes/builtin.py
ADDED
@@ -0,0 +1,101 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
|
2 |
+
|
3 |
+
from .catalog import MeshInfo, register_meshes
|
4 |
+
|
5 |
+
DENSEPOSE_MESHES_DIR = "https://dl.fbaipublicfiles.com/densepose/meshes/"
|
6 |
+
|
7 |
+
MESHES = [
|
8 |
+
MeshInfo(
|
9 |
+
name="smpl_27554",
|
10 |
+
data="smpl_27554.pkl",
|
11 |
+
geodists="geodists/geodists_smpl_27554.pkl",
|
12 |
+
symmetry="symmetry/symmetry_smpl_27554.pkl",
|
13 |
+
texcoords="texcoords/texcoords_smpl_27554.pkl",
|
14 |
+
),
|
15 |
+
MeshInfo(
|
16 |
+
name="chimp_5029",
|
17 |
+
data="chimp_5029.pkl",
|
18 |
+
geodists="geodists/geodists_chimp_5029.pkl",
|
19 |
+
symmetry="symmetry/symmetry_chimp_5029.pkl",
|
20 |
+
texcoords="texcoords/texcoords_chimp_5029.pkl",
|
21 |
+
),
|
22 |
+
MeshInfo(
|
23 |
+
name="cat_5001",
|
24 |
+
data="cat_5001.pkl",
|
25 |
+
geodists="geodists/geodists_cat_5001.pkl",
|
26 |
+
symmetry="symmetry/symmetry_cat_5001.pkl",
|
27 |
+
texcoords="texcoords/texcoords_cat_5001.pkl",
|
28 |
+
),
|
29 |
+
MeshInfo(
|
30 |
+
name="cat_7466",
|
31 |
+
data="cat_7466.pkl",
|
32 |
+
geodists="geodists/geodists_cat_7466.pkl",
|
33 |
+
symmetry="symmetry/symmetry_cat_7466.pkl",
|
34 |
+
texcoords="texcoords/texcoords_cat_7466.pkl",
|
35 |
+
),
|
36 |
+
MeshInfo(
|
37 |
+
name="sheep_5004",
|
38 |
+
data="sheep_5004.pkl",
|
39 |
+
geodists="geodists/geodists_sheep_5004.pkl",
|
40 |
+
symmetry="symmetry/symmetry_sheep_5004.pkl",
|
41 |
+
texcoords="texcoords/texcoords_sheep_5004.pkl",
|
42 |
+
),
|
43 |
+
MeshInfo(
|
44 |
+
name="zebra_5002",
|
45 |
+
data="zebra_5002.pkl",
|
46 |
+
geodists="geodists/geodists_zebra_5002.pkl",
|
47 |
+
symmetry="symmetry/symmetry_zebra_5002.pkl",
|
48 |
+
texcoords="texcoords/texcoords_zebra_5002.pkl",
|
49 |
+
),
|
50 |
+
MeshInfo(
|
51 |
+
name="horse_5004",
|
52 |
+
data="horse_5004.pkl",
|
53 |
+
geodists="geodists/geodists_horse_5004.pkl",
|
54 |
+
symmetry="symmetry/symmetry_horse_5004.pkl",
|
55 |
+
texcoords="texcoords/texcoords_zebra_5002.pkl",
|
56 |
+
),
|
57 |
+
MeshInfo(
|
58 |
+
name="giraffe_5002",
|
59 |
+
data="giraffe_5002.pkl",
|
60 |
+
geodists="geodists/geodists_giraffe_5002.pkl",
|
61 |
+
symmetry="symmetry/symmetry_giraffe_5002.pkl",
|
62 |
+
texcoords="texcoords/texcoords_giraffe_5002.pkl",
|
63 |
+
),
|
64 |
+
MeshInfo(
|
65 |
+
name="elephant_5002",
|
66 |
+
data="elephant_5002.pkl",
|
67 |
+
geodists="geodists/geodists_elephant_5002.pkl",
|
68 |
+
symmetry="symmetry/symmetry_elephant_5002.pkl",
|
69 |
+
texcoords="texcoords/texcoords_elephant_5002.pkl",
|
70 |
+
),
|
71 |
+
MeshInfo(
|
72 |
+
name="dog_5002",
|
73 |
+
data="dog_5002.pkl",
|
74 |
+
geodists="geodists/geodists_dog_5002.pkl",
|
75 |
+
symmetry="symmetry/symmetry_dog_5002.pkl",
|
76 |
+
texcoords="texcoords/texcoords_dog_5002.pkl",
|
77 |
+
),
|
78 |
+
MeshInfo(
|
79 |
+
name="dog_7466",
|
80 |
+
data="dog_7466.pkl",
|
81 |
+
geodists="geodists/geodists_dog_7466.pkl",
|
82 |
+
symmetry="symmetry/symmetry_dog_7466.pkl",
|
83 |
+
texcoords="texcoords/texcoords_dog_7466.pkl",
|
84 |
+
),
|
85 |
+
MeshInfo(
|
86 |
+
name="cow_5002",
|
87 |
+
data="cow_5002.pkl",
|
88 |
+
geodists="geodists/geodists_cow_5002.pkl",
|
89 |
+
symmetry="symmetry/symmetry_cow_5002.pkl",
|
90 |
+
texcoords="texcoords/texcoords_cow_5002.pkl",
|
91 |
+
),
|
92 |
+
MeshInfo(
|
93 |
+
name="bear_4936",
|
94 |
+
data="bear_4936.pkl",
|
95 |
+
geodists="geodists/geodists_bear_4936.pkl",
|
96 |
+
symmetry="symmetry/symmetry_bear_4936.pkl",
|
97 |
+
texcoords="texcoords/texcoords_bear_4936.pkl",
|
98 |
+
),
|
99 |
+
]
|
100 |
+
|
101 |
+
register_meshes(MESHES, DENSEPOSE_MESHES_DIR)
|
3rdparty/densepose/data/meshes/catalog.py
ADDED
@@ -0,0 +1,71 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
|
2 |
+
|
3 |
+
import logging
|
4 |
+
from collections import UserDict
|
5 |
+
from dataclasses import dataclass
|
6 |
+
from typing import Iterable, Optional
|
7 |
+
|
8 |
+
from ..utils import maybe_prepend_base_path
|
9 |
+
|
10 |
+
|
11 |
+
@dataclass
|
12 |
+
class MeshInfo:
|
13 |
+
name: str
|
14 |
+
data: str
|
15 |
+
geodists: Optional[str] = None
|
16 |
+
symmetry: Optional[str] = None
|
17 |
+
texcoords: Optional[str] = None
|
18 |
+
|
19 |
+
|
20 |
+
class _MeshCatalog(UserDict):
|
21 |
+
def __init__(self, *args, **kwargs):
|
22 |
+
super().__init__(*args, **kwargs)
|
23 |
+
self.mesh_ids = {}
|
24 |
+
self.mesh_names = {}
|
25 |
+
self.max_mesh_id = -1
|
26 |
+
|
27 |
+
def __setitem__(self, key, value):
|
28 |
+
if key in self:
|
29 |
+
logger = logging.getLogger(__name__)
|
30 |
+
logger.warning(
|
31 |
+
f"Overwriting mesh catalog entry '{key}': old value {self[key]}"
|
32 |
+
f", new value {value}"
|
33 |
+
)
|
34 |
+
mesh_id = self.mesh_ids[key]
|
35 |
+
else:
|
36 |
+
self.max_mesh_id += 1
|
37 |
+
mesh_id = self.max_mesh_id
|
38 |
+
super().__setitem__(key, value)
|
39 |
+
self.mesh_ids[key] = mesh_id
|
40 |
+
self.mesh_names[mesh_id] = key
|
41 |
+
|
42 |
+
def get_mesh_id(self, shape_name: str) -> int:
|
43 |
+
return self.mesh_ids[shape_name]
|
44 |
+
|
45 |
+
def get_mesh_name(self, mesh_id: int) -> str:
|
46 |
+
return self.mesh_names[mesh_id]
|
47 |
+
|
48 |
+
|
49 |
+
MeshCatalog = _MeshCatalog()
|
50 |
+
|
51 |
+
|
52 |
+
def register_mesh(mesh_info: MeshInfo, base_path: Optional[str]) -> None:
|
53 |
+
geodists, symmetry, texcoords = mesh_info.geodists, mesh_info.symmetry, mesh_info.texcoords
|
54 |
+
if geodists:
|
55 |
+
geodists = maybe_prepend_base_path(base_path, geodists)
|
56 |
+
if symmetry:
|
57 |
+
symmetry = maybe_prepend_base_path(base_path, symmetry)
|
58 |
+
if texcoords:
|
59 |
+
texcoords = maybe_prepend_base_path(base_path, texcoords)
|
60 |
+
MeshCatalog[mesh_info.name] = MeshInfo(
|
61 |
+
name=mesh_info.name,
|
62 |
+
data=maybe_prepend_base_path(base_path, mesh_info.data),
|
63 |
+
geodists=geodists,
|
64 |
+
symmetry=symmetry,
|
65 |
+
texcoords=texcoords,
|
66 |
+
)
|
67 |
+
|
68 |
+
|
69 |
+
def register_meshes(mesh_infos: Iterable[MeshInfo], base_path: Optional[str]) -> None:
|
70 |
+
for mesh_info in mesh_infos:
|
71 |
+
register_mesh(mesh_info, base_path)
|
3rdparty/densepose/data/samplers/__init__.py
ADDED
@@ -0,0 +1,8 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
|
3 |
+
from .densepose_uniform import DensePoseUniformSampler
|
4 |
+
from .densepose_confidence_based import DensePoseConfidenceBasedSampler
|
5 |
+
from .densepose_cse_uniform import DensePoseCSEUniformSampler
|
6 |
+
from .densepose_cse_confidence_based import DensePoseCSEConfidenceBasedSampler
|
7 |
+
from .mask_from_densepose import MaskFromDensePoseSampler
|
8 |
+
from .prediction_to_gt import PredictionToGroundTruthSampler
|
3rdparty/densepose/data/samplers/densepose_base.py
ADDED
@@ -0,0 +1,203 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
|
3 |
+
from typing import Any, Dict, List, Tuple
|
4 |
+
import torch
|
5 |
+
from torch.nn import functional as F
|
6 |
+
|
7 |
+
from detectron2.structures import BoxMode, Instances
|
8 |
+
|
9 |
+
from densepose.converters import ToChartResultConverter
|
10 |
+
from densepose.converters.base import IntTupleBox, make_int_box
|
11 |
+
from densepose.structures import DensePoseDataRelative, DensePoseList
|
12 |
+
|
13 |
+
|
14 |
+
class DensePoseBaseSampler:
|
15 |
+
"""
|
16 |
+
Base DensePose sampler to produce DensePose data from DensePose predictions.
|
17 |
+
Samples for each class are drawn according to some distribution over all pixels estimated
|
18 |
+
to belong to that class.
|
19 |
+
"""
|
20 |
+
|
21 |
+
def __init__(self, count_per_class: int = 8):
|
22 |
+
"""
|
23 |
+
Constructor
|
24 |
+
|
25 |
+
Args:
|
26 |
+
count_per_class (int): the sampler produces at most `count_per_class`
|
27 |
+
samples for each category
|
28 |
+
"""
|
29 |
+
self.count_per_class = count_per_class
|
30 |
+
|
31 |
+
def __call__(self, instances: Instances) -> DensePoseList:
|
32 |
+
"""
|
33 |
+
Convert DensePose predictions (an instance of `DensePoseChartPredictorOutput`)
|
34 |
+
into DensePose annotations data (an instance of `DensePoseList`)
|
35 |
+
"""
|
36 |
+
boxes_xyxy_abs = instances.pred_boxes.tensor.clone().cpu()
|
37 |
+
boxes_xywh_abs = BoxMode.convert(boxes_xyxy_abs, BoxMode.XYXY_ABS, BoxMode.XYWH_ABS)
|
38 |
+
dp_datas = []
|
39 |
+
for i in range(len(boxes_xywh_abs)):
|
40 |
+
annotation_i = self._sample(instances[i], make_int_box(boxes_xywh_abs[i]))
|
41 |
+
annotation_i[DensePoseDataRelative.S_KEY] = self._resample_mask( # pyre-ignore[6]
|
42 |
+
instances[i].pred_densepose
|
43 |
+
)
|
44 |
+
dp_datas.append(DensePoseDataRelative(annotation_i))
|
45 |
+
# create densepose annotations on CPU
|
46 |
+
dp_list = DensePoseList(dp_datas, boxes_xyxy_abs, instances.image_size)
|
47 |
+
return dp_list
|
48 |
+
|
49 |
+
def _sample(self, instance: Instances, bbox_xywh: IntTupleBox) -> Dict[str, List[Any]]:
|
50 |
+
"""
|
51 |
+
Sample DensPoseDataRelative from estimation results
|
52 |
+
"""
|
53 |
+
labels, dp_result = self._produce_labels_and_results(instance)
|
54 |
+
annotation = {
|
55 |
+
DensePoseDataRelative.X_KEY: [],
|
56 |
+
DensePoseDataRelative.Y_KEY: [],
|
57 |
+
DensePoseDataRelative.U_KEY: [],
|
58 |
+
DensePoseDataRelative.V_KEY: [],
|
59 |
+
DensePoseDataRelative.I_KEY: [],
|
60 |
+
}
|
61 |
+
n, h, w = dp_result.shape
|
62 |
+
for part_id in range(1, DensePoseDataRelative.N_PART_LABELS + 1):
|
63 |
+
# indices - tuple of 3 1D tensors of size k
|
64 |
+
# 0: index along the first dimension N
|
65 |
+
# 1: index along H dimension
|
66 |
+
# 2: index along W dimension
|
67 |
+
indices = torch.nonzero(labels.expand(n, h, w) == part_id, as_tuple=True)
|
68 |
+
# values - an array of size [n, k]
|
69 |
+
# n: number of channels (U, V, confidences)
|
70 |
+
# k: number of points labeled with part_id
|
71 |
+
values = dp_result[indices].view(n, -1)
|
72 |
+
k = values.shape[1]
|
73 |
+
count = min(self.count_per_class, k)
|
74 |
+
if count <= 0:
|
75 |
+
continue
|
76 |
+
index_sample = self._produce_index_sample(values, count)
|
77 |
+
sampled_values = values[:, index_sample]
|
78 |
+
sampled_y = indices[1][index_sample] + 0.5
|
79 |
+
sampled_x = indices[2][index_sample] + 0.5
|
80 |
+
# prepare / normalize data
|
81 |
+
x = (sampled_x / w * 256.0).cpu().tolist()
|
82 |
+
y = (sampled_y / h * 256.0).cpu().tolist()
|
83 |
+
u = sampled_values[0].clamp(0, 1).cpu().tolist()
|
84 |
+
v = sampled_values[1].clamp(0, 1).cpu().tolist()
|
85 |
+
fine_segm_labels = [part_id] * count
|
86 |
+
# extend annotations
|
87 |
+
annotation[DensePoseDataRelative.X_KEY].extend(x)
|
88 |
+
annotation[DensePoseDataRelative.Y_KEY].extend(y)
|
89 |
+
annotation[DensePoseDataRelative.U_KEY].extend(u)
|
90 |
+
annotation[DensePoseDataRelative.V_KEY].extend(v)
|
91 |
+
annotation[DensePoseDataRelative.I_KEY].extend(fine_segm_labels)
|
92 |
+
return annotation
|
93 |
+
|
94 |
+
def _produce_index_sample(self, values: torch.Tensor, count: int):
|
95 |
+
"""
|
96 |
+
Abstract method to produce a sample of indices to select data
|
97 |
+
To be implemented in descendants
|
98 |
+
|
99 |
+
Args:
|
100 |
+
values (torch.Tensor): an array of size [n, k] that contains
|
101 |
+
estimated values (U, V, confidences);
|
102 |
+
n: number of channels (U, V, confidences)
|
103 |
+
k: number of points labeled with part_id
|
104 |
+
count (int): number of samples to produce, should be positive and <= k
|
105 |
+
|
106 |
+
Return:
|
107 |
+
list(int): indices of values (along axis 1) selected as a sample
|
108 |
+
"""
|
109 |
+
raise NotImplementedError
|
110 |
+
|
111 |
+
def _produce_labels_and_results(self, instance: Instances) -> Tuple[torch.Tensor, torch.Tensor]:
|
112 |
+
"""
|
113 |
+
Method to get labels and DensePose results from an instance
|
114 |
+
|
115 |
+
Args:
|
116 |
+
instance (Instances): an instance of `DensePoseChartPredictorOutput`
|
117 |
+
|
118 |
+
Return:
|
119 |
+
labels (torch.Tensor): shape [H, W], DensePose segmentation labels
|
120 |
+
dp_result (torch.Tensor): shape [2, H, W], stacked DensePose results u and v
|
121 |
+
"""
|
122 |
+
converter = ToChartResultConverter
|
123 |
+
chart_result = converter.convert(instance.pred_densepose, instance.pred_boxes)
|
124 |
+
labels, dp_result = chart_result.labels.cpu(), chart_result.uv.cpu()
|
125 |
+
return labels, dp_result
|
126 |
+
|
127 |
+
def _resample_mask(self, output: Any) -> torch.Tensor:
|
128 |
+
"""
|
129 |
+
Convert DensePose predictor output to segmentation annotation - tensors of size
|
130 |
+
(256, 256) and type `int64`.
|
131 |
+
|
132 |
+
Args:
|
133 |
+
output: DensePose predictor output with the following attributes:
|
134 |
+
- coarse_segm: tensor of size [N, D, H, W] with unnormalized coarse
|
135 |
+
segmentation scores
|
136 |
+
- fine_segm: tensor of size [N, C, H, W] with unnormalized fine
|
137 |
+
segmentation scores
|
138 |
+
Return:
|
139 |
+
Tensor of size (S, S) and type `int64` with coarse segmentation annotations,
|
140 |
+
where S = DensePoseDataRelative.MASK_SIZE
|
141 |
+
"""
|
142 |
+
sz = DensePoseDataRelative.MASK_SIZE
|
143 |
+
S = (
|
144 |
+
F.interpolate(output.coarse_segm, (sz, sz), mode="bilinear", align_corners=False)
|
145 |
+
.argmax(dim=1)
|
146 |
+
.long()
|
147 |
+
)
|
148 |
+
I = (
|
149 |
+
(
|
150 |
+
F.interpolate(
|
151 |
+
output.fine_segm,
|
152 |
+
(sz, sz),
|
153 |
+
mode="bilinear",
|
154 |
+
align_corners=False,
|
155 |
+
).argmax(dim=1)
|
156 |
+
* (S > 0).long()
|
157 |
+
)
|
158 |
+
.squeeze()
|
159 |
+
.cpu()
|
160 |
+
)
|
161 |
+
# Map fine segmentation results to coarse segmentation ground truth
|
162 |
+
# TODO: extract this into separate classes
|
163 |
+
# coarse segmentation: 1 = Torso, 2 = Right Hand, 3 = Left Hand,
|
164 |
+
# 4 = Left Foot, 5 = Right Foot, 6 = Upper Leg Right, 7 = Upper Leg Left,
|
165 |
+
# 8 = Lower Leg Right, 9 = Lower Leg Left, 10 = Upper Arm Left,
|
166 |
+
# 11 = Upper Arm Right, 12 = Lower Arm Left, 13 = Lower Arm Right,
|
167 |
+
# 14 = Head
|
168 |
+
# fine segmentation: 1, 2 = Torso, 3 = Right Hand, 4 = Left Hand,
|
169 |
+
# 5 = Left Foot, 6 = Right Foot, 7, 9 = Upper Leg Right,
|
170 |
+
# 8, 10 = Upper Leg Left, 11, 13 = Lower Leg Right,
|
171 |
+
# 12, 14 = Lower Leg Left, 15, 17 = Upper Arm Left,
|
172 |
+
# 16, 18 = Upper Arm Right, 19, 21 = Lower Arm Left,
|
173 |
+
# 20, 22 = Lower Arm Right, 23, 24 = Head
|
174 |
+
FINE_TO_COARSE_SEGMENTATION = {
|
175 |
+
1: 1,
|
176 |
+
2: 1,
|
177 |
+
3: 2,
|
178 |
+
4: 3,
|
179 |
+
5: 4,
|
180 |
+
6: 5,
|
181 |
+
7: 6,
|
182 |
+
8: 7,
|
183 |
+
9: 6,
|
184 |
+
10: 7,
|
185 |
+
11: 8,
|
186 |
+
12: 9,
|
187 |
+
13: 8,
|
188 |
+
14: 9,
|
189 |
+
15: 10,
|
190 |
+
16: 11,
|
191 |
+
17: 10,
|
192 |
+
18: 11,
|
193 |
+
19: 12,
|
194 |
+
20: 13,
|
195 |
+
21: 12,
|
196 |
+
22: 13,
|
197 |
+
23: 14,
|
198 |
+
24: 14,
|
199 |
+
}
|
200 |
+
mask = torch.zeros((sz, sz), dtype=torch.int64, device=torch.device("cpu"))
|
201 |
+
for i in range(DensePoseDataRelative.N_PART_LABELS):
|
202 |
+
mask[I == i + 1] = FINE_TO_COARSE_SEGMENTATION[i + 1]
|
203 |
+
return mask
|
3rdparty/densepose/data/samplers/densepose_confidence_based.py
ADDED
@@ -0,0 +1,108 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
|
3 |
+
import random
|
4 |
+
from typing import Optional, Tuple
|
5 |
+
import torch
|
6 |
+
|
7 |
+
from densepose.converters import ToChartResultConverterWithConfidences
|
8 |
+
|
9 |
+
from .densepose_base import DensePoseBaseSampler
|
10 |
+
|
11 |
+
|
12 |
+
class DensePoseConfidenceBasedSampler(DensePoseBaseSampler):
|
13 |
+
"""
|
14 |
+
Samples DensePose data from DensePose predictions.
|
15 |
+
Samples for each class are drawn using confidence value estimates.
|
16 |
+
"""
|
17 |
+
|
18 |
+
def __init__(
|
19 |
+
self,
|
20 |
+
confidence_channel: str,
|
21 |
+
count_per_class: int = 8,
|
22 |
+
search_count_multiplier: Optional[float] = None,
|
23 |
+
search_proportion: Optional[float] = None,
|
24 |
+
):
|
25 |
+
"""
|
26 |
+
Constructor
|
27 |
+
|
28 |
+
Args:
|
29 |
+
confidence_channel (str): confidence channel to use for sampling;
|
30 |
+
possible values:
|
31 |
+
"sigma_2": confidences for UV values
|
32 |
+
"fine_segm_confidence": confidences for fine segmentation
|
33 |
+
"coarse_segm_confidence": confidences for coarse segmentation
|
34 |
+
(default: "sigma_2")
|
35 |
+
count_per_class (int): the sampler produces at most `count_per_class`
|
36 |
+
samples for each category (default: 8)
|
37 |
+
search_count_multiplier (float or None): if not None, the total number
|
38 |
+
of the most confident estimates of a given class to consider is
|
39 |
+
defined as `min(search_count_multiplier * count_per_class, N)`,
|
40 |
+
where `N` is the total number of estimates of the class; cannot be
|
41 |
+
specified together with `search_proportion` (default: None)
|
42 |
+
search_proportion (float or None): if not None, the total number of the
|
43 |
+
of the most confident estimates of a given class to consider is
|
44 |
+
defined as `min(max(search_proportion * N, count_per_class), N)`,
|
45 |
+
where `N` is the total number of estimates of the class; cannot be
|
46 |
+
specified together with `search_count_multiplier` (default: None)
|
47 |
+
"""
|
48 |
+
super().__init__(count_per_class)
|
49 |
+
self.confidence_channel = confidence_channel
|
50 |
+
self.search_count_multiplier = search_count_multiplier
|
51 |
+
self.search_proportion = search_proportion
|
52 |
+
assert (search_count_multiplier is None) or (search_proportion is None), (
|
53 |
+
f"Cannot specify both search_count_multiplier (={search_count_multiplier})"
|
54 |
+
f"and search_proportion (={search_proportion})"
|
55 |
+
)
|
56 |
+
|
57 |
+
def _produce_index_sample(self, values: torch.Tensor, count: int):
|
58 |
+
"""
|
59 |
+
Produce a sample of indices to select data based on confidences
|
60 |
+
|
61 |
+
Args:
|
62 |
+
values (torch.Tensor): an array of size [n, k] that contains
|
63 |
+
estimated values (U, V, confidences);
|
64 |
+
n: number of channels (U, V, confidences)
|
65 |
+
k: number of points labeled with part_id
|
66 |
+
count (int): number of samples to produce, should be positive and <= k
|
67 |
+
|
68 |
+
Return:
|
69 |
+
list(int): indices of values (along axis 1) selected as a sample
|
70 |
+
"""
|
71 |
+
k = values.shape[1]
|
72 |
+
if k == count:
|
73 |
+
index_sample = list(range(k))
|
74 |
+
else:
|
75 |
+
# take the best count * search_count_multiplier pixels,
|
76 |
+
# sample from them uniformly
|
77 |
+
# (here best = smallest variance)
|
78 |
+
_, sorted_confidence_indices = torch.sort(values[2])
|
79 |
+
if self.search_count_multiplier is not None:
|
80 |
+
search_count = min(int(count * self.search_count_multiplier), k)
|
81 |
+
elif self.search_proportion is not None:
|
82 |
+
search_count = min(max(int(k * self.search_proportion), count), k)
|
83 |
+
else:
|
84 |
+
search_count = min(count, k)
|
85 |
+
sample_from_top = random.sample(range(search_count), count)
|
86 |
+
index_sample = sorted_confidence_indices[:search_count][sample_from_top]
|
87 |
+
return index_sample
|
88 |
+
|
89 |
+
def _produce_labels_and_results(self, instance) -> Tuple[torch.Tensor, torch.Tensor]:
|
90 |
+
"""
|
91 |
+
Method to get labels and DensePose results from an instance, with confidences
|
92 |
+
|
93 |
+
Args:
|
94 |
+
instance (Instances): an instance of `DensePoseChartPredictorOutputWithConfidences`
|
95 |
+
|
96 |
+
Return:
|
97 |
+
labels (torch.Tensor): shape [H, W], DensePose segmentation labels
|
98 |
+
dp_result (torch.Tensor): shape [3, H, W], DensePose results u and v
|
99 |
+
stacked with the confidence channel
|
100 |
+
"""
|
101 |
+
converter = ToChartResultConverterWithConfidences
|
102 |
+
chart_result = converter.convert(instance.pred_densepose, instance.pred_boxes)
|
103 |
+
labels, dp_result = chart_result.labels.cpu(), chart_result.uv.cpu()
|
104 |
+
dp_result = torch.cat(
|
105 |
+
(dp_result, getattr(chart_result, self.confidence_channel)[None].cpu())
|
106 |
+
)
|
107 |
+
|
108 |
+
return labels, dp_result
|
3rdparty/densepose/data/samplers/densepose_cse_base.py
ADDED
@@ -0,0 +1,139 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
|
3 |
+
from typing import Any, Dict, List, Tuple
|
4 |
+
import torch
|
5 |
+
from torch.nn import functional as F
|
6 |
+
|
7 |
+
from detectron2.config import CfgNode
|
8 |
+
from detectron2.structures import Instances
|
9 |
+
|
10 |
+
from densepose.converters.base import IntTupleBox
|
11 |
+
from densepose.data.utils import get_class_to_mesh_name_mapping
|
12 |
+
from densepose.modeling.cse.utils import squared_euclidean_distance_matrix
|
13 |
+
from densepose.structures import DensePoseDataRelative
|
14 |
+
|
15 |
+
from .densepose_base import DensePoseBaseSampler
|
16 |
+
|
17 |
+
|
18 |
+
class DensePoseCSEBaseSampler(DensePoseBaseSampler):
|
19 |
+
"""
|
20 |
+
Base DensePose sampler to produce DensePose data from DensePose predictions.
|
21 |
+
Samples for each class are drawn according to some distribution over all pixels estimated
|
22 |
+
to belong to that class.
|
23 |
+
"""
|
24 |
+
|
25 |
+
def __init__(
|
26 |
+
self,
|
27 |
+
cfg: CfgNode,
|
28 |
+
use_gt_categories: bool,
|
29 |
+
embedder: torch.nn.Module,
|
30 |
+
count_per_class: int = 8,
|
31 |
+
):
|
32 |
+
"""
|
33 |
+
Constructor
|
34 |
+
|
35 |
+
Args:
|
36 |
+
cfg (CfgNode): the config of the model
|
37 |
+
embedder (torch.nn.Module): necessary to compute mesh vertex embeddings
|
38 |
+
count_per_class (int): the sampler produces at most `count_per_class`
|
39 |
+
samples for each category
|
40 |
+
"""
|
41 |
+
super().__init__(count_per_class)
|
42 |
+
self.embedder = embedder
|
43 |
+
self.class_to_mesh_name = get_class_to_mesh_name_mapping(cfg)
|
44 |
+
self.use_gt_categories = use_gt_categories
|
45 |
+
|
46 |
+
def _sample(self, instance: Instances, bbox_xywh: IntTupleBox) -> Dict[str, List[Any]]:
|
47 |
+
"""
|
48 |
+
Sample DensPoseDataRelative from estimation results
|
49 |
+
"""
|
50 |
+
if self.use_gt_categories:
|
51 |
+
instance_class = instance.dataset_classes.tolist()[0]
|
52 |
+
else:
|
53 |
+
instance_class = instance.pred_classes.tolist()[0]
|
54 |
+
mesh_name = self.class_to_mesh_name[instance_class]
|
55 |
+
|
56 |
+
annotation = {
|
57 |
+
DensePoseDataRelative.X_KEY: [],
|
58 |
+
DensePoseDataRelative.Y_KEY: [],
|
59 |
+
DensePoseDataRelative.VERTEX_IDS_KEY: [],
|
60 |
+
DensePoseDataRelative.MESH_NAME_KEY: mesh_name,
|
61 |
+
}
|
62 |
+
|
63 |
+
mask, embeddings, other_values = self._produce_mask_and_results(instance, bbox_xywh)
|
64 |
+
indices = torch.nonzero(mask, as_tuple=True)
|
65 |
+
selected_embeddings = embeddings.permute(1, 2, 0)[indices].cpu()
|
66 |
+
values = other_values[:, indices[0], indices[1]]
|
67 |
+
k = values.shape[1]
|
68 |
+
|
69 |
+
count = min(self.count_per_class, k)
|
70 |
+
if count <= 0:
|
71 |
+
return annotation
|
72 |
+
|
73 |
+
index_sample = self._produce_index_sample(values, count)
|
74 |
+
closest_vertices = squared_euclidean_distance_matrix(
|
75 |
+
selected_embeddings[index_sample], self.embedder(mesh_name)
|
76 |
+
)
|
77 |
+
closest_vertices = torch.argmin(closest_vertices, dim=1)
|
78 |
+
|
79 |
+
sampled_y = indices[0][index_sample] + 0.5
|
80 |
+
sampled_x = indices[1][index_sample] + 0.5
|
81 |
+
# prepare / normalize data
|
82 |
+
_, _, w, h = bbox_xywh
|
83 |
+
x = (sampled_x / w * 256.0).cpu().tolist()
|
84 |
+
y = (sampled_y / h * 256.0).cpu().tolist()
|
85 |
+
# extend annotations
|
86 |
+
annotation[DensePoseDataRelative.X_KEY].extend(x)
|
87 |
+
annotation[DensePoseDataRelative.Y_KEY].extend(y)
|
88 |
+
annotation[DensePoseDataRelative.VERTEX_IDS_KEY].extend(closest_vertices.cpu().tolist())
|
89 |
+
return annotation
|
90 |
+
|
91 |
+
def _produce_mask_and_results(
|
92 |
+
self, instance: Instances, bbox_xywh: IntTupleBox
|
93 |
+
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
94 |
+
"""
|
95 |
+
Method to get labels and DensePose results from an instance
|
96 |
+
|
97 |
+
Args:
|
98 |
+
instance (Instances): an instance of `DensePoseEmbeddingPredictorOutput`
|
99 |
+
bbox_xywh (IntTupleBox): the corresponding bounding box
|
100 |
+
|
101 |
+
Return:
|
102 |
+
mask (torch.Tensor): shape [H, W], DensePose segmentation mask
|
103 |
+
embeddings (Tuple[torch.Tensor]): a tensor of shape [D, H, W],
|
104 |
+
DensePose CSE Embeddings
|
105 |
+
other_values (Tuple[torch.Tensor]): a tensor of shape [0, H, W],
|
106 |
+
for potential other values
|
107 |
+
"""
|
108 |
+
densepose_output = instance.pred_densepose
|
109 |
+
S = densepose_output.coarse_segm
|
110 |
+
E = densepose_output.embedding
|
111 |
+
_, _, w, h = bbox_xywh
|
112 |
+
embeddings = F.interpolate(E, size=(h, w), mode="bilinear")[0]
|
113 |
+
coarse_segm_resized = F.interpolate(S, size=(h, w), mode="bilinear")[0]
|
114 |
+
mask = coarse_segm_resized.argmax(0) > 0
|
115 |
+
other_values = torch.empty((0, h, w), device=E.device)
|
116 |
+
return mask, embeddings, other_values
|
117 |
+
|
118 |
+
def _resample_mask(self, output: Any) -> torch.Tensor:
|
119 |
+
"""
|
120 |
+
Convert DensePose predictor output to segmentation annotation - tensors of size
|
121 |
+
(256, 256) and type `int64`.
|
122 |
+
|
123 |
+
Args:
|
124 |
+
output: DensePose predictor output with the following attributes:
|
125 |
+
- coarse_segm: tensor of size [N, D, H, W] with unnormalized coarse
|
126 |
+
segmentation scores
|
127 |
+
Return:
|
128 |
+
Tensor of size (S, S) and type `int64` with coarse segmentation annotations,
|
129 |
+
where S = DensePoseDataRelative.MASK_SIZE
|
130 |
+
"""
|
131 |
+
sz = DensePoseDataRelative.MASK_SIZE
|
132 |
+
mask = (
|
133 |
+
F.interpolate(output.coarse_segm, (sz, sz), mode="bilinear", align_corners=False)
|
134 |
+
.argmax(dim=1)
|
135 |
+
.long()
|
136 |
+
.squeeze()
|
137 |
+
.cpu()
|
138 |
+
)
|
139 |
+
return mask
|
3rdparty/densepose/data/samplers/densepose_cse_confidence_based.py
ADDED
@@ -0,0 +1,119 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
|
3 |
+
import random
|
4 |
+
from typing import Optional, Tuple
|
5 |
+
import torch
|
6 |
+
from torch.nn import functional as F
|
7 |
+
|
8 |
+
from detectron2.config import CfgNode
|
9 |
+
from detectron2.structures import Instances
|
10 |
+
|
11 |
+
from densepose.converters.base import IntTupleBox
|
12 |
+
|
13 |
+
from .densepose_cse_base import DensePoseCSEBaseSampler
|
14 |
+
|
15 |
+
|
16 |
+
class DensePoseCSEConfidenceBasedSampler(DensePoseCSEBaseSampler):
|
17 |
+
"""
|
18 |
+
Samples DensePose data from DensePose predictions.
|
19 |
+
Samples for each class are drawn using confidence value estimates.
|
20 |
+
"""
|
21 |
+
|
22 |
+
def __init__(
|
23 |
+
self,
|
24 |
+
cfg: CfgNode,
|
25 |
+
use_gt_categories: bool,
|
26 |
+
embedder: torch.nn.Module,
|
27 |
+
confidence_channel: str,
|
28 |
+
count_per_class: int = 8,
|
29 |
+
search_count_multiplier: Optional[float] = None,
|
30 |
+
search_proportion: Optional[float] = None,
|
31 |
+
):
|
32 |
+
"""
|
33 |
+
Constructor
|
34 |
+
|
35 |
+
Args:
|
36 |
+
cfg (CfgNode): the config of the model
|
37 |
+
embedder (torch.nn.Module): necessary to compute mesh vertex embeddings
|
38 |
+
confidence_channel (str): confidence channel to use for sampling;
|
39 |
+
possible values:
|
40 |
+
"coarse_segm_confidence": confidences for coarse segmentation
|
41 |
+
(default: "coarse_segm_confidence")
|
42 |
+
count_per_class (int): the sampler produces at most `count_per_class`
|
43 |
+
samples for each category (default: 8)
|
44 |
+
search_count_multiplier (float or None): if not None, the total number
|
45 |
+
of the most confident estimates of a given class to consider is
|
46 |
+
defined as `min(search_count_multiplier * count_per_class, N)`,
|
47 |
+
where `N` is the total number of estimates of the class; cannot be
|
48 |
+
specified together with `search_proportion` (default: None)
|
49 |
+
search_proportion (float or None): if not None, the total number of the
|
50 |
+
of the most confident estimates of a given class to consider is
|
51 |
+
defined as `min(max(search_proportion * N, count_per_class), N)`,
|
52 |
+
where `N` is the total number of estimates of the class; cannot be
|
53 |
+
specified together with `search_count_multiplier` (default: None)
|
54 |
+
"""
|
55 |
+
super().__init__(cfg, use_gt_categories, embedder, count_per_class)
|
56 |
+
self.confidence_channel = confidence_channel
|
57 |
+
self.search_count_multiplier = search_count_multiplier
|
58 |
+
self.search_proportion = search_proportion
|
59 |
+
assert (search_count_multiplier is None) or (search_proportion is None), (
|
60 |
+
f"Cannot specify both search_count_multiplier (={search_count_multiplier})"
|
61 |
+
f"and search_proportion (={search_proportion})"
|
62 |
+
)
|
63 |
+
|
64 |
+
def _produce_index_sample(self, values: torch.Tensor, count: int):
|
65 |
+
"""
|
66 |
+
Produce a sample of indices to select data based on confidences
|
67 |
+
|
68 |
+
Args:
|
69 |
+
values (torch.Tensor): a tensor of length k that contains confidences
|
70 |
+
k: number of points labeled with part_id
|
71 |
+
count (int): number of samples to produce, should be positive and <= k
|
72 |
+
|
73 |
+
Return:
|
74 |
+
list(int): indices of values (along axis 1) selected as a sample
|
75 |
+
"""
|
76 |
+
k = values.shape[1]
|
77 |
+
if k == count:
|
78 |
+
index_sample = list(range(k))
|
79 |
+
else:
|
80 |
+
# take the best count * search_count_multiplier pixels,
|
81 |
+
# sample from them uniformly
|
82 |
+
# (here best = smallest variance)
|
83 |
+
_, sorted_confidence_indices = torch.sort(values[0])
|
84 |
+
if self.search_count_multiplier is not None:
|
85 |
+
search_count = min(int(count * self.search_count_multiplier), k)
|
86 |
+
elif self.search_proportion is not None:
|
87 |
+
search_count = min(max(int(k * self.search_proportion), count), k)
|
88 |
+
else:
|
89 |
+
search_count = min(count, k)
|
90 |
+
sample_from_top = random.sample(range(search_count), count)
|
91 |
+
index_sample = sorted_confidence_indices[-search_count:][sample_from_top]
|
92 |
+
return index_sample
|
93 |
+
|
94 |
+
def _produce_mask_and_results(
|
95 |
+
self, instance: Instances, bbox_xywh: IntTupleBox
|
96 |
+
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
|
97 |
+
"""
|
98 |
+
Method to get labels and DensePose results from an instance
|
99 |
+
|
100 |
+
Args:
|
101 |
+
instance (Instances): an instance of
|
102 |
+
`DensePoseEmbeddingPredictorOutputWithConfidences`
|
103 |
+
bbox_xywh (IntTupleBox): the corresponding bounding box
|
104 |
+
|
105 |
+
Return:
|
106 |
+
mask (torch.Tensor): shape [H, W], DensePose segmentation mask
|
107 |
+
embeddings (Tuple[torch.Tensor]): a tensor of shape [D, H, W]
|
108 |
+
DensePose CSE Embeddings
|
109 |
+
other_values: a tensor of shape [1, H, W], DensePose CSE confidence
|
110 |
+
"""
|
111 |
+
_, _, w, h = bbox_xywh
|
112 |
+
densepose_output = instance.pred_densepose
|
113 |
+
mask, embeddings, _ = super()._produce_mask_and_results(instance, bbox_xywh)
|
114 |
+
other_values = F.interpolate(
|
115 |
+
getattr(densepose_output, self.confidence_channel),
|
116 |
+
size=(h, w),
|
117 |
+
mode="bilinear",
|
118 |
+
)[0].cpu()
|
119 |
+
return mask, embeddings, other_values
|
3rdparty/densepose/data/samplers/densepose_cse_uniform.py
ADDED
@@ -0,0 +1,12 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
|
3 |
+
from .densepose_cse_base import DensePoseCSEBaseSampler
|
4 |
+
from .densepose_uniform import DensePoseUniformSampler
|
5 |
+
|
6 |
+
|
7 |
+
class DensePoseCSEUniformSampler(DensePoseCSEBaseSampler, DensePoseUniformSampler):
|
8 |
+
"""
|
9 |
+
Uniform Sampler for CSE
|
10 |
+
"""
|
11 |
+
|
12 |
+
pass
|
3rdparty/densepose/data/samplers/densepose_uniform.py
ADDED
@@ -0,0 +1,41 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
|
3 |
+
import random
|
4 |
+
import torch
|
5 |
+
|
6 |
+
from .densepose_base import DensePoseBaseSampler
|
7 |
+
|
8 |
+
|
9 |
+
class DensePoseUniformSampler(DensePoseBaseSampler):
|
10 |
+
"""
|
11 |
+
Samples DensePose data from DensePose predictions.
|
12 |
+
Samples for each class are drawn uniformly over all pixels estimated
|
13 |
+
to belong to that class.
|
14 |
+
"""
|
15 |
+
|
16 |
+
def __init__(self, count_per_class: int = 8):
|
17 |
+
"""
|
18 |
+
Constructor
|
19 |
+
|
20 |
+
Args:
|
21 |
+
count_per_class (int): the sampler produces at most `count_per_class`
|
22 |
+
samples for each category
|
23 |
+
"""
|
24 |
+
super().__init__(count_per_class)
|
25 |
+
|
26 |
+
def _produce_index_sample(self, values: torch.Tensor, count: int):
|
27 |
+
"""
|
28 |
+
Produce a uniform sample of indices to select data
|
29 |
+
|
30 |
+
Args:
|
31 |
+
values (torch.Tensor): an array of size [n, k] that contains
|
32 |
+
estimated values (U, V, confidences);
|
33 |
+
n: number of channels (U, V, confidences)
|
34 |
+
k: number of points labeled with part_id
|
35 |
+
count (int): number of samples to produce, should be positive and <= k
|
36 |
+
|
37 |
+
Return:
|
38 |
+
list(int): indices of values (along axis 1) selected as a sample
|
39 |
+
"""
|
40 |
+
k = values.shape[1]
|
41 |
+
return random.sample(range(k), count)
|
3rdparty/densepose/data/samplers/mask_from_densepose.py
ADDED
@@ -0,0 +1,28 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
|
3 |
+
from detectron2.structures import BitMasks, Instances
|
4 |
+
|
5 |
+
from densepose.converters import ToMaskConverter
|
6 |
+
|
7 |
+
|
8 |
+
class MaskFromDensePoseSampler:
|
9 |
+
"""
|
10 |
+
Produce mask GT from DensePose predictions
|
11 |
+
This sampler simply converts DensePose predictions to BitMasks
|
12 |
+
that a contain a bool tensor of the size of the input image
|
13 |
+
"""
|
14 |
+
|
15 |
+
def __call__(self, instances: Instances) -> BitMasks:
|
16 |
+
"""
|
17 |
+
Converts predicted data from `instances` into the GT mask data
|
18 |
+
|
19 |
+
Args:
|
20 |
+
instances (Instances): predicted results, expected to have `pred_densepose` field
|
21 |
+
|
22 |
+
Returns:
|
23 |
+
Boolean Tensor of the size of the input image that has non-zero
|
24 |
+
values at pixels that are estimated to belong to the detected object
|
25 |
+
"""
|
26 |
+
return ToMaskConverter.convert(
|
27 |
+
instances.pred_densepose, instances.pred_boxes, instances.image_size
|
28 |
+
)
|
3rdparty/densepose/data/samplers/prediction_to_gt.py
ADDED
@@ -0,0 +1,98 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
|
3 |
+
from dataclasses import dataclass
|
4 |
+
from typing import Any, Callable, Dict, List, Optional
|
5 |
+
|
6 |
+
from detectron2.structures import Instances
|
7 |
+
|
8 |
+
ModelOutput = Dict[str, Any]
|
9 |
+
SampledData = Dict[str, Any]
|
10 |
+
|
11 |
+
|
12 |
+
@dataclass
|
13 |
+
class _Sampler:
|
14 |
+
"""
|
15 |
+
Sampler registry entry that contains:
|
16 |
+
- src (str): source field to sample from (deleted after sampling)
|
17 |
+
- dst (Optional[str]): destination field to sample to, if not None
|
18 |
+
- func (Optional[Callable: Any -> Any]): function that performs sampling,
|
19 |
+
if None, reference copy is performed
|
20 |
+
"""
|
21 |
+
|
22 |
+
src: str
|
23 |
+
dst: Optional[str]
|
24 |
+
func: Optional[Callable[[Any], Any]]
|
25 |
+
|
26 |
+
|
27 |
+
class PredictionToGroundTruthSampler:
|
28 |
+
"""
|
29 |
+
Sampler implementation that converts predictions to GT using registered
|
30 |
+
samplers for different fields of `Instances`.
|
31 |
+
"""
|
32 |
+
|
33 |
+
def __init__(self, dataset_name: str = ""):
|
34 |
+
self.dataset_name = dataset_name
|
35 |
+
self._samplers = {}
|
36 |
+
self.register_sampler("pred_boxes", "gt_boxes", None)
|
37 |
+
self.register_sampler("pred_classes", "gt_classes", None)
|
38 |
+
# delete scores
|
39 |
+
self.register_sampler("scores")
|
40 |
+
|
41 |
+
def __call__(self, model_output: List[ModelOutput]) -> List[SampledData]:
|
42 |
+
"""
|
43 |
+
Transform model output into ground truth data through sampling
|
44 |
+
|
45 |
+
Args:
|
46 |
+
model_output (Dict[str, Any]): model output
|
47 |
+
Returns:
|
48 |
+
Dict[str, Any]: sampled data
|
49 |
+
"""
|
50 |
+
for model_output_i in model_output:
|
51 |
+
instances: Instances = model_output_i["instances"]
|
52 |
+
# transform data in each field
|
53 |
+
for _, sampler in self._samplers.items():
|
54 |
+
if not instances.has(sampler.src) or sampler.dst is None:
|
55 |
+
continue
|
56 |
+
if sampler.func is None:
|
57 |
+
instances.set(sampler.dst, instances.get(sampler.src))
|
58 |
+
else:
|
59 |
+
instances.set(sampler.dst, sampler.func(instances))
|
60 |
+
# delete model output data that was transformed
|
61 |
+
for _, sampler in self._samplers.items():
|
62 |
+
if sampler.src != sampler.dst and instances.has(sampler.src):
|
63 |
+
instances.remove(sampler.src)
|
64 |
+
model_output_i["dataset"] = self.dataset_name
|
65 |
+
return model_output
|
66 |
+
|
67 |
+
def register_sampler(
|
68 |
+
self,
|
69 |
+
prediction_attr: str,
|
70 |
+
gt_attr: Optional[str] = None,
|
71 |
+
func: Optional[Callable[[Any], Any]] = None,
|
72 |
+
):
|
73 |
+
"""
|
74 |
+
Register sampler for a field
|
75 |
+
|
76 |
+
Args:
|
77 |
+
prediction_attr (str): field to replace with a sampled value
|
78 |
+
gt_attr (Optional[str]): field to store the sampled value to, if not None
|
79 |
+
func (Optional[Callable: Any -> Any]): sampler function
|
80 |
+
"""
|
81 |
+
self._samplers[(prediction_attr, gt_attr)] = _Sampler(
|
82 |
+
src=prediction_attr, dst=gt_attr, func=func
|
83 |
+
)
|
84 |
+
|
85 |
+
def remove_sampler(
|
86 |
+
self,
|
87 |
+
prediction_attr: str,
|
88 |
+
gt_attr: Optional[str] = None,
|
89 |
+
):
|
90 |
+
"""
|
91 |
+
Remove sampler for a field
|
92 |
+
|
93 |
+
Args:
|
94 |
+
prediction_attr (str): field to replace with a sampled value
|
95 |
+
gt_attr (Optional[str]): field to store the sampled value to, if not None
|
96 |
+
"""
|
97 |
+
assert (prediction_attr, gt_attr) in self._samplers
|
98 |
+
del self._samplers[(prediction_attr, gt_attr)]
|
3rdparty/densepose/data/transform/__init__.py
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
|
3 |
+
from .image import ImageResizeTransform
|
3rdparty/densepose/data/transform/image.py
ADDED
@@ -0,0 +1,39 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
|
3 |
+
import torch
|
4 |
+
|
5 |
+
|
6 |
+
class ImageResizeTransform:
|
7 |
+
"""
|
8 |
+
Transform that resizes images loaded from a dataset
|
9 |
+
(BGR data in NCHW channel order, typically uint8) to a format ready to be
|
10 |
+
consumed by DensePose training (BGR float32 data in NCHW channel order)
|
11 |
+
"""
|
12 |
+
|
13 |
+
def __init__(self, min_size: int = 800, max_size: int = 1333):
|
14 |
+
self.min_size = min_size
|
15 |
+
self.max_size = max_size
|
16 |
+
|
17 |
+
def __call__(self, images: torch.Tensor) -> torch.Tensor:
|
18 |
+
"""
|
19 |
+
Args:
|
20 |
+
images (torch.Tensor): tensor of size [N, 3, H, W] that contains
|
21 |
+
BGR data (typically in uint8)
|
22 |
+
Returns:
|
23 |
+
images (torch.Tensor): tensor of size [N, 3, H1, W1] where
|
24 |
+
H1 and W1 are chosen to respect the specified min and max sizes
|
25 |
+
and preserve the original aspect ratio, the data channels
|
26 |
+
follow BGR order and the data type is `torch.float32`
|
27 |
+
"""
|
28 |
+
# resize with min size
|
29 |
+
images = images.float()
|
30 |
+
min_size = min(images.shape[-2:])
|
31 |
+
max_size = max(images.shape[-2:])
|
32 |
+
scale = min(self.min_size / min_size, self.max_size / max_size)
|
33 |
+
images = torch.nn.functional.interpolate(
|
34 |
+
images,
|
35 |
+
scale_factor=scale,
|
36 |
+
mode="bilinear",
|
37 |
+
align_corners=False,
|
38 |
+
)
|
39 |
+
return images
|
3rdparty/densepose/data/utils.py
ADDED
@@ -0,0 +1,38 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
|
3 |
+
import os
|
4 |
+
from typing import Dict, Optional
|
5 |
+
|
6 |
+
from detectron2.config import CfgNode
|
7 |
+
|
8 |
+
|
9 |
+
def is_relative_local_path(path: str) -> bool:
|
10 |
+
path_str = os.fsdecode(path)
|
11 |
+
return ("://" not in path_str) and not os.path.isabs(path)
|
12 |
+
|
13 |
+
|
14 |
+
def maybe_prepend_base_path(base_path: Optional[str], path: str):
|
15 |
+
"""
|
16 |
+
Prepends the provided path with a base path prefix if:
|
17 |
+
1) base path is not None;
|
18 |
+
2) path is a local path
|
19 |
+
"""
|
20 |
+
if base_path is None:
|
21 |
+
return path
|
22 |
+
if is_relative_local_path(path):
|
23 |
+
return os.path.join(base_path, path)
|
24 |
+
return path
|
25 |
+
|
26 |
+
|
27 |
+
def get_class_to_mesh_name_mapping(cfg: CfgNode) -> Dict[int, str]:
|
28 |
+
return {
|
29 |
+
int(class_id): mesh_name
|
30 |
+
for class_id, mesh_name in cfg.DATASETS.CLASS_TO_MESH_NAME_MAPPING.items()
|
31 |
+
}
|
32 |
+
|
33 |
+
|
34 |
+
def get_category_to_class_mapping(dataset_cfg: CfgNode) -> Dict[str, int]:
|
35 |
+
return {
|
36 |
+
category: int(class_id)
|
37 |
+
for category, class_id in dataset_cfg.CATEGORY_TO_CLASS_MAPPING.items()
|
38 |
+
}
|
3rdparty/densepose/data/video/__init__.py
ADDED
@@ -0,0 +1,17 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
|
3 |
+
from .frame_selector import (
|
4 |
+
FrameSelectionStrategy,
|
5 |
+
RandomKFramesSelector,
|
6 |
+
FirstKFramesSelector,
|
7 |
+
LastKFramesSelector,
|
8 |
+
FrameTsList,
|
9 |
+
FrameSelector,
|
10 |
+
)
|
11 |
+
|
12 |
+
from .video_keyframe_dataset import (
|
13 |
+
VideoKeyframeDataset,
|
14 |
+
video_list_from_file,
|
15 |
+
list_keyframes,
|
16 |
+
read_keyframes,
|
17 |
+
)
|
3rdparty/densepose/data/video/frame_selector.py
ADDED
@@ -0,0 +1,87 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
|
3 |
+
import random
|
4 |
+
from collections.abc import Callable
|
5 |
+
from enum import Enum
|
6 |
+
from typing import Callable as TCallable
|
7 |
+
from typing import List
|
8 |
+
|
9 |
+
FrameTsList = List[int]
|
10 |
+
FrameSelector = TCallable[[FrameTsList], FrameTsList]
|
11 |
+
|
12 |
+
|
13 |
+
class FrameSelectionStrategy(Enum):
|
14 |
+
"""
|
15 |
+
Frame selection strategy used with videos:
|
16 |
+
- "random_k": select k random frames
|
17 |
+
- "first_k": select k first frames
|
18 |
+
- "last_k": select k last frames
|
19 |
+
- "all": select all frames
|
20 |
+
"""
|
21 |
+
|
22 |
+
# fmt: off
|
23 |
+
RANDOM_K = "random_k"
|
24 |
+
FIRST_K = "first_k"
|
25 |
+
LAST_K = "last_k"
|
26 |
+
ALL = "all"
|
27 |
+
# fmt: on
|
28 |
+
|
29 |
+
|
30 |
+
class RandomKFramesSelector(Callable): # pyre-ignore[39]
|
31 |
+
"""
|
32 |
+
Selector that retains at most `k` random frames
|
33 |
+
"""
|
34 |
+
|
35 |
+
def __init__(self, k: int):
|
36 |
+
self.k = k
|
37 |
+
|
38 |
+
def __call__(self, frame_tss: FrameTsList) -> FrameTsList:
|
39 |
+
"""
|
40 |
+
Select `k` random frames
|
41 |
+
|
42 |
+
Args:
|
43 |
+
frames_tss (List[int]): timestamps of input frames
|
44 |
+
Returns:
|
45 |
+
List[int]: timestamps of selected frames
|
46 |
+
"""
|
47 |
+
return random.sample(frame_tss, min(self.k, len(frame_tss)))
|
48 |
+
|
49 |
+
|
50 |
+
class FirstKFramesSelector(Callable): # pyre-ignore[39]
|
51 |
+
"""
|
52 |
+
Selector that retains at most `k` first frames
|
53 |
+
"""
|
54 |
+
|
55 |
+
def __init__(self, k: int):
|
56 |
+
self.k = k
|
57 |
+
|
58 |
+
def __call__(self, frame_tss: FrameTsList) -> FrameTsList:
|
59 |
+
"""
|
60 |
+
Select `k` first frames
|
61 |
+
|
62 |
+
Args:
|
63 |
+
frames_tss (List[int]): timestamps of input frames
|
64 |
+
Returns:
|
65 |
+
List[int]: timestamps of selected frames
|
66 |
+
"""
|
67 |
+
return frame_tss[: self.k]
|
68 |
+
|
69 |
+
|
70 |
+
class LastKFramesSelector(Callable): # pyre-ignore[39]
|
71 |
+
"""
|
72 |
+
Selector that retains at most `k` last frames from video data
|
73 |
+
"""
|
74 |
+
|
75 |
+
def __init__(self, k: int):
|
76 |
+
self.k = k
|
77 |
+
|
78 |
+
def __call__(self, frame_tss: FrameTsList) -> FrameTsList:
|
79 |
+
"""
|
80 |
+
Select `k` last frames
|
81 |
+
|
82 |
+
Args:
|
83 |
+
frames_tss (List[int]): timestamps of input frames
|
84 |
+
Returns:
|
85 |
+
List[int]: timestamps of selected frames
|
86 |
+
"""
|
87 |
+
return frame_tss[-self.k :]
|
3rdparty/densepose/data/video/video_keyframe_dataset.py
ADDED
@@ -0,0 +1,300 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# -*- coding: utf-8 -*-
|
2 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
3 |
+
|
4 |
+
import csv
|
5 |
+
import logging
|
6 |
+
import numpy as np
|
7 |
+
from typing import Any, Callable, Dict, List, Optional, Union
|
8 |
+
import av
|
9 |
+
import torch
|
10 |
+
from torch.utils.data.dataset import Dataset
|
11 |
+
|
12 |
+
from detectron2.utils.file_io import PathManager
|
13 |
+
|
14 |
+
from ..utils import maybe_prepend_base_path
|
15 |
+
from .frame_selector import FrameSelector, FrameTsList
|
16 |
+
|
17 |
+
FrameList = List[av.frame.Frame] # pyre-ignore[16]
|
18 |
+
FrameTransform = Callable[[torch.Tensor], torch.Tensor]
|
19 |
+
|
20 |
+
|
21 |
+
def list_keyframes(video_fpath: str, video_stream_idx: int = 0) -> FrameTsList:
|
22 |
+
"""
|
23 |
+
Traverses all keyframes of a video file. Returns a list of keyframe
|
24 |
+
timestamps. Timestamps are counts in timebase units.
|
25 |
+
|
26 |
+
Args:
|
27 |
+
video_fpath (str): Video file path
|
28 |
+
video_stream_idx (int): Video stream index (default: 0)
|
29 |
+
Returns:
|
30 |
+
List[int]: list of keyframe timestaps (timestamp is a count in timebase
|
31 |
+
units)
|
32 |
+
"""
|
33 |
+
try:
|
34 |
+
with PathManager.open(video_fpath, "rb") as io:
|
35 |
+
container = av.open(io, mode="r")
|
36 |
+
stream = container.streams.video[video_stream_idx]
|
37 |
+
keyframes = []
|
38 |
+
pts = -1
|
39 |
+
# Note: even though we request forward seeks for keyframes, sometimes
|
40 |
+
# a keyframe in backwards direction is returned. We introduce tolerance
|
41 |
+
# as a max count of ignored backward seeks
|
42 |
+
tolerance_backward_seeks = 2
|
43 |
+
while True:
|
44 |
+
try:
|
45 |
+
container.seek(pts + 1, backward=False, any_frame=False, stream=stream)
|
46 |
+
except av.AVError as e:
|
47 |
+
# the exception occurs when the video length is exceeded,
|
48 |
+
# we then return whatever data we've already collected
|
49 |
+
logger = logging.getLogger(__name__)
|
50 |
+
logger.debug(
|
51 |
+
f"List keyframes: Error seeking video file {video_fpath}, "
|
52 |
+
f"video stream {video_stream_idx}, pts {pts + 1}, AV error: {e}"
|
53 |
+
)
|
54 |
+
return keyframes
|
55 |
+
except OSError as e:
|
56 |
+
logger = logging.getLogger(__name__)
|
57 |
+
logger.warning(
|
58 |
+
f"List keyframes: Error seeking video file {video_fpath}, "
|
59 |
+
f"video stream {video_stream_idx}, pts {pts + 1}, OS error: {e}"
|
60 |
+
)
|
61 |
+
return []
|
62 |
+
packet = next(container.demux(video=video_stream_idx))
|
63 |
+
if packet.pts is not None and packet.pts <= pts:
|
64 |
+
logger = logging.getLogger(__name__)
|
65 |
+
logger.warning(
|
66 |
+
f"Video file {video_fpath}, stream {video_stream_idx}: "
|
67 |
+
f"bad seek for packet {pts + 1} (got packet {packet.pts}), "
|
68 |
+
f"tolerance {tolerance_backward_seeks}."
|
69 |
+
)
|
70 |
+
tolerance_backward_seeks -= 1
|
71 |
+
if tolerance_backward_seeks == 0:
|
72 |
+
return []
|
73 |
+
pts += 1
|
74 |
+
continue
|
75 |
+
tolerance_backward_seeks = 2
|
76 |
+
pts = packet.pts
|
77 |
+
if pts is None:
|
78 |
+
return keyframes
|
79 |
+
if packet.is_keyframe:
|
80 |
+
keyframes.append(pts)
|
81 |
+
return keyframes
|
82 |
+
except OSError as e:
|
83 |
+
logger = logging.getLogger(__name__)
|
84 |
+
logger.warning(
|
85 |
+
f"List keyframes: Error opening video file container {video_fpath}, " f"OS error: {e}"
|
86 |
+
)
|
87 |
+
except RuntimeError as e:
|
88 |
+
logger = logging.getLogger(__name__)
|
89 |
+
logger.warning(
|
90 |
+
f"List keyframes: Error opening video file container {video_fpath}, "
|
91 |
+
f"Runtime error: {e}"
|
92 |
+
)
|
93 |
+
return []
|
94 |
+
|
95 |
+
|
96 |
+
def read_keyframes(
|
97 |
+
video_fpath: str, keyframes: FrameTsList, video_stream_idx: int = 0
|
98 |
+
) -> FrameList: # pyre-ignore[11]
|
99 |
+
"""
|
100 |
+
Reads keyframe data from a video file.
|
101 |
+
|
102 |
+
Args:
|
103 |
+
video_fpath (str): Video file path
|
104 |
+
keyframes (List[int]): List of keyframe timestamps (as counts in
|
105 |
+
timebase units to be used in container seek operations)
|
106 |
+
video_stream_idx (int): Video stream index (default: 0)
|
107 |
+
Returns:
|
108 |
+
List[Frame]: list of frames that correspond to the specified timestamps
|
109 |
+
"""
|
110 |
+
try:
|
111 |
+
with PathManager.open(video_fpath, "rb") as io:
|
112 |
+
container = av.open(io)
|
113 |
+
stream = container.streams.video[video_stream_idx]
|
114 |
+
frames = []
|
115 |
+
for pts in keyframes:
|
116 |
+
try:
|
117 |
+
container.seek(pts, any_frame=False, stream=stream)
|
118 |
+
frame = next(container.decode(video=0))
|
119 |
+
frames.append(frame)
|
120 |
+
except av.AVError as e:
|
121 |
+
logger = logging.getLogger(__name__)
|
122 |
+
logger.warning(
|
123 |
+
f"Read keyframes: Error seeking video file {video_fpath}, "
|
124 |
+
f"video stream {video_stream_idx}, pts {pts}, AV error: {e}"
|
125 |
+
)
|
126 |
+
container.close()
|
127 |
+
return frames
|
128 |
+
except OSError as e:
|
129 |
+
logger = logging.getLogger(__name__)
|
130 |
+
logger.warning(
|
131 |
+
f"Read keyframes: Error seeking video file {video_fpath}, "
|
132 |
+
f"video stream {video_stream_idx}, pts {pts}, OS error: {e}"
|
133 |
+
)
|
134 |
+
container.close()
|
135 |
+
return frames
|
136 |
+
except StopIteration:
|
137 |
+
logger = logging.getLogger(__name__)
|
138 |
+
logger.warning(
|
139 |
+
f"Read keyframes: Error decoding frame from {video_fpath}, "
|
140 |
+
f"video stream {video_stream_idx}, pts {pts}"
|
141 |
+
)
|
142 |
+
container.close()
|
143 |
+
return frames
|
144 |
+
|
145 |
+
container.close()
|
146 |
+
return frames
|
147 |
+
except OSError as e:
|
148 |
+
logger = logging.getLogger(__name__)
|
149 |
+
logger.warning(
|
150 |
+
f"Read keyframes: Error opening video file container {video_fpath}, OS error: {e}"
|
151 |
+
)
|
152 |
+
except RuntimeError as e:
|
153 |
+
logger = logging.getLogger(__name__)
|
154 |
+
logger.warning(
|
155 |
+
f"Read keyframes: Error opening video file container {video_fpath}, Runtime error: {e}"
|
156 |
+
)
|
157 |
+
return []
|
158 |
+
|
159 |
+
|
160 |
+
def video_list_from_file(video_list_fpath: str, base_path: Optional[str] = None):
|
161 |
+
"""
|
162 |
+
Create a list of paths to video files from a text file.
|
163 |
+
|
164 |
+
Args:
|
165 |
+
video_list_fpath (str): path to a plain text file with the list of videos
|
166 |
+
base_path (str): base path for entries from the video list (default: None)
|
167 |
+
"""
|
168 |
+
video_list = []
|
169 |
+
with PathManager.open(video_list_fpath, "r") as io:
|
170 |
+
for line in io:
|
171 |
+
video_list.append(maybe_prepend_base_path(base_path, str(line.strip())))
|
172 |
+
return video_list
|
173 |
+
|
174 |
+
|
175 |
+
def read_keyframe_helper_data(fpath: str):
|
176 |
+
"""
|
177 |
+
Read keyframe data from a file in CSV format: the header should contain
|
178 |
+
"video_id" and "keyframes" fields. Value specifications are:
|
179 |
+
video_id: int
|
180 |
+
keyframes: list(int)
|
181 |
+
Example of contents:
|
182 |
+
video_id,keyframes
|
183 |
+
2,"[1,11,21,31,41,51,61,71,81]"
|
184 |
+
|
185 |
+
Args:
|
186 |
+
fpath (str): File containing keyframe data
|
187 |
+
|
188 |
+
Return:
|
189 |
+
video_id_to_keyframes (dict: int -> list(int)): for a given video ID it
|
190 |
+
contains a list of keyframes for that video
|
191 |
+
"""
|
192 |
+
video_id_to_keyframes = {}
|
193 |
+
try:
|
194 |
+
with PathManager.open(fpath, "r") as io:
|
195 |
+
csv_reader = csv.reader(io)
|
196 |
+
header = next(csv_reader)
|
197 |
+
video_id_idx = header.index("video_id")
|
198 |
+
keyframes_idx = header.index("keyframes")
|
199 |
+
for row in csv_reader:
|
200 |
+
video_id = int(row[video_id_idx])
|
201 |
+
assert (
|
202 |
+
video_id not in video_id_to_keyframes
|
203 |
+
), f"Duplicate keyframes entry for video {fpath}"
|
204 |
+
video_id_to_keyframes[video_id] = (
|
205 |
+
[int(v) for v in row[keyframes_idx][1:-1].split(",")]
|
206 |
+
if len(row[keyframes_idx]) > 2
|
207 |
+
else []
|
208 |
+
)
|
209 |
+
except Exception as e:
|
210 |
+
logger = logging.getLogger(__name__)
|
211 |
+
logger.warning(f"Error reading keyframe helper data from {fpath}: {e}")
|
212 |
+
return video_id_to_keyframes
|
213 |
+
|
214 |
+
|
215 |
+
class VideoKeyframeDataset(Dataset):
|
216 |
+
"""
|
217 |
+
Dataset that provides keyframes for a set of videos.
|
218 |
+
"""
|
219 |
+
|
220 |
+
_EMPTY_FRAMES = torch.empty((0, 3, 1, 1))
|
221 |
+
|
222 |
+
def __init__(
|
223 |
+
self,
|
224 |
+
video_list: List[str],
|
225 |
+
category_list: Union[str, List[str], None] = None,
|
226 |
+
frame_selector: Optional[FrameSelector] = None,
|
227 |
+
transform: Optional[FrameTransform] = None,
|
228 |
+
keyframe_helper_fpath: Optional[str] = None,
|
229 |
+
):
|
230 |
+
"""
|
231 |
+
Dataset constructor
|
232 |
+
|
233 |
+
Args:
|
234 |
+
video_list (List[str]): list of paths to video files
|
235 |
+
category_list (Union[str, List[str], None]): list of animal categories for each
|
236 |
+
video file. If it is a string, or None, this applies to all videos
|
237 |
+
frame_selector (Callable: KeyFrameList -> KeyFrameList):
|
238 |
+
selects keyframes to process, keyframes are given by
|
239 |
+
packet timestamps in timebase counts. If None, all keyframes
|
240 |
+
are selected (default: None)
|
241 |
+
transform (Callable: torch.Tensor -> torch.Tensor):
|
242 |
+
transforms a batch of RGB images (tensors of size [B, 3, H, W]),
|
243 |
+
returns a tensor of the same size. If None, no transform is
|
244 |
+
applied (default: None)
|
245 |
+
|
246 |
+
"""
|
247 |
+
if type(category_list) == list:
|
248 |
+
self.category_list = category_list
|
249 |
+
else:
|
250 |
+
self.category_list = [category_list] * len(video_list)
|
251 |
+
assert len(video_list) == len(
|
252 |
+
self.category_list
|
253 |
+
), "length of video and category lists must be equal"
|
254 |
+
self.video_list = video_list
|
255 |
+
self.frame_selector = frame_selector
|
256 |
+
self.transform = transform
|
257 |
+
self.keyframe_helper_data = (
|
258 |
+
read_keyframe_helper_data(keyframe_helper_fpath)
|
259 |
+
if keyframe_helper_fpath is not None
|
260 |
+
else None
|
261 |
+
)
|
262 |
+
|
263 |
+
def __getitem__(self, idx: int) -> Dict[str, Any]:
|
264 |
+
"""
|
265 |
+
Gets selected keyframes from a given video
|
266 |
+
|
267 |
+
Args:
|
268 |
+
idx (int): video index in the video list file
|
269 |
+
Returns:
|
270 |
+
A dictionary containing two keys:
|
271 |
+
images (torch.Tensor): tensor of size [N, H, W, 3] or of size
|
272 |
+
defined by the transform that contains keyframes data
|
273 |
+
categories (List[str]): categories of the frames
|
274 |
+
"""
|
275 |
+
categories = [self.category_list[idx]]
|
276 |
+
fpath = self.video_list[idx]
|
277 |
+
keyframes = (
|
278 |
+
list_keyframes(fpath)
|
279 |
+
if self.keyframe_helper_data is None or idx not in self.keyframe_helper_data
|
280 |
+
else self.keyframe_helper_data[idx]
|
281 |
+
)
|
282 |
+
transform = self.transform
|
283 |
+
frame_selector = self.frame_selector
|
284 |
+
if not keyframes:
|
285 |
+
return {"images": self._EMPTY_FRAMES, "categories": []}
|
286 |
+
if frame_selector is not None:
|
287 |
+
keyframes = frame_selector(keyframes)
|
288 |
+
frames = read_keyframes(fpath, keyframes)
|
289 |
+
if not frames:
|
290 |
+
return {"images": self._EMPTY_FRAMES, "categories": []}
|
291 |
+
frames = np.stack([frame.to_rgb().to_ndarray() for frame in frames])
|
292 |
+
frames = torch.as_tensor(frames, device=torch.device("cpu"))
|
293 |
+
frames = frames[..., [2, 1, 0]] # RGB -> BGR
|
294 |
+
frames = frames.permute(0, 3, 1, 2).float() # NHWC -> NCHW
|
295 |
+
if transform is not None:
|
296 |
+
frames = transform(frames)
|
297 |
+
return {"images": frames, "categories": categories}
|
298 |
+
|
299 |
+
def __len__(self):
|
300 |
+
return len(self.video_list)
|
3rdparty/densepose/engine/__init__.py
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
|
3 |
+
from .trainer import Trainer
|
3rdparty/densepose/engine/trainer.py
ADDED
@@ -0,0 +1,258 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved
|
2 |
+
|
3 |
+
import logging
|
4 |
+
import os
|
5 |
+
from collections import OrderedDict
|
6 |
+
from typing import List, Optional, Union
|
7 |
+
import torch
|
8 |
+
from torch import nn
|
9 |
+
|
10 |
+
from detectron2.checkpoint import DetectionCheckpointer
|
11 |
+
from detectron2.config import CfgNode
|
12 |
+
from detectron2.engine import DefaultTrainer
|
13 |
+
from detectron2.evaluation import (
|
14 |
+
DatasetEvaluator,
|
15 |
+
DatasetEvaluators,
|
16 |
+
inference_on_dataset,
|
17 |
+
print_csv_format,
|
18 |
+
)
|
19 |
+
from detectron2.solver.build import get_default_optimizer_params, maybe_add_gradient_clipping
|
20 |
+
from detectron2.utils import comm
|
21 |
+
from detectron2.utils.events import EventWriter, get_event_storage
|
22 |
+
|
23 |
+
from densepose import DensePoseDatasetMapperTTA, DensePoseGeneralizedRCNNWithTTA, load_from_cfg
|
24 |
+
from densepose.data import (
|
25 |
+
DatasetMapper,
|
26 |
+
build_combined_loader,
|
27 |
+
build_detection_test_loader,
|
28 |
+
build_detection_train_loader,
|
29 |
+
build_inference_based_loaders,
|
30 |
+
has_inference_based_loaders,
|
31 |
+
)
|
32 |
+
from densepose.evaluation.d2_evaluator_adapter import Detectron2COCOEvaluatorAdapter
|
33 |
+
from densepose.evaluation.evaluator import DensePoseCOCOEvaluator, build_densepose_evaluator_storage
|
34 |
+
from densepose.modeling.cse import Embedder
|
35 |
+
|
36 |
+
|
37 |
+
class SampleCountingLoader:
|
38 |
+
def __init__(self, loader):
|
39 |
+
self.loader = loader
|
40 |
+
|
41 |
+
def __iter__(self):
|
42 |
+
it = iter(self.loader)
|
43 |
+
storage = get_event_storage()
|
44 |
+
while True:
|
45 |
+
try:
|
46 |
+
batch = next(it)
|
47 |
+
num_inst_per_dataset = {}
|
48 |
+
for data in batch:
|
49 |
+
dataset_name = data["dataset"]
|
50 |
+
if dataset_name not in num_inst_per_dataset:
|
51 |
+
num_inst_per_dataset[dataset_name] = 0
|
52 |
+
num_inst = len(data["instances"])
|
53 |
+
num_inst_per_dataset[dataset_name] += num_inst
|
54 |
+
for dataset_name in num_inst_per_dataset:
|
55 |
+
storage.put_scalar(f"batch/{dataset_name}", num_inst_per_dataset[dataset_name])
|
56 |
+
yield batch
|
57 |
+
except StopIteration:
|
58 |
+
break
|
59 |
+
|
60 |
+
|
61 |
+
class SampleCountMetricPrinter(EventWriter):
|
62 |
+
def __init__(self):
|
63 |
+
self.logger = logging.getLogger(__name__)
|
64 |
+
|
65 |
+
def write(self):
|
66 |
+
storage = get_event_storage()
|
67 |
+
batch_stats_strs = []
|
68 |
+
for key, buf in storage.histories().items():
|
69 |
+
if key.startswith("batch/"):
|
70 |
+
batch_stats_strs.append(f"{key} {buf.avg(20)}")
|
71 |
+
self.logger.info(", ".join(batch_stats_strs))
|
72 |
+
|
73 |
+
|
74 |
+
class Trainer(DefaultTrainer):
|
75 |
+
@classmethod
|
76 |
+
def extract_embedder_from_model(cls, model: nn.Module) -> Optional[Embedder]:
|
77 |
+
if isinstance(model, nn.parallel.DistributedDataParallel):
|
78 |
+
model = model.module
|
79 |
+
if hasattr(model, "roi_heads") and hasattr(model.roi_heads, "embedder"):
|
80 |
+
return model.roi_heads.embedder
|
81 |
+
return None
|
82 |
+
|
83 |
+
# TODO: the only reason to copy the base class code here is to pass the embedder from
|
84 |
+
# the model to the evaluator; that should be refactored to avoid unnecessary copy-pasting
|
85 |
+
@classmethod
|
86 |
+
def test(
|
87 |
+
cls,
|
88 |
+
cfg: CfgNode,
|
89 |
+
model: nn.Module,
|
90 |
+
evaluators: Optional[Union[DatasetEvaluator, List[DatasetEvaluator]]] = None,
|
91 |
+
):
|
92 |
+
"""
|
93 |
+
Args:
|
94 |
+
cfg (CfgNode):
|
95 |
+
model (nn.Module):
|
96 |
+
evaluators (DatasetEvaluator, list[DatasetEvaluator] or None): if None, will call
|
97 |
+
:meth:`build_evaluator`. Otherwise, must have the same length as
|
98 |
+
``cfg.DATASETS.TEST``.
|
99 |
+
|
100 |
+
Returns:
|
101 |
+
dict: a dict of result metrics
|
102 |
+
"""
|
103 |
+
logger = logging.getLogger(__name__)
|
104 |
+
if isinstance(evaluators, DatasetEvaluator):
|
105 |
+
evaluators = [evaluators]
|
106 |
+
if evaluators is not None:
|
107 |
+
assert len(cfg.DATASETS.TEST) == len(evaluators), "{} != {}".format(
|
108 |
+
len(cfg.DATASETS.TEST), len(evaluators)
|
109 |
+
)
|
110 |
+
|
111 |
+
results = OrderedDict()
|
112 |
+
for idx, dataset_name in enumerate(cfg.DATASETS.TEST):
|
113 |
+
data_loader = cls.build_test_loader(cfg, dataset_name)
|
114 |
+
# When evaluators are passed in as arguments,
|
115 |
+
# implicitly assume that evaluators can be created before data_loader.
|
116 |
+
if evaluators is not None:
|
117 |
+
evaluator = evaluators[idx]
|
118 |
+
else:
|
119 |
+
try:
|
120 |
+
embedder = cls.extract_embedder_from_model(model)
|
121 |
+
evaluator = cls.build_evaluator(cfg, dataset_name, embedder=embedder)
|
122 |
+
except NotImplementedError:
|
123 |
+
logger.warn(
|
124 |
+
"No evaluator found. Use `DefaultTrainer.test(evaluators=)`, "
|
125 |
+
"or implement its `build_evaluator` method."
|
126 |
+
)
|
127 |
+
results[dataset_name] = {}
|
128 |
+
continue
|
129 |
+
if cfg.DENSEPOSE_EVALUATION.DISTRIBUTED_INFERENCE or comm.is_main_process():
|
130 |
+
results_i = inference_on_dataset(model, data_loader, evaluator)
|
131 |
+
else:
|
132 |
+
results_i = {}
|
133 |
+
results[dataset_name] = results_i
|
134 |
+
if comm.is_main_process():
|
135 |
+
assert isinstance(
|
136 |
+
results_i, dict
|
137 |
+
), "Evaluator must return a dict on the main process. Got {} instead.".format(
|
138 |
+
results_i
|
139 |
+
)
|
140 |
+
logger.info("Evaluation results for {} in csv format:".format(dataset_name))
|
141 |
+
print_csv_format(results_i)
|
142 |
+
|
143 |
+
if len(results) == 1:
|
144 |
+
results = list(results.values())[0]
|
145 |
+
return results
|
146 |
+
|
147 |
+
@classmethod
|
148 |
+
def build_evaluator(
|
149 |
+
cls,
|
150 |
+
cfg: CfgNode,
|
151 |
+
dataset_name: str,
|
152 |
+
output_folder: Optional[str] = None,
|
153 |
+
embedder: Optional[Embedder] = None,
|
154 |
+
) -> DatasetEvaluators:
|
155 |
+
if output_folder is None:
|
156 |
+
output_folder = os.path.join(cfg.OUTPUT_DIR, "inference")
|
157 |
+
evaluators = []
|
158 |
+
distributed = cfg.DENSEPOSE_EVALUATION.DISTRIBUTED_INFERENCE
|
159 |
+
# Note: we currently use COCO evaluator for both COCO and LVIS datasets
|
160 |
+
# to have compatible metrics. LVIS bbox evaluator could also be used
|
161 |
+
# with an adapter to properly handle filtered / mapped categories
|
162 |
+
# evaluator_type = MetadataCatalog.get(dataset_name).evaluator_type
|
163 |
+
# if evaluator_type == "coco":
|
164 |
+
# evaluators.append(COCOEvaluator(dataset_name, output_dir=output_folder))
|
165 |
+
# elif evaluator_type == "lvis":
|
166 |
+
# evaluators.append(LVISEvaluator(dataset_name, output_dir=output_folder))
|
167 |
+
evaluators.append(
|
168 |
+
Detectron2COCOEvaluatorAdapter(
|
169 |
+
dataset_name, output_dir=output_folder, distributed=distributed
|
170 |
+
)
|
171 |
+
)
|
172 |
+
if cfg.MODEL.DENSEPOSE_ON:
|
173 |
+
storage = build_densepose_evaluator_storage(cfg, output_folder)
|
174 |
+
evaluators.append(
|
175 |
+
DensePoseCOCOEvaluator(
|
176 |
+
dataset_name,
|
177 |
+
distributed,
|
178 |
+
output_folder,
|
179 |
+
evaluator_type=cfg.DENSEPOSE_EVALUATION.TYPE,
|
180 |
+
min_iou_threshold=cfg.DENSEPOSE_EVALUATION.MIN_IOU_THRESHOLD,
|
181 |
+
storage=storage,
|
182 |
+
embedder=embedder,
|
183 |
+
should_evaluate_mesh_alignment=cfg.DENSEPOSE_EVALUATION.EVALUATE_MESH_ALIGNMENT,
|
184 |
+
mesh_alignment_mesh_names=cfg.DENSEPOSE_EVALUATION.MESH_ALIGNMENT_MESH_NAMES,
|
185 |
+
)
|
186 |
+
)
|
187 |
+
return DatasetEvaluators(evaluators)
|
188 |
+
|
189 |
+
@classmethod
|
190 |
+
def build_optimizer(cls, cfg: CfgNode, model: nn.Module):
|
191 |
+
params = get_default_optimizer_params(
|
192 |
+
model,
|
193 |
+
base_lr=cfg.SOLVER.BASE_LR,
|
194 |
+
weight_decay_norm=cfg.SOLVER.WEIGHT_DECAY_NORM,
|
195 |
+
bias_lr_factor=cfg.SOLVER.BIAS_LR_FACTOR,
|
196 |
+
weight_decay_bias=cfg.SOLVER.WEIGHT_DECAY_BIAS,
|
197 |
+
overrides={
|
198 |
+
"features": {
|
199 |
+
"lr": cfg.SOLVER.BASE_LR * cfg.MODEL.ROI_DENSEPOSE_HEAD.CSE.FEATURES_LR_FACTOR,
|
200 |
+
},
|
201 |
+
"embeddings": {
|
202 |
+
"lr": cfg.SOLVER.BASE_LR * cfg.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBEDDING_LR_FACTOR,
|
203 |
+
},
|
204 |
+
},
|
205 |
+
)
|
206 |
+
optimizer = torch.optim.SGD(
|
207 |
+
params,
|
208 |
+
cfg.SOLVER.BASE_LR,
|
209 |
+
momentum=cfg.SOLVER.MOMENTUM,
|
210 |
+
nesterov=cfg.SOLVER.NESTEROV,
|
211 |
+
weight_decay=cfg.SOLVER.WEIGHT_DECAY,
|
212 |
+
)
|
213 |
+
# pyre-fixme[6]: For 2nd param expected `Type[Optimizer]` but got `SGD`.
|
214 |
+
return maybe_add_gradient_clipping(cfg, optimizer)
|
215 |
+
|
216 |
+
@classmethod
|
217 |
+
def build_test_loader(cls, cfg: CfgNode, dataset_name):
|
218 |
+
return build_detection_test_loader(cfg, dataset_name, mapper=DatasetMapper(cfg, False))
|
219 |
+
|
220 |
+
@classmethod
|
221 |
+
def build_train_loader(cls, cfg: CfgNode):
|
222 |
+
data_loader = build_detection_train_loader(cfg, mapper=DatasetMapper(cfg, True))
|
223 |
+
if not has_inference_based_loaders(cfg):
|
224 |
+
return data_loader
|
225 |
+
model = cls.build_model(cfg)
|
226 |
+
model.to(cfg.BOOTSTRAP_MODEL.DEVICE)
|
227 |
+
DetectionCheckpointer(model).resume_or_load(cfg.BOOTSTRAP_MODEL.WEIGHTS, resume=False)
|
228 |
+
inference_based_loaders, ratios = build_inference_based_loaders(cfg, model)
|
229 |
+
loaders = [data_loader] + inference_based_loaders
|
230 |
+
ratios = [1.0] + ratios
|
231 |
+
combined_data_loader = build_combined_loader(cfg, loaders, ratios)
|
232 |
+
sample_counting_loader = SampleCountingLoader(combined_data_loader)
|
233 |
+
return sample_counting_loader
|
234 |
+
|
235 |
+
def build_writers(self):
|
236 |
+
writers = super().build_writers()
|
237 |
+
writers.append(SampleCountMetricPrinter())
|
238 |
+
return writers
|
239 |
+
|
240 |
+
@classmethod
|
241 |
+
def test_with_TTA(cls, cfg: CfgNode, model):
|
242 |
+
logger = logging.getLogger("detectron2.trainer")
|
243 |
+
# In the end of training, run an evaluation with TTA
|
244 |
+
# Only support some R-CNN models.
|
245 |
+
logger.info("Running inference with test-time augmentation ...")
|
246 |
+
transform_data = load_from_cfg(cfg)
|
247 |
+
model = DensePoseGeneralizedRCNNWithTTA(
|
248 |
+
cfg, model, transform_data, DensePoseDatasetMapperTTA(cfg)
|
249 |
+
)
|
250 |
+
evaluators = [
|
251 |
+
cls.build_evaluator(
|
252 |
+
cfg, name, output_folder=os.path.join(cfg.OUTPUT_DIR, "inference_TTA")
|
253 |
+
)
|
254 |
+
for name in cfg.DATASETS.TEST
|
255 |
+
]
|
256 |
+
res = cls.test(cfg, model, evaluators) # pyre-ignore[6]
|
257 |
+
res = OrderedDict({k + "_TTA": v for k, v in res.items()})
|
258 |
+
return res
|
3rdparty/densepose/evaluation/__init__.py
ADDED
@@ -0,0 +1,3 @@
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
|
3 |
+
from .evaluator import DensePoseCOCOEvaluator
|
3rdparty/densepose/evaluation/d2_evaluator_adapter.py
ADDED
@@ -0,0 +1,50 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
|
3 |
+
from detectron2.data.catalog import Metadata
|
4 |
+
from detectron2.evaluation import COCOEvaluator
|
5 |
+
|
6 |
+
from densepose.data.datasets.coco import (
|
7 |
+
get_contiguous_id_to_category_id_map,
|
8 |
+
maybe_filter_categories_cocoapi,
|
9 |
+
)
|
10 |
+
|
11 |
+
|
12 |
+
def _maybe_add_iscrowd_annotations(cocoapi) -> None:
|
13 |
+
for ann in cocoapi.dataset["annotations"]:
|
14 |
+
if "iscrowd" not in ann:
|
15 |
+
ann["iscrowd"] = 0
|
16 |
+
|
17 |
+
|
18 |
+
class Detectron2COCOEvaluatorAdapter(COCOEvaluator):
|
19 |
+
def __init__(
|
20 |
+
self,
|
21 |
+
dataset_name,
|
22 |
+
output_dir=None,
|
23 |
+
distributed=True,
|
24 |
+
):
|
25 |
+
super().__init__(dataset_name, output_dir=output_dir, distributed=distributed)
|
26 |
+
maybe_filter_categories_cocoapi(dataset_name, self._coco_api)
|
27 |
+
_maybe_add_iscrowd_annotations(self._coco_api)
|
28 |
+
# substitute category metadata to account for categories
|
29 |
+
# that are mapped to the same contiguous id
|
30 |
+
if hasattr(self._metadata, "thing_dataset_id_to_contiguous_id"):
|
31 |
+
self._maybe_substitute_metadata()
|
32 |
+
|
33 |
+
def _maybe_substitute_metadata(self):
|
34 |
+
cont_id_2_cat_id = get_contiguous_id_to_category_id_map(self._metadata)
|
35 |
+
cat_id_2_cont_id = self._metadata.thing_dataset_id_to_contiguous_id
|
36 |
+
if len(cont_id_2_cat_id) == len(cat_id_2_cont_id):
|
37 |
+
return
|
38 |
+
|
39 |
+
cat_id_2_cont_id_injective = {}
|
40 |
+
for cat_id, cont_id in cat_id_2_cont_id.items():
|
41 |
+
if (cont_id in cont_id_2_cat_id) and (cont_id_2_cat_id[cont_id] == cat_id):
|
42 |
+
cat_id_2_cont_id_injective[cat_id] = cont_id
|
43 |
+
|
44 |
+
metadata_new = Metadata(name=self._metadata.name)
|
45 |
+
for key, value in self._metadata.__dict__.items():
|
46 |
+
if key == "thing_dataset_id_to_contiguous_id":
|
47 |
+
setattr(metadata_new, key, cat_id_2_cont_id_injective)
|
48 |
+
else:
|
49 |
+
setattr(metadata_new, key, value)
|
50 |
+
self._metadata = metadata_new
|
3rdparty/densepose/evaluation/densepose_coco_evaluation.py
ADDED
@@ -0,0 +1,1303 @@
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
1 |
+
# Copyright (c) Facebook, Inc. and its affiliates.
|
2 |
+
# All rights reserved.
|
3 |
+
#
|
4 |
+
# This source code is licensed under the license found in the
|
5 |
+
# LICENSE file in the root directory of this source tree.
|
6 |
+
# This is a modified version of cocoeval.py where we also have the densepose evaluation.
|
7 |
+
|
8 |
+
__author__ = "tsungyi"
|
9 |
+
|
10 |
+
import copy
|
11 |
+
import datetime
|
12 |
+
import logging
|
13 |
+
import numpy as np
|
14 |
+
import pickle
|
15 |
+
import time
|
16 |
+
from collections import defaultdict
|
17 |
+
from enum import Enum
|
18 |
+
from typing import Any, Dict, Tuple
|
19 |
+
import scipy.spatial.distance as ssd
|
20 |
+
import torch
|
21 |
+
import torch.nn.functional as F
|
22 |
+
from pycocotools import mask as maskUtils
|
23 |
+
from scipy.io import loadmat
|
24 |
+
from scipy.ndimage import zoom as spzoom
|
25 |
+
|
26 |
+
from detectron2.utils.file_io import PathManager
|
27 |
+
|
28 |
+
from densepose.converters.chart_output_to_chart_result import resample_uv_tensors_to_bbox
|
29 |
+
from densepose.converters.segm_to_mask import (
|
30 |
+
resample_coarse_segm_tensor_to_bbox,
|
31 |
+
resample_fine_and_coarse_segm_tensors_to_bbox,
|
32 |
+
)
|
33 |
+
from densepose.modeling.cse.utils import squared_euclidean_distance_matrix
|
34 |
+
from densepose.structures import DensePoseDataRelative
|
35 |
+
from densepose.structures.mesh import create_mesh
|
36 |
+
|
37 |
+
logger = logging.getLogger(__name__)
|
38 |
+
|
39 |
+
|
40 |
+
class DensePoseEvalMode(str, Enum):
|
41 |
+
# use both masks and geodesic distances (GPS * IOU) to compute scores
|
42 |
+
GPSM = "gpsm"
|
43 |
+
# use only geodesic distances (GPS) to compute scores
|
44 |
+
GPS = "gps"
|
45 |
+
# use only masks (IOU) to compute scores
|
46 |
+
IOU = "iou"
|
47 |
+
|
48 |
+
|
49 |
+
class DensePoseDataMode(str, Enum):
|
50 |
+
# use estimated IUV data (default mode)
|
51 |
+
IUV_DT = "iuvdt"
|
52 |
+
# use ground truth IUV data
|
53 |
+
IUV_GT = "iuvgt"
|
54 |
+
# use ground truth labels I and set UV to 0
|
55 |
+
I_GT_UV_0 = "igtuv0"
|
56 |
+
# use ground truth labels I and estimated UV coordinates
|
57 |
+
I_GT_UV_DT = "igtuvdt"
|
58 |
+
# use estimated labels I and set UV to 0
|
59 |
+
I_DT_UV_0 = "idtuv0"
|
60 |
+
|
61 |
+
|
62 |
+
class DensePoseCocoEval:
|
63 |
+
# Interface for evaluating detection on the Microsoft COCO dataset.
|
64 |
+
#
|
65 |
+
# The usage for CocoEval is as follows:
|
66 |
+
# cocoGt=..., cocoDt=... # load dataset and results
|
67 |
+
# E = CocoEval(cocoGt,cocoDt); # initialize CocoEval object
|
68 |
+
# E.params.recThrs = ...; # set parameters as desired
|
69 |
+
# E.evaluate(); # run per image evaluation
|
70 |
+
# E.accumulate(); # accumulate per image results
|
71 |
+
# E.summarize(); # display summary metrics of results
|
72 |
+
# For example usage see evalDemo.m and http://mscoco.org/.
|
73 |
+
#
|
74 |
+
# The evaluation parameters are as follows (defaults in brackets):
|
75 |
+
# imgIds - [all] N img ids to use for evaluation
|
76 |
+
# catIds - [all] K cat ids to use for evaluation
|
77 |
+
# iouThrs - [.5:.05:.95] T=10 IoU thresholds for evaluation
|
78 |
+
# recThrs - [0:.01:1] R=101 recall thresholds for evaluation
|
79 |
+
# areaRng - [...] A=4 object area ranges for evaluation
|
80 |
+
# maxDets - [1 10 100] M=3 thresholds on max detections per image
|
81 |
+
# iouType - ['segm'] set iouType to 'segm', 'bbox', 'keypoints' or 'densepose'
|
82 |
+
# iouType replaced the now DEPRECATED useSegm parameter.
|
83 |
+
# useCats - [1] if true use category labels for evaluation
|
84 |
+
# Note: if useCats=0 category labels are ignored as in proposal scoring.
|
85 |
+
# Note: multiple areaRngs [Ax2] and maxDets [Mx1] can be specified.
|
86 |
+
#
|
87 |
+
# evaluate(): evaluates detections on every image and every category and
|
88 |
+
# concats the results into the "evalImgs" with fields:
|
89 |
+
# dtIds - [1xD] id for each of the D detections (dt)
|
90 |
+
# gtIds - [1xG] id for each of the G ground truths (gt)
|
91 |
+
# dtMatches - [TxD] matching gt id at each IoU or 0
|
92 |
+
# gtMatches - [TxG] matching dt id at each IoU or 0
|
93 |
+
# dtScores - [1xD] confidence of each dt
|
94 |
+
# gtIgnore - [1xG] ignore flag for each gt
|
95 |
+
# dtIgnore - [TxD] ignore flag for each dt at each IoU
|
96 |
+
#
|
97 |
+
# accumulate(): accumulates the per-image, per-category evaluation
|
98 |
+
# results in "evalImgs" into the dictionary "eval" with fields:
|
99 |
+
# params - parameters used for evaluation
|
100 |
+
# date - date evaluation was performed
|
101 |
+
# counts - [T,R,K,A,M] parameter dimensions (see above)
|
102 |
+
# precision - [TxRxKxAxM] precision for every evaluation setting
|
103 |
+
# recall - [TxKxAxM] max recall for every evaluation setting
|
104 |
+
# Note: precision and recall==-1 for settings with no gt objects.
|
105 |
+
#
|
106 |
+
# See also coco, mask, pycocoDemo, pycocoEvalDemo
|
107 |
+
#
|
108 |
+
# Microsoft COCO Toolbox. version 2.0
|
109 |
+
# Data, paper, and tutorials available at: http://mscoco.org/
|
110 |
+
# Code written by Piotr Dollar and Tsung-Yi Lin, 2015.
|
111 |
+
# Licensed under the Simplified BSD License [see coco/license.txt]
|
112 |
+
def __init__(
|
113 |
+
self,
|
114 |
+
cocoGt=None,
|
115 |
+
cocoDt=None,
|
116 |
+
iouType: str = "densepose",
|
117 |
+
multi_storage=None,
|
118 |
+
embedder=None,
|
119 |
+
dpEvalMode: DensePoseEvalMode = DensePoseEvalMode.GPS,
|
120 |
+
dpDataMode: DensePoseDataMode = DensePoseDataMode.IUV_DT,
|
121 |
+
):
|
122 |
+
"""
|
123 |
+
Initialize CocoEval using coco APIs for gt and dt
|
124 |
+
:param cocoGt: coco object with ground truth annotations
|
125 |
+
:param cocoDt: coco object with detection results
|
126 |
+
:return: None
|
127 |
+
"""
|
128 |
+
self.cocoGt = cocoGt # ground truth COCO API
|
129 |
+
self.cocoDt = cocoDt # detections COCO API
|
130 |
+
self.multi_storage = multi_storage
|
131 |
+
self.embedder = embedder
|
132 |
+
self._dpEvalMode = dpEvalMode
|
133 |
+
self._dpDataMode = dpDataMode
|
134 |
+
self.evalImgs = defaultdict(list) # per-image per-category eval results [KxAxI]
|
135 |
+
self.eval = {} # accumulated evaluation results
|
136 |
+
self._gts = defaultdict(list) # gt for evaluation
|
137 |
+
self._dts = defaultdict(list) # dt for evaluation
|
138 |
+
self.params = Params(iouType=iouType) # parameters
|
139 |
+
self._paramsEval = {} # parameters for evaluation
|
140 |
+
self.stats = [] # result summarization
|
141 |
+
self.ious = {} # ious between all gts and dts
|
142 |
+
if cocoGt is not None:
|
143 |
+
self.params.imgIds = sorted(cocoGt.getImgIds())
|
144 |
+
self.params.catIds = sorted(cocoGt.getCatIds())
|
145 |
+
self.ignoreThrBB = 0.7
|
146 |
+
self.ignoreThrUV = 0.9
|
147 |
+
|
148 |
+
def _loadGEval(self):
|
149 |
+
smpl_subdiv_fpath = PathManager.get_local_path(
|
150 |
+
"https://dl.fbaipublicfiles.com/densepose/data/SMPL_subdiv.mat"
|
151 |
+
)
|
152 |
+
pdist_transform_fpath = PathManager.get_local_path(
|
153 |
+
"https://dl.fbaipublicfiles.com/densepose/data/SMPL_SUBDIV_TRANSFORM.mat"
|
154 |
+
)
|
155 |
+
pdist_matrix_fpath = PathManager.get_local_path(
|
156 |
+
"https://dl.fbaipublicfiles.com/densepose/data/Pdist_matrix.pkl", timeout_sec=120
|
157 |
+
)
|
158 |
+
SMPL_subdiv = loadmat(smpl_subdiv_fpath)
|
159 |
+
self.PDIST_transform = loadmat(pdist_transform_fpath)
|
160 |
+
self.PDIST_transform = self.PDIST_transform["index"].squeeze()
|
161 |
+
UV = np.array([SMPL_subdiv["U_subdiv"], SMPL_subdiv["V_subdiv"]]).squeeze()
|
162 |
+
ClosestVertInds = np.arange(UV.shape[1]) + 1
|
163 |
+
self.Part_UVs = []
|
164 |
+
self.Part_ClosestVertInds = []
|
165 |
+
for i in np.arange(24):
|
166 |
+
self.Part_UVs.append(UV[:, SMPL_subdiv["Part_ID_subdiv"].squeeze() == (i + 1)])
|
167 |
+
self.Part_ClosestVertInds.append(
|
168 |
+
ClosestVertInds[SMPL_subdiv["Part_ID_subdiv"].squeeze() == (i + 1)]
|
169 |
+
)
|
170 |
+
|
171 |
+
with open(pdist_matrix_fpath, "rb") as hFile:
|
172 |
+
arrays = pickle.load(hFile, encoding="latin1")
|
173 |
+
self.Pdist_matrix = arrays["Pdist_matrix"]
|
174 |
+
self.Part_ids = np.array(SMPL_subdiv["Part_ID_subdiv"].squeeze())
|
175 |
+
# Mean geodesic distances for parts.
|
176 |
+
self.Mean_Distances = np.array([0, 0.351, 0.107, 0.126, 0.237, 0.173, 0.142, 0.128, 0.150])
|
177 |
+
# Coarse Part labels.
|
178 |
+
self.CoarseParts = np.array(
|
179 |
+
[0, 1, 1, 2, 2, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5, 6, 6, 6, 6, 7, 7, 7, 7, 8, 8]
|
180 |
+
)
|
181 |
+
|
182 |
+
def _prepare(self):
|
183 |
+
"""
|
184 |
+
Prepare ._gts and ._dts for evaluation based on params
|
185 |
+
:return: None
|
186 |
+
"""
|
187 |
+
|
188 |
+
def _toMask(anns, coco):
|
189 |
+
# modify ann['segmentation'] by reference
|
190 |
+
for ann in anns:
|
191 |
+
# safeguard for invalid segmentation annotation;
|
192 |
+
# annotations containing empty lists exist in the posetrack
|
193 |
+
# dataset. This is not a correct segmentation annotation
|
194 |
+
# in terms of COCO format; we need to deal with it somehow
|
195 |
+
segm = ann["segmentation"]
|
196 |
+
if type(segm) == list and len(segm) == 0:
|
197 |
+
ann["segmentation"] = None
|
198 |
+
continue
|
199 |
+
rle = coco.annToRLE(ann)
|
200 |
+
ann["segmentation"] = rle
|
201 |
+
|
202 |
+
def _getIgnoreRegion(iid, coco):
|
203 |
+
img = coco.imgs[iid]
|
204 |
+
|
205 |
+
if "ignore_regions_x" not in img.keys():
|
206 |
+
return None
|
207 |
+
|
208 |
+
if len(img["ignore_regions_x"]) == 0:
|
209 |
+
return None
|
210 |
+
|
211 |
+
rgns_merged = [
|
212 |
+
[v for xy in zip(region_x, region_y) for v in xy]
|
213 |
+
for region_x, region_y in zip(img["ignore_regions_x"], img["ignore_regions_y"])
|
214 |
+
]
|
215 |
+
rles = maskUtils.frPyObjects(rgns_merged, img["height"], img["width"])
|
216 |
+
rle = maskUtils.merge(rles)
|
217 |
+
return maskUtils.decode(rle)
|
218 |
+
|
219 |
+
def _checkIgnore(dt, iregion):
|
220 |
+
if iregion is None:
|
221 |
+
return True
|
222 |
+
|
223 |
+
bb = np.array(dt["bbox"]).astype(int)
|
224 |
+
x1, y1, x2, y2 = bb[0], bb[1], bb[0] + bb[2], bb[1] + bb[3]
|
225 |
+
x2 = min([x2, iregion.shape[1]])
|
226 |
+
y2 = min([y2, iregion.shape[0]])
|
227 |
+
|
228 |
+
if bb[2] * bb[3] == 0:
|
229 |
+
return False
|
230 |
+
|
231 |
+
crop_iregion = iregion[y1:y2, x1:x2]
|
232 |
+
|
233 |
+
if crop_iregion.sum() == 0:
|
234 |
+
return True
|
235 |
+
|
236 |
+
if "densepose" not in dt.keys(): # filtering boxes
|
237 |
+
return crop_iregion.sum() / bb[2] / bb[3] < self.ignoreThrBB
|
238 |
+
|
239 |
+
# filtering UVs
|
240 |
+
ignoremask = np.require(crop_iregion, requirements=["F"])
|
241 |
+
mask = self._extract_mask(dt)
|
242 |
+
uvmask = np.require(np.asarray(mask > 0), dtype=np.uint8, requirements=["F"])
|
243 |
+
uvmask_ = maskUtils.encode(uvmask)
|
244 |
+
ignoremask_ = maskUtils.encode(ignoremask)
|
245 |
+
uviou = maskUtils.iou([uvmask_], [ignoremask_], [1])[0]
|
246 |
+
return uviou < self.ignoreThrUV
|
247 |
+
|
248 |
+
p = self.params
|
249 |
+
|
250 |
+
if p.useCats:
|
251 |
+
gts = self.cocoGt.loadAnns(self.cocoGt.getAnnIds(imgIds=p.imgIds, catIds=p.catIds))
|
252 |
+
dts = self.cocoDt.loadAnns(self.cocoDt.getAnnIds(imgIds=p.imgIds, catIds=p.catIds))
|
253 |
+
else:
|
254 |
+
gts = self.cocoGt.loadAnns(self.cocoGt.getAnnIds(imgIds=p.imgIds))
|
255 |
+
dts = self.cocoDt.loadAnns(self.cocoDt.getAnnIds(imgIds=p.imgIds))
|
256 |
+
|
257 |
+
imns = self.cocoGt.loadImgs(p.imgIds)
|
258 |
+
self.size_mapping = {}
|
259 |
+
for im in imns:
|
260 |
+
self.size_mapping[im["id"]] = [im["height"], im["width"]]
|
261 |
+
|
262 |
+
# if iouType == 'uv', add point gt annotations
|
263 |
+
if p.iouType == "densepose":
|
264 |
+
self._loadGEval()
|
265 |
+
|
266 |
+
# convert ground truth to mask if iouType == 'segm'
|
267 |
+
if p.iouType == "segm":
|
268 |
+
_toMask(gts, self.cocoGt)
|
269 |
+
_toMask(dts, self.cocoDt)
|
270 |
+
|
271 |
+
# set ignore flag
|
272 |
+
for gt in gts:
|
273 |
+
gt["ignore"] = gt["ignore"] if "ignore" in gt else 0
|
274 |
+
gt["ignore"] = "iscrowd" in gt and gt["iscrowd"]
|
275 |
+
if p.iouType == "keypoints":
|
276 |
+
gt["ignore"] = (gt["num_keypoints"] == 0) or gt["ignore"]
|
277 |
+
if p.iouType == "densepose":
|
278 |
+
gt["ignore"] = ("dp_x" in gt) == 0
|
279 |
+
if p.iouType == "segm":
|
280 |
+
gt["ignore"] = gt["segmentation"] is None
|
281 |
+
|
282 |
+
self._gts = defaultdict(list) # gt for evaluation
|
283 |
+
self._dts = defaultdict(list) # dt for evaluation
|
284 |
+
self._igrgns = defaultdict(list)
|
285 |
+
|
286 |
+
for gt in gts:
|
287 |
+
iid = gt["image_id"]
|
288 |
+
if iid not in self._igrgns.keys():
|
289 |
+
self._igrgns[iid] = _getIgnoreRegion(iid, self.cocoGt)
|
290 |
+
if _checkIgnore(gt, self._igrgns[iid]):
|
291 |
+
self._gts[iid, gt["category_id"]].append(gt)
|
292 |
+
for dt in dts:
|
293 |
+
iid = dt["image_id"]
|
294 |
+
if (iid not in self._igrgns) or _checkIgnore(dt, self._igrgns[iid]):
|
295 |
+
self._dts[iid, dt["category_id"]].append(dt)
|
296 |
+
|
297 |
+
self.evalImgs = defaultdict(list) # per-image per-category evaluation results
|
298 |
+
self.eval = {} # accumulated evaluation results
|
299 |
+
|
300 |
+
def evaluate(self):
|
301 |
+
"""
|
302 |
+
Run per image evaluation on given images and store results (a list of dict) in self.evalImgs
|
303 |
+
:return: None
|
304 |
+
"""
|
305 |
+
tic = time.time()
|
306 |
+
logger.info("Running per image DensePose evaluation... {}".format(self.params.iouType))
|
307 |
+
p = self.params
|
308 |
+
# add backward compatibility if useSegm is specified in params
|
309 |
+
if p.useSegm is not None:
|
310 |
+
p.iouType = "segm" if p.useSegm == 1 else "bbox"
|
311 |
+
logger.info("useSegm (deprecated) is not None. Running DensePose evaluation")
|
312 |
+
p.imgIds = list(np.unique(p.imgIds))
|
313 |
+
if p.useCats:
|
314 |
+
p.catIds = list(np.unique(p.catIds))
|
315 |
+
p.maxDets = sorted(p.maxDets)
|
316 |
+
self.params = p
|
317 |
+
|
318 |
+
self._prepare()
|
319 |
+
# loop through images, area range, max detection number
|
320 |
+
catIds = p.catIds if p.useCats else [-1]
|
321 |
+
|
322 |
+
if p.iouType in ["segm", "bbox"]:
|
323 |
+
computeIoU = self.computeIoU
|
324 |
+
elif p.iouType == "keypoints":
|
325 |
+
computeIoU = self.computeOks
|
326 |
+
elif p.iouType == "densepose":
|
327 |
+
computeIoU = self.computeOgps
|
328 |
+
if self._dpEvalMode in {DensePoseEvalMode.GPSM, DensePoseEvalMode.IOU}:
|
329 |
+
self.real_ious = {
|
330 |
+
(imgId, catId): self.computeDPIoU(imgId, catId)
|
331 |
+
for imgId in p.imgIds
|
332 |
+
for catId in catIds
|
333 |
+
}
|
334 |
+
|
335 |
+
self.ious = {
|
336 |
+
(imgId, catId): computeIoU(imgId, catId) for imgId in p.imgIds for catId in catIds
|
337 |
+
}
|
338 |
+
|
339 |
+
evaluateImg = self.evaluateImg
|
340 |
+
maxDet = p.maxDets[-1]
|
341 |
+
self.evalImgs = [
|
342 |
+
evaluateImg(imgId, catId, areaRng, maxDet)
|
343 |
+
for catId in catIds
|
344 |
+
for areaRng in p.areaRng
|
345 |
+
for imgId in p.imgIds
|
346 |
+
]
|
347 |
+
self._paramsEval = copy.deepcopy(self.params)
|
348 |
+
toc = time.time()
|
349 |
+
logger.info("DensePose evaluation DONE (t={:0.2f}s).".format(toc - tic))
|
350 |
+
|
351 |
+
def getDensePoseMask(self, polys):
|
352 |
+
maskGen = np.zeros([256, 256])
|
353 |
+
stop = min(len(polys) + 1, 15)
|
354 |
+
for i in range(1, stop):
|
355 |
+
if polys[i - 1]:
|
356 |
+
currentMask = maskUtils.decode(polys[i - 1])
|
357 |
+
maskGen[currentMask > 0] = i
|
358 |
+
return maskGen
|
359 |
+
|
360 |
+
def _generate_rlemask_on_image(self, mask, imgId, data):
|
361 |
+
bbox_xywh = np.array(data["bbox"])
|
362 |
+
x, y, w, h = bbox_xywh
|
363 |
+
im_h, im_w = self.size_mapping[imgId]
|
364 |
+
im_mask = np.zeros((im_h, im_w), dtype=np.uint8)
|
365 |
+
if mask is not None:
|
366 |
+
x0 = max(int(x), 0)
|
367 |
+
x1 = min(int(x + w), im_w, int(x) + mask.shape[1])
|
368 |
+
y0 = max(int(y), 0)
|
369 |
+
y1 = min(int(y + h), im_h, int(y) + mask.shape[0])
|
370 |
+
y = int(y)
|
371 |
+
x = int(x)
|
372 |
+
im_mask[y0:y1, x0:x1] = mask[y0 - y : y1 - y, x0 - x : x1 - x]
|
373 |
+
im_mask = np.require(np.asarray(im_mask > 0), dtype=np.uint8, requirements=["F"])
|
374 |
+
rle_mask = maskUtils.encode(np.array(im_mask[:, :, np.newaxis], order="F"))[0]
|
375 |
+
return rle_mask
|
376 |
+
|
377 |
+
def computeDPIoU(self, imgId, catId):
|
378 |
+
p = self.params
|
379 |
+
if p.useCats:
|
380 |
+
gt = self._gts[imgId, catId]
|
381 |
+
dt = self._dts[imgId, catId]
|
382 |
+
else:
|
383 |
+
gt = [_ for cId in p.catIds for _ in self._gts[imgId, cId]]
|
384 |
+
dt = [_ for cId in p.catIds for _ in self._dts[imgId, cId]]
|
385 |
+
if len(gt) == 0 and len(dt) == 0:
|
386 |
+
return []
|
387 |
+
inds = np.argsort([-d["score"] for d in dt], kind="mergesort")
|
388 |
+
dt = [dt[i] for i in inds]
|
389 |
+
if len(dt) > p.maxDets[-1]:
|
390 |
+
dt = dt[0 : p.maxDets[-1]]
|
391 |
+
|
392 |
+
gtmasks = []
|
393 |
+
for g in gt:
|
394 |
+
if DensePoseDataRelative.S_KEY in g:
|
395 |
+
# convert DensePose mask to a binary mask
|
396 |
+
mask = np.minimum(self.getDensePoseMask(g[DensePoseDataRelative.S_KEY]), 1.0)
|
397 |
+
_, _, w, h = g["bbox"]
|
398 |
+
scale_x = float(max(w, 1)) / mask.shape[1]
|
399 |
+
scale_y = float(max(h, 1)) / mask.shape[0]
|
400 |
+
mask = spzoom(mask, (scale_y, scale_x), order=1, prefilter=False)
|
401 |
+
mask = np.array(mask > 0.5, dtype=np.uint8)
|
402 |
+
rle_mask = self._generate_rlemask_on_image(mask, imgId, g)
|
403 |
+
elif "segmentation" in g:
|
404 |
+
segmentation = g["segmentation"]
|
405 |
+
if isinstance(segmentation, list) and segmentation:
|
406 |
+
# polygons
|
407 |
+
im_h, im_w = self.size_mapping[imgId]
|
408 |
+
rles = maskUtils.frPyObjects(segmentation, im_h, im_w)
|
409 |
+
rle_mask = maskUtils.merge(rles)
|
410 |
+
elif isinstance(segmentation, dict):
|
411 |
+
if isinstance(segmentation["counts"], list):
|
412 |
+
# uncompressed RLE
|
413 |
+
im_h, im_w = self.size_mapping[imgId]
|
414 |
+
rle_mask = maskUtils.frPyObjects(segmentation, im_h, im_w)
|
415 |
+
else:
|
416 |
+
# compressed RLE
|
417 |
+
rle_mask = segmentation
|
418 |
+
else:
|
419 |
+
rle_mask = self._generate_rlemask_on_image(None, imgId, g)
|
420 |
+
else:
|
421 |
+
rle_mask = self._generate_rlemask_on_image(None, imgId, g)
|
422 |
+
gtmasks.append(rle_mask)
|
423 |
+
|
424 |
+
dtmasks = []
|
425 |
+
for d in dt:
|
426 |
+
mask = self._extract_mask(d)
|
427 |
+
mask = np.require(np.asarray(mask > 0), dtype=np.uint8, requirements=["F"])
|
428 |
+
rle_mask = self._generate_rlemask_on_image(mask, imgId, d)
|
429 |
+
dtmasks.append(rle_mask)
|
430 |
+
|
431 |
+
# compute iou between each dt and gt region
|
432 |
+
iscrowd = [int(o.get("iscrowd", 0)) for o in gt]
|
433 |
+
iousDP = maskUtils.iou(dtmasks, gtmasks, iscrowd)
|
434 |
+
return iousDP
|
435 |
+
|
436 |
+
def computeIoU(self, imgId, catId):
|
437 |
+
p = self.params
|
438 |
+
if p.useCats:
|
439 |
+
gt = self._gts[imgId, catId]
|
440 |
+
dt = self._dts[imgId, catId]
|
441 |
+
else:
|
442 |
+
gt = [_ for cId in p.catIds for _ in self._gts[imgId, cId]]
|
443 |
+
dt = [_ for cId in p.catIds for _ in self._dts[imgId, cId]]
|
444 |
+
if len(gt) == 0 and len(dt) == 0:
|
445 |
+
return []
|
446 |
+
inds = np.argsort([-d["score"] for d in dt], kind="mergesort")
|
447 |
+
dt = [dt[i] for i in inds]
|
448 |
+
if len(dt) > p.maxDets[-1]:
|
449 |
+
dt = dt[0 : p.maxDets[-1]]
|
450 |
+
|
451 |
+
if p.iouType == "segm":
|
452 |
+
g = [g["segmentation"] for g in gt if g["segmentation"] is not None]
|
453 |
+
d = [d["segmentation"] for d in dt if d["segmentation"] is not None]
|
454 |
+
elif p.iouType == "bbox":
|
455 |
+
g = [g["bbox"] for g in gt]
|
456 |
+
d = [d["bbox"] for d in dt]
|
457 |
+
else:
|
458 |
+
raise Exception("unknown iouType for iou computation")
|
459 |
+
|
460 |
+
# compute iou between each dt and gt region
|
461 |
+
iscrowd = [int(o.get("iscrowd", 0)) for o in gt]
|
462 |
+
ious = maskUtils.iou(d, g, iscrowd)
|
463 |
+
return ious
|
464 |
+
|
465 |
+
def computeOks(self, imgId, catId):
|
466 |
+
p = self.params
|
467 |
+
# dimension here should be Nxm
|
468 |
+
gts = self._gts[imgId, catId]
|
469 |
+
dts = self._dts[imgId, catId]
|
470 |
+
inds = np.argsort([-d["score"] for d in dts], kind="mergesort")
|
471 |
+
dts = [dts[i] for i in inds]
|
472 |
+
if len(dts) > p.maxDets[-1]:
|
473 |
+
dts = dts[0 : p.maxDets[-1]]
|
474 |
+
# if len(gts) == 0 and len(dts) == 0:
|
475 |
+
if len(gts) == 0 or len(dts) == 0:
|
476 |
+
return []
|
477 |
+
ious = np.zeros((len(dts), len(gts)))
|
478 |
+
sigmas = (
|
479 |
+
np.array(
|
480 |
+
[
|
481 |
+
0.26,
|
482 |
+
0.25,
|
483 |
+
0.25,
|
484 |
+
0.35,
|
485 |
+
0.35,
|
486 |
+
0.79,
|
487 |
+
0.79,
|
488 |
+
0.72,
|
489 |
+
0.72,
|
490 |
+
0.62,
|
491 |
+
0.62,
|
492 |
+
1.07,
|
493 |
+
1.07,
|
494 |
+
0.87,
|
495 |
+
0.87,
|
496 |
+
0.89,
|
497 |
+
0.89,
|
498 |
+
]
|
499 |
+
)
|
500 |
+
/ 10.0
|
501 |
+
)
|
502 |
+
vars = (sigmas * 2) ** 2
|
503 |
+
k = len(sigmas)
|
504 |
+
# compute oks between each detection and ground truth object
|
505 |
+
for j, gt in enumerate(gts):
|
506 |
+
# create bounds for ignore regions(double the gt bbox)
|
507 |
+
g = np.array(gt["keypoints"])
|
508 |
+
xg = g[0::3]
|
509 |
+
yg = g[1::3]
|
510 |
+
vg = g[2::3]
|
511 |
+
k1 = np.count_nonzero(vg > 0)
|
512 |
+
bb = gt["bbox"]
|
513 |
+
x0 = bb[0] - bb[2]
|
514 |
+
x1 = bb[0] + bb[2] * 2
|
515 |
+
y0 = bb[1] - bb[3]
|
516 |
+
y1 = bb[1] + bb[3] * 2
|
517 |
+
for i, dt in enumerate(dts):
|
518 |
+
d = np.array(dt["keypoints"])
|
519 |
+
xd = d[0::3]
|
520 |
+
yd = d[1::3]
|
521 |
+
if k1 > 0:
|
522 |
+
# measure the per-keypoint distance if keypoints visible
|
523 |
+
dx = xd - xg
|
524 |
+
dy = yd - yg
|
525 |
+
else:
|
526 |
+
# measure minimum distance to keypoints in (x0,y0) & (x1,y1)
|
527 |
+
z = np.zeros(k)
|
528 |
+
dx = np.max((z, x0 - xd), axis=0) + np.max((z, xd - x1), axis=0)
|
529 |
+
dy = np.max((z, y0 - yd), axis=0) + np.max((z, yd - y1), axis=0)
|
530 |
+
e = (dx**2 + dy**2) / vars / (gt["area"] + np.spacing(1)) / 2
|
531 |
+
if k1 > 0:
|
532 |
+
e = e[vg > 0]
|
533 |
+
ious[i, j] = np.sum(np.exp(-e)) / e.shape[0]
|
534 |
+
return ious
|
535 |
+
|
536 |
+
def _extract_mask(self, dt: Dict[str, Any]) -> np.ndarray:
|
537 |
+
if "densepose" in dt:
|
538 |
+
densepose_results_quantized = dt["densepose"]
|
539 |
+
return densepose_results_quantized.labels_uv_uint8[0].numpy()
|
540 |
+
elif "cse_mask" in dt:
|
541 |
+
return dt["cse_mask"]
|
542 |
+
elif "coarse_segm" in dt:
|
543 |
+
dy = max(int(dt["bbox"][3]), 1)
|
544 |
+
dx = max(int(dt["bbox"][2]), 1)
|
545 |
+
return (
|
546 |
+
F.interpolate(
|
547 |
+
dt["coarse_segm"].unsqueeze(0),
|
548 |
+
(dy, dx),
|
549 |
+
mode="bilinear",
|
550 |
+
align_corners=False,
|
551 |
+
)
|
552 |
+
.squeeze(0)
|
553 |
+
.argmax(0)
|
554 |
+
.numpy()
|
555 |
+
.astype(np.uint8)
|
556 |
+
)
|
557 |
+
elif "record_id" in dt:
|
558 |
+
assert (
|
559 |
+
self.multi_storage is not None
|
560 |
+
), f"Storage record id encountered in a detection {dt}, but no storage provided!"
|
561 |
+
record = self.multi_storage.get(dt["rank"], dt["record_id"])
|
562 |
+
coarse_segm = record["coarse_segm"]
|
563 |
+
dy = max(int(dt["bbox"][3]), 1)
|
564 |
+
dx = max(int(dt["bbox"][2]), 1)
|
565 |
+
return (
|
566 |
+
F.interpolate(
|
567 |
+
coarse_segm.unsqueeze(0),
|
568 |
+
(dy, dx),
|
569 |
+
mode="bilinear",
|
570 |
+
align_corners=False,
|
571 |
+
)
|
572 |
+
.squeeze(0)
|
573 |
+
.argmax(0)
|
574 |
+
.numpy()
|
575 |
+
.astype(np.uint8)
|
576 |
+
)
|
577 |
+
else:
|
578 |
+
raise Exception(f"No mask data in the detection: {dt}")
|
579 |
+
raise ValueError('The prediction dict needs to contain either "densepose" or "cse_mask"')
|
580 |
+
|
581 |
+
def _extract_iuv(
|
582 |
+
self, densepose_data: np.ndarray, py: np.ndarray, px: np.ndarray, gt: Dict[str, Any]
|
583 |
+
) -> Tuple[np.ndarray, np.ndarray, np.ndarray]:
|
584 |
+
"""
|
585 |
+
Extract arrays of I, U and V values at given points as numpy arrays
|
586 |
+
given the data mode stored in self._dpDataMode
|
587 |
+
"""
|
588 |
+
if self._dpDataMode == DensePoseDataMode.IUV_DT:
|
589 |
+
# estimated labels and UV (default)
|
590 |
+
ipoints = densepose_data[0, py, px]
|
591 |
+
upoints = densepose_data[1, py, px] / 255.0 # convert from uint8 by /255.
|
592 |
+
vpoints = densepose_data[2, py, px] / 255.0
|
593 |
+
elif self._dpDataMode == DensePoseDataMode.IUV_GT:
|
594 |
+
# ground truth
|
595 |
+
ipoints = np.array(gt["dp_I"])
|
596 |
+
upoints = np.array(gt["dp_U"])
|
597 |
+
vpoints = np.array(gt["dp_V"])
|
598 |
+
elif self._dpDataMode == DensePoseDataMode.I_GT_UV_0:
|
599 |
+
# ground truth labels, UV = 0
|
600 |
+
ipoints = np.array(gt["dp_I"])
|
601 |
+
upoints = upoints * 0.0
|
602 |
+
vpoints = vpoints * 0.0
|
603 |
+
elif self._dpDataMode == DensePoseDataMode.I_GT_UV_DT:
|
604 |
+
# ground truth labels, estimated UV
|
605 |
+
ipoints = np.array(gt["dp_I"])
|
606 |
+
upoints = densepose_data[1, py, px] / 255.0 # convert from uint8 by /255.
|
607 |
+
vpoints = densepose_data[2, py, px] / 255.0
|
608 |
+
elif self._dpDataMode == DensePoseDataMode.I_DT_UV_0:
|
609 |
+
# estimated labels, UV = 0
|
610 |
+
ipoints = densepose_data[0, py, px]
|
611 |
+
upoints = upoints * 0.0
|
612 |
+
vpoints = vpoints * 0.0
|
613 |
+
else:
|
614 |
+
raise ValueError(f"Unknown data mode: {self._dpDataMode}")
|
615 |
+
return ipoints, upoints, vpoints
|
616 |
+
|
617 |
+
def computeOgps_single_pair(self, dt, gt, py, px, pt_mask):
|
618 |
+
if "densepose" in dt:
|
619 |
+
ipoints, upoints, vpoints = self.extract_iuv_from_quantized(dt, gt, py, px, pt_mask)
|
620 |
+
return self.computeOgps_single_pair_iuv(dt, gt, ipoints, upoints, vpoints)
|
621 |
+
elif "u" in dt:
|
622 |
+
ipoints, upoints, vpoints = self.extract_iuv_from_raw(dt, gt, py, px, pt_mask)
|
623 |
+
return self.computeOgps_single_pair_iuv(dt, gt, ipoints, upoints, vpoints)
|
624 |
+
elif "record_id" in dt:
|
625 |
+
assert (
|
626 |
+
self.multi_storage is not None
|
627 |
+
), f"Storage record id encountered in detection {dt}, but no storage provided!"
|
628 |
+
record = self.multi_storage.get(dt["rank"], dt["record_id"])
|
629 |
+
record["bbox"] = dt["bbox"]
|
630 |
+
if "u" in record:
|
631 |
+
ipoints, upoints, vpoints = self.extract_iuv_from_raw(record, gt, py, px, pt_mask)
|
632 |
+
return self.computeOgps_single_pair_iuv(dt, gt, ipoints, upoints, vpoints)
|
633 |
+
elif "embedding" in record:
|
634 |
+
return self.computeOgps_single_pair_cse(
|
635 |
+
dt,
|
636 |
+
gt,
|
637 |
+
py,
|
638 |
+
px,
|
639 |
+
pt_mask,
|
640 |
+
record["coarse_segm"],
|
641 |
+
record["embedding"],
|
642 |
+
record["bbox"],
|
643 |
+
)
|
644 |
+
else:
|
645 |
+
raise Exception(f"Unknown record format: {record}")
|
646 |
+
elif "embedding" in dt:
|
647 |
+
return self.computeOgps_single_pair_cse(
|
648 |
+
dt, gt, py, px, pt_mask, dt["coarse_segm"], dt["embedding"], dt["bbox"]
|
649 |
+
)
|
650 |
+
raise Exception(f"Unknown detection format: {dt}")
|
651 |
+
|
652 |
+
def extract_iuv_from_quantized(self, dt, gt, py, px, pt_mask):
|
653 |
+
densepose_results_quantized = dt["densepose"]
|
654 |
+
ipoints, upoints, vpoints = self._extract_iuv(
|
655 |
+
densepose_results_quantized.labels_uv_uint8.numpy(), py, px, gt
|
656 |
+
)
|
657 |
+
ipoints[pt_mask == -1] = 0
|
658 |
+
return ipoints, upoints, vpoints
|
659 |
+
|
660 |
+
def extract_iuv_from_raw(self, dt, gt, py, px, pt_mask):
|
661 |
+
labels_dt = resample_fine_and_coarse_segm_tensors_to_bbox(
|
662 |
+
dt["fine_segm"].unsqueeze(0),
|
663 |
+
dt["coarse_segm"].unsqueeze(0),
|
664 |
+
dt["bbox"],
|
665 |
+
)
|
666 |
+
uv = resample_uv_tensors_to_bbox(
|
667 |
+
dt["u"].unsqueeze(0), dt["v"].unsqueeze(0), labels_dt.squeeze(0), dt["bbox"]
|
668 |
+
)
|
669 |
+
labels_uv_uint8 = torch.cat((labels_dt.byte(), (uv * 255).clamp(0, 255).byte()))
|
670 |
+
ipoints, upoints, vpoints = self._extract_iuv(labels_uv_uint8.numpy(), py, px, gt)
|
671 |
+
ipoints[pt_mask == -1] = 0
|
672 |
+
return ipoints, upoints, vpoints
|
673 |
+
|
674 |
+
def computeOgps_single_pair_iuv(self, dt, gt, ipoints, upoints, vpoints):
|
675 |
+
cVertsGT, ClosestVertsGTTransformed = self.findAllClosestVertsGT(gt)
|
676 |
+
cVerts = self.findAllClosestVertsUV(upoints, vpoints, ipoints)
|
677 |
+
# Get pairwise geodesic distances between gt and estimated mesh points.
|
678 |
+
dist = self.getDistancesUV(ClosestVertsGTTransformed, cVerts)
|
679 |
+
# Compute the Ogps measure.
|
680 |
+
# Find the mean geodesic normalization distance for
|
681 |
+
# each GT point, based on which part it is on.
|
682 |
+
Current_Mean_Distances = self.Mean_Distances[
|
683 |
+
self.CoarseParts[self.Part_ids[cVertsGT[cVertsGT > 0].astype(int) - 1]]
|
684 |
+
]
|
685 |
+
return dist, Current_Mean_Distances
|
686 |
+
|
687 |
+
def computeOgps_single_pair_cse(
|
688 |
+
self, dt, gt, py, px, pt_mask, coarse_segm, embedding, bbox_xywh_abs
|
689 |
+
):
|
690 |
+
# 0-based mesh vertex indices
|
691 |
+
cVertsGT = torch.as_tensor(gt["dp_vertex"], dtype=torch.int64)
|
692 |
+
# label for each pixel of the bbox, [H, W] tensor of long
|
693 |
+
labels_dt = resample_coarse_segm_tensor_to_bbox(
|
694 |
+
coarse_segm.unsqueeze(0), bbox_xywh_abs
|
695 |
+
).squeeze(0)
|
696 |
+
x, y, w, h = bbox_xywh_abs
|
697 |
+
# embedding for each pixel of the bbox, [D, H, W] tensor of float32
|
698 |
+
embedding = F.interpolate(
|
699 |
+
embedding.unsqueeze(0), (int(h), int(w)), mode="bilinear", align_corners=False
|
700 |
+
).squeeze(0)
|
701 |
+
# valid locations py, px
|
702 |
+
py_pt = torch.from_numpy(py[pt_mask > -1])
|
703 |
+
px_pt = torch.from_numpy(px[pt_mask > -1])
|
704 |
+
cVerts = torch.ones_like(cVertsGT) * -1
|
705 |
+
cVerts[pt_mask > -1] = self.findClosestVertsCse(
|
706 |
+
embedding, py_pt, px_pt, labels_dt, gt["ref_model"]
|
707 |
+
)
|
708 |
+
# Get pairwise geodesic distances between gt and estimated mesh points.
|
709 |
+
dist = self.getDistancesCse(cVertsGT, cVerts, gt["ref_model"])
|
710 |
+
# normalize distances
|
711 |
+
if (gt["ref_model"] == "smpl_27554") and ("dp_I" in gt):
|
712 |
+
Current_Mean_Distances = self.Mean_Distances[
|
713 |
+
self.CoarseParts[np.array(gt["dp_I"], dtype=int)]
|
714 |
+
]
|
715 |
+
else:
|
716 |
+
Current_Mean_Distances = 0.255
|
717 |
+
return dist, Current_Mean_Distances
|
718 |
+
|
719 |
+
def computeOgps(self, imgId, catId):
|
720 |
+
p = self.params
|
721 |
+
# dimension here should be Nxm
|
722 |
+
g = self._gts[imgId, catId]
|
723 |
+
d = self._dts[imgId, catId]
|
724 |
+
inds = np.argsort([-d_["score"] for d_ in d], kind="mergesort")
|
725 |
+
d = [d[i] for i in inds]
|
726 |
+
if len(d) > p.maxDets[-1]:
|
727 |
+
d = d[0 : p.maxDets[-1]]
|
728 |
+
# if len(gts) == 0 and len(dts) == 0:
|
729 |
+
if len(g) == 0 or len(d) == 0:
|
730 |
+
return []
|
731 |
+
ious = np.zeros((len(d), len(g)))
|
732 |
+
# compute opgs between each detection and ground truth object
|
733 |
+
# sigma = self.sigma #0.255 # dist = 0.3m corresponds to ogps = 0.5
|
734 |
+
# 1 # dist = 0.3m corresponds to ogps = 0.96
|
735 |
+
# 1.45 # dist = 1.7m (person height) corresponds to ogps = 0.5)
|
736 |
+
for j, gt in enumerate(g):
|
737 |
+
if not gt["ignore"]:
|
738 |
+
g_ = gt["bbox"]
|
739 |
+
for i, dt in enumerate(d):
|
740 |
+
#
|
741 |
+
dy = int(dt["bbox"][3])
|
742 |
+
dx = int(dt["bbox"][2])
|
743 |
+
dp_x = np.array(gt["dp_x"]) * g_[2] / 255.0
|
744 |
+
dp_y = np.array(gt["dp_y"]) * g_[3] / 255.0
|
745 |
+
py = (dp_y + g_[1] - dt["bbox"][1]).astype(int)
|
746 |
+
px = (dp_x + g_[0] - dt["bbox"][0]).astype(int)
|
747 |
+
#
|
748 |
+
pts = np.zeros(len(px))
|
749 |
+
pts[px >= dx] = -1
|
750 |
+
pts[py >= dy] = -1
|
751 |
+
pts[px < 0] = -1
|
752 |
+
pts[py < 0] = -1
|
753 |
+
if len(pts) < 1:
|
754 |
+
ogps = 0.0
|
755 |
+
elif np.max(pts) == -1:
|
756 |
+
ogps = 0.0
|
757 |
+
else:
|
758 |
+
px[pts == -1] = 0
|
759 |
+
py[pts == -1] = 0
|
760 |
+
dists_between_matches, dist_norm_coeffs = self.computeOgps_single_pair(
|
761 |
+
dt, gt, py, px, pts
|
762 |
+
)
|
763 |
+
# Compute gps
|
764 |
+
ogps_values = np.exp(
|
765 |
+
-(dists_between_matches**2) / (2 * (dist_norm_coeffs**2))
|
766 |
+
)
|
767 |
+
#
|
768 |
+
ogps = np.mean(ogps_values) if len(ogps_values) > 0 else 0.0
|
769 |
+
ious[i, j] = ogps
|
770 |
+
|
771 |
+
gbb = [gt["bbox"] for gt in g]
|
772 |
+
dbb = [dt["bbox"] for dt in d]
|
773 |
+
|
774 |
+
# compute iou between each dt and gt region
|
775 |
+
iscrowd = [int(o.get("iscrowd", 0)) for o in g]
|
776 |
+
ious_bb = maskUtils.iou(dbb, gbb, iscrowd)
|
777 |
+
return ious, ious_bb
|
778 |
+
|
779 |
+
def evaluateImg(self, imgId, catId, aRng, maxDet):
|
780 |
+
"""
|
781 |
+
perform evaluation for single category and image
|
782 |
+
:return: dict (single image results)
|
783 |
+
"""
|
784 |
+
|
785 |
+
p = self.params
|
786 |
+
if p.useCats:
|
787 |
+
gt = self._gts[imgId, catId]
|
788 |
+
dt = self._dts[imgId, catId]
|
789 |
+
else:
|
790 |
+
gt = [_ for cId in p.catIds for _ in self._gts[imgId, cId]]
|
791 |
+
dt = [_ for cId in p.catIds for _ in self._dts[imgId, cId]]
|
792 |
+
if len(gt) == 0 and len(dt) == 0:
|
793 |
+
return None
|
794 |
+
|
795 |
+
for g in gt:
|
796 |
+
# g['_ignore'] = g['ignore']
|
797 |
+
if g["ignore"] or (g["area"] < aRng[0] or g["area"] > aRng[1]):
|
798 |
+
g["_ignore"] = True
|
799 |
+
else:
|
800 |
+
g["_ignore"] = False
|
801 |
+
|
802 |
+
# sort dt highest score first, sort gt ignore last
|
803 |
+
gtind = np.argsort([g["_ignore"] for g in gt], kind="mergesort")
|
804 |
+
gt = [gt[i] for i in gtind]
|
805 |
+
dtind = np.argsort([-d["score"] for d in dt], kind="mergesort")
|
806 |
+
dt = [dt[i] for i in dtind[0:maxDet]]
|
807 |
+
iscrowd = [int(o.get("iscrowd", 0)) for o in gt]
|
808 |
+
# load computed ious
|
809 |
+
if p.iouType == "densepose":
|
810 |
+
# print('Checking the length', len(self.ious[imgId, catId]))
|
811 |
+
# if len(self.ious[imgId, catId]) == 0:
|
812 |
+
# print(self.ious[imgId, catId])
|
813 |
+
ious = (
|
814 |
+
self.ious[imgId, catId][0][:, gtind]
|
815 |
+
if len(self.ious[imgId, catId]) > 0
|
816 |
+
else self.ious[imgId, catId]
|
817 |
+
)
|
818 |
+
ioubs = (
|
819 |
+
self.ious[imgId, catId][1][:, gtind]
|
820 |
+
if len(self.ious[imgId, catId]) > 0
|
821 |
+
else self.ious[imgId, catId]
|
822 |
+
)
|
823 |
+
if self._dpEvalMode in {DensePoseEvalMode.GPSM, DensePoseEvalMode.IOU}:
|
824 |
+
iousM = (
|
825 |
+
self.real_ious[imgId, catId][:, gtind]
|
826 |
+
if len(self.real_ious[imgId, catId]) > 0
|
827 |
+
else self.real_ious[imgId, catId]
|
828 |
+
)
|
829 |
+
else:
|
830 |
+
ious = (
|
831 |
+
self.ious[imgId, catId][:, gtind]
|
832 |
+
if len(self.ious[imgId, catId]) > 0
|
833 |
+
else self.ious[imgId, catId]
|
834 |
+
)
|
835 |
+
|
836 |
+
T = len(p.iouThrs)
|
837 |
+
G = len(gt)
|
838 |
+
D = len(dt)
|
839 |
+
gtm = np.zeros((T, G))
|
840 |
+
dtm = np.zeros((T, D))
|
841 |
+
gtIg = np.array([g["_ignore"] for g in gt])
|
842 |
+
dtIg = np.zeros((T, D))
|
843 |
+
if np.all(gtIg) and p.iouType == "densepose":
|
844 |
+
dtIg = np.logical_or(dtIg, True)
|
845 |
+
|
846 |
+
if len(ious) > 0: # and not p.iouType == 'densepose':
|
847 |
+
for tind, t in enumerate(p.iouThrs):
|
848 |
+
for dind, d in enumerate(dt):
|
849 |
+
# information about best match so far (m=-1 -> unmatched)
|
850 |
+
iou = min([t, 1 - 1e-10])
|
851 |
+
m = -1
|
852 |
+
for gind, _g in enumerate(gt):
|
853 |
+
# if this gt already matched, and not a crowd, continue
|
854 |
+
if gtm[tind, gind] > 0 and not iscrowd[gind]:
|
855 |
+
continue
|
856 |
+
# if dt matched to reg gt, and on ignore gt, stop
|
857 |
+
if m > -1 and gtIg[m] == 0 and gtIg[gind] == 1:
|
858 |
+
break
|
859 |
+
if p.iouType == "densepose":
|
860 |
+
if self._dpEvalMode == DensePoseEvalMode.GPSM:
|
861 |
+
new_iou = np.sqrt(iousM[dind, gind] * ious[dind, gind])
|
862 |
+
elif self._dpEvalMode == DensePoseEvalMode.IOU:
|
863 |
+
new_iou = iousM[dind, gind]
|
864 |
+
elif self._dpEvalMode == DensePoseEvalMode.GPS:
|
865 |
+
new_iou = ious[dind, gind]
|
866 |
+
else:
|
867 |
+
new_iou = ious[dind, gind]
|
868 |
+
if new_iou < iou:
|
869 |
+
continue
|
870 |
+
if new_iou == 0.0:
|
871 |
+
continue
|
872 |
+
# if match successful and best so far, store appropriately
|
873 |
+
iou = new_iou
|
874 |
+
m = gind
|
875 |
+
# if match made store id of match for both dt and gt
|
876 |
+
if m == -1:
|
877 |
+
continue
|
878 |
+
dtIg[tind, dind] = gtIg[m]
|
879 |
+
dtm[tind, dind] = gt[m]["id"]
|
880 |
+
gtm[tind, m] = d["id"]
|
881 |
+
|
882 |
+
if p.iouType == "densepose":
|
883 |
+
if not len(ioubs) == 0:
|
884 |
+
for dind, d in enumerate(dt):
|
885 |
+
# information about best match so far (m=-1 -> unmatched)
|
886 |
+
if dtm[tind, dind] == 0:
|
887 |
+
ioub = 0.8
|
888 |
+
m = -1
|
889 |
+
for gind, _g in enumerate(gt):
|
890 |
+
# if this gt already matched, and not a crowd, continue
|
891 |
+
if gtm[tind, gind] > 0 and not iscrowd[gind]:
|
892 |
+
continue
|
893 |
+
# continue to next gt unless better match made
|
894 |
+
if ioubs[dind, gind] < ioub:
|
895 |
+
continue
|
896 |
+
# if match successful and best so far, store appropriately
|
897 |
+
ioub = ioubs[dind, gind]
|
898 |
+
m = gind
|
899 |
+
# if match made store id of match for both dt and gt
|
900 |
+
if m > -1:
|
901 |
+
dtIg[:, dind] = gtIg[m]
|
902 |
+
if gtIg[m]:
|
903 |
+
dtm[tind, dind] = gt[m]["id"]
|
904 |
+
gtm[tind, m] = d["id"]
|
905 |
+
# set unmatched detections outside of area range to ignore
|
906 |
+
a = np.array([d["area"] < aRng[0] or d["area"] > aRng[1] for d in dt]).reshape((1, len(dt)))
|
907 |
+
dtIg = np.logical_or(dtIg, np.logical_and(dtm == 0, np.repeat(a, T, 0)))
|
908 |
+
# store results for given image and category
|
909 |
+
# print('Done with the function', len(self.ious[imgId, catId]))
|
910 |
+
return {
|
911 |
+
"image_id": imgId,
|
912 |
+
"category_id": catId,
|
913 |
+
"aRng": aRng,
|
914 |
+
"maxDet": maxDet,
|
915 |
+
"dtIds": [d["id"] for d in dt],
|
916 |
+
"gtIds": [g["id"] for g in gt],
|
917 |
+
"dtMatches": dtm,
|
918 |
+
"gtMatches": gtm,
|
919 |
+
"dtScores": [d["score"] for d in dt],
|
920 |
+
"gtIgnore": gtIg,
|
921 |
+
"dtIgnore": dtIg,
|
922 |
+
}
|
923 |
+
|
924 |
+
def accumulate(self, p=None):
|
925 |
+
"""
|
926 |
+
Accumulate per image evaluation results and store the result in self.eval
|
927 |
+
:param p: input params for evaluation
|
928 |
+
:return: None
|
929 |
+
"""
|
930 |
+
logger.info("Accumulating evaluation results...")
|
931 |
+
tic = time.time()
|
932 |
+
if not self.evalImgs:
|
933 |
+
logger.info("Please run evaluate() first")
|
934 |
+
# allows input customized parameters
|
935 |
+
if p is None:
|
936 |
+
p = self.params
|
937 |
+
p.catIds = p.catIds if p.useCats == 1 else [-1]
|
938 |
+
T = len(p.iouThrs)
|
939 |
+
R = len(p.recThrs)
|
940 |
+
K = len(p.catIds) if p.useCats else 1
|
941 |
+
A = len(p.areaRng)
|
942 |
+
M = len(p.maxDets)
|
943 |
+
precision = -(np.ones((T, R, K, A, M))) # -1 for the precision of absent categories
|
944 |
+
recall = -(np.ones((T, K, A, M)))
|
945 |
+
|
946 |
+
# create dictionary for future indexing
|
947 |
+
logger.info("Categories: {}".format(p.catIds))
|
948 |
+
_pe = self._paramsEval
|
949 |
+
catIds = _pe.catIds if _pe.useCats else [-1]
|
950 |
+
setK = set(catIds)
|
951 |
+
setA = set(map(tuple, _pe.areaRng))
|
952 |
+
setM = set(_pe.maxDets)
|
953 |
+
setI = set(_pe.imgIds)
|
954 |
+
# get inds to evaluate
|
955 |
+
k_list = [n for n, k in enumerate(p.catIds) if k in setK]
|
956 |
+
m_list = [m for n, m in enumerate(p.maxDets) if m in setM]
|
957 |
+
a_list = [n for n, a in enumerate(map(lambda x: tuple(x), p.areaRng)) if a in setA]
|
958 |
+
i_list = [n for n, i in enumerate(p.imgIds) if i in setI]
|
959 |
+
I0 = len(_pe.imgIds)
|
960 |
+
A0 = len(_pe.areaRng)
|
961 |
+
# retrieve E at each category, area range, and max number of detections
|
962 |
+
for k, k0 in enumerate(k_list):
|
963 |
+
Nk = k0 * A0 * I0
|
964 |
+
for a, a0 in enumerate(a_list):
|
965 |
+
Na = a0 * I0
|
966 |
+
for m, maxDet in enumerate(m_list):
|
967 |
+
E = [self.evalImgs[Nk + Na + i] for i in i_list]
|
968 |
+
E = [e for e in E if e is not None]
|
969 |
+
if len(E) == 0:
|
970 |
+
continue
|
971 |
+
dtScores = np.concatenate([e["dtScores"][0:maxDet] for e in E])
|
972 |
+
|
973 |
+
# different sorting method generates slightly different results.
|
974 |
+
# mergesort is used to be consistent as Matlab implementation.
|
975 |
+
inds = np.argsort(-dtScores, kind="mergesort")
|
976 |
+
|
977 |
+
dtm = np.concatenate([e["dtMatches"][:, 0:maxDet] for e in E], axis=1)[:, inds]
|
978 |
+
dtIg = np.concatenate([e["dtIgnore"][:, 0:maxDet] for e in E], axis=1)[:, inds]
|
979 |
+
gtIg = np.concatenate([e["gtIgnore"] for e in E])
|
980 |
+
npig = np.count_nonzero(gtIg == 0)
|
981 |
+
if npig == 0:
|
982 |
+
continue
|
983 |
+
tps = np.logical_and(dtm, np.logical_not(dtIg))
|
984 |
+
fps = np.logical_and(np.logical_not(dtm), np.logical_not(dtIg))
|
985 |
+
tp_sum = np.cumsum(tps, axis=1).astype(dtype=float)
|
986 |
+
fp_sum = np.cumsum(fps, axis=1).astype(dtype=float)
|
987 |
+
for t, (tp, fp) in enumerate(zip(tp_sum, fp_sum)):
|
988 |
+
tp = np.array(tp)
|
989 |
+
fp = np.array(fp)
|
990 |
+
nd = len(tp)
|
991 |
+
rc = tp / npig
|
992 |
+
pr = tp / (fp + tp + np.spacing(1))
|
993 |
+
q = np.zeros((R,))
|
994 |
+
|
995 |
+
if nd:
|
996 |
+
recall[t, k, a, m] = rc[-1]
|
997 |
+
else:
|
998 |
+
recall[t, k, a, m] = 0
|
999 |
+
|
1000 |
+
# numpy is slow without cython optimization for accessing elements
|
1001 |
+
# use python array gets significant speed improvement
|
1002 |
+
pr = pr.tolist()
|
1003 |
+
q = q.tolist()
|
1004 |
+
|
1005 |
+
for i in range(nd - 1, 0, -1):
|
1006 |
+
if pr[i] > pr[i - 1]:
|
1007 |
+
pr[i - 1] = pr[i]
|
1008 |
+
|
1009 |
+
inds = np.searchsorted(rc, p.recThrs, side="left")
|
1010 |
+
try:
|
1011 |
+
for ri, pi in enumerate(inds):
|
1012 |
+
q[ri] = pr[pi]
|
1013 |
+
except Exception:
|
1014 |
+
pass
|
1015 |
+
precision[t, :, k, a, m] = np.array(q)
|
1016 |
+
logger.info(
|
1017 |
+
"Final: max precision {}, min precision {}".format(np.max(precision), np.min(precision))
|
1018 |
+
)
|
1019 |
+
self.eval = {
|
1020 |
+
"params": p,
|
1021 |
+
"counts": [T, R, K, A, M],
|
1022 |
+
"date": datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S"),
|
1023 |
+
"precision": precision,
|
1024 |
+
"recall": recall,
|
1025 |
+
}
|
1026 |
+
toc = time.time()
|
1027 |
+
logger.info("DONE (t={:0.2f}s).".format(toc - tic))
|
1028 |
+
|
1029 |
+
def summarize(self):
|
1030 |
+
"""
|
1031 |
+
Compute and display summary metrics for evaluation results.
|
1032 |
+
Note this function can *only* be applied on the default parameter setting
|
1033 |
+
"""
|
1034 |
+
|
1035 |
+
def _summarize(ap=1, iouThr=None, areaRng="all", maxDets=100):
|
1036 |
+
p = self.params
|
1037 |
+
iStr = " {:<18} {} @[ {}={:<9} | area={:>6s} | maxDets={:>3d} ] = {:0.3f}"
|
1038 |
+
titleStr = "Average Precision" if ap == 1 else "Average Recall"
|
1039 |
+
typeStr = "(AP)" if ap == 1 else "(AR)"
|
1040 |
+
measure = "IoU"
|
1041 |
+
if self.params.iouType == "keypoints":
|
1042 |
+
measure = "OKS"
|
1043 |
+
elif self.params.iouType == "densepose":
|
1044 |
+
measure = "OGPS"
|
1045 |
+
iouStr = (
|
1046 |
+
"{:0.2f}:{:0.2f}".format(p.iouThrs[0], p.iouThrs[-1])
|
1047 |
+
if iouThr is None
|
1048 |
+
else "{:0.2f}".format(iouThr)
|
1049 |
+
)
|
1050 |
+
|
1051 |
+
aind = [i for i, aRng in enumerate(p.areaRngLbl) if aRng == areaRng]
|
1052 |
+
mind = [i for i, mDet in enumerate(p.maxDets) if mDet == maxDets]
|
1053 |
+
if ap == 1:
|
1054 |
+
# dimension of precision: [TxRxKxAxM]
|
1055 |
+
s = self.eval["precision"]
|
1056 |
+
# IoU
|
1057 |
+
if iouThr is not None:
|
1058 |
+
t = np.where(np.abs(iouThr - p.iouThrs) < 0.001)[0]
|
1059 |
+
s = s[t]
|
1060 |
+
s = s[:, :, :, aind, mind]
|
1061 |
+
else:
|
1062 |
+
# dimension of recall: [TxKxAxM]
|
1063 |
+
s = self.eval["recall"]
|
1064 |
+
if iouThr is not None:
|
1065 |
+
t = np.where(np.abs(iouThr - p.iouThrs) < 0.001)[0]
|
1066 |
+
s = s[t]
|
1067 |
+
s = s[:, :, aind, mind]
|
1068 |
+
if len(s[s > -1]) == 0:
|
1069 |
+
mean_s = -1
|
1070 |
+
else:
|
1071 |
+
mean_s = np.mean(s[s > -1])
|
1072 |
+
logger.info(iStr.format(titleStr, typeStr, measure, iouStr, areaRng, maxDets, mean_s))
|
1073 |
+
return mean_s
|
1074 |
+
|
1075 |
+
def _summarizeDets():
|
1076 |
+
stats = np.zeros((12,))
|
1077 |
+
stats[0] = _summarize(1)
|
1078 |
+
stats[1] = _summarize(1, iouThr=0.5, maxDets=self.params.maxDets[2])
|
1079 |
+
stats[2] = _summarize(1, iouThr=0.75, maxDets=self.params.maxDets[2])
|
1080 |
+
stats[3] = _summarize(1, areaRng="small", maxDets=self.params.maxDets[2])
|
1081 |
+
stats[4] = _summarize(1, areaRng="medium", maxDets=self.params.maxDets[2])
|
1082 |
+
stats[5] = _summarize(1, areaRng="large", maxDets=self.params.maxDets[2])
|
1083 |
+
stats[6] = _summarize(0, maxDets=self.params.maxDets[0])
|
1084 |
+
stats[7] = _summarize(0, maxDets=self.params.maxDets[1])
|
1085 |
+
stats[8] = _summarize(0, maxDets=self.params.maxDets[2])
|
1086 |
+
stats[9] = _summarize(0, areaRng="small", maxDets=self.params.maxDets[2])
|
1087 |
+
stats[10] = _summarize(0, areaRng="medium", maxDets=self.params.maxDets[2])
|
1088 |
+
stats[11] = _summarize(0, areaRng="large", maxDets=self.params.maxDets[2])
|
1089 |
+
return stats
|
1090 |
+
|
1091 |
+
def _summarizeKps():
|
1092 |
+
stats = np.zeros((10,))
|
1093 |
+
stats[0] = _summarize(1, maxDets=20)
|
1094 |
+
stats[1] = _summarize(1, maxDets=20, iouThr=0.5)
|
1095 |
+
stats[2] = _summarize(1, maxDets=20, iouThr=0.75)
|
1096 |
+
stats[3] = _summarize(1, maxDets=20, areaRng="medium")
|
1097 |
+
stats[4] = _summarize(1, maxDets=20, areaRng="large")
|
1098 |
+
stats[5] = _summarize(0, maxDets=20)
|
1099 |
+
stats[6] = _summarize(0, maxDets=20, iouThr=0.5)
|
1100 |
+
stats[7] = _summarize(0, maxDets=20, iouThr=0.75)
|
1101 |
+
stats[8] = _summarize(0, maxDets=20, areaRng="medium")
|
1102 |
+
stats[9] = _summarize(0, maxDets=20, areaRng="large")
|
1103 |
+
return stats
|
1104 |
+
|
1105 |
+
def _summarizeUvs():
|
1106 |
+
stats = [_summarize(1, maxDets=self.params.maxDets[0])]
|
1107 |
+
min_threshold = self.params.iouThrs.min()
|
1108 |
+
if min_threshold <= 0.201:
|
1109 |
+
stats += [_summarize(1, maxDets=self.params.maxDets[0], iouThr=0.2)]
|
1110 |
+
if min_threshold <= 0.301:
|
1111 |
+
stats += [_summarize(1, maxDets=self.params.maxDets[0], iouThr=0.3)]
|
1112 |
+
if min_threshold <= 0.401:
|
1113 |
+
stats += [_summarize(1, maxDets=self.params.maxDets[0], iouThr=0.4)]
|
1114 |
+
stats += [
|
1115 |
+
_summarize(1, maxDets=self.params.maxDets[0], iouThr=0.5),
|
1116 |
+
_summarize(1, maxDets=self.params.maxDets[0], iouThr=0.75),
|
1117 |
+
_summarize(1, maxDets=self.params.maxDets[0], areaRng="medium"),
|
1118 |
+
_summarize(1, maxDets=self.params.maxDets[0], areaRng="large"),
|
1119 |
+
_summarize(0, maxDets=self.params.maxDets[0]),
|
1120 |
+
_summarize(0, maxDets=self.params.maxDets[0], iouThr=0.5),
|
1121 |
+
_summarize(0, maxDets=self.params.maxDets[0], iouThr=0.75),
|
1122 |
+
_summarize(0, maxDets=self.params.maxDets[0], areaRng="medium"),
|
1123 |
+
_summarize(0, maxDets=self.params.maxDets[0], areaRng="large"),
|
1124 |
+
]
|
1125 |
+
return np.array(stats)
|
1126 |
+
|
1127 |
+
def _summarizeUvsOld():
|
1128 |
+
stats = np.zeros((18,))
|
1129 |
+
stats[0] = _summarize(1, maxDets=self.params.maxDets[0])
|
1130 |
+
stats[1] = _summarize(1, maxDets=self.params.maxDets[0], iouThr=0.5)
|
1131 |
+
stats[2] = _summarize(1, maxDets=self.params.maxDets[0], iouThr=0.55)
|
1132 |
+
stats[3] = _summarize(1, maxDets=self.params.maxDets[0], iouThr=0.60)
|
1133 |
+
stats[4] = _summarize(1, maxDets=self.params.maxDets[0], iouThr=0.65)
|
1134 |
+
stats[5] = _summarize(1, maxDets=self.params.maxDets[0], iouThr=0.70)
|
1135 |
+
stats[6] = _summarize(1, maxDets=self.params.maxDets[0], iouThr=0.75)
|
1136 |
+
stats[7] = _summarize(1, maxDets=self.params.maxDets[0], iouThr=0.80)
|
1137 |
+
stats[8] = _summarize(1, maxDets=self.params.maxDets[0], iouThr=0.85)
|
1138 |
+
stats[9] = _summarize(1, maxDets=self.params.maxDets[0], iouThr=0.90)
|
1139 |
+
stats[10] = _summarize(1, maxDets=self.params.maxDets[0], iouThr=0.95)
|
1140 |
+
stats[11] = _summarize(1, maxDets=self.params.maxDets[0], areaRng="medium")
|
1141 |
+
stats[12] = _summarize(1, maxDets=self.params.maxDets[0], areaRng="large")
|
1142 |
+
stats[13] = _summarize(0, maxDets=self.params.maxDets[0])
|
1143 |
+
stats[14] = _summarize(0, maxDets=self.params.maxDets[0], iouThr=0.5)
|
1144 |
+
stats[15] = _summarize(0, maxDets=self.params.maxDets[0], iouThr=0.75)
|
1145 |
+
stats[16] = _summarize(0, maxDets=self.params.maxDets[0], areaRng="medium")
|
1146 |
+
stats[17] = _summarize(0, maxDets=self.params.maxDets[0], areaRng="large")
|
1147 |
+
return stats
|
1148 |
+
|
1149 |
+
if not self.eval:
|
1150 |
+
raise Exception("Please run accumulate() first")
|
1151 |
+
iouType = self.params.iouType
|
1152 |
+
if iouType in ["segm", "bbox"]:
|
1153 |
+
summarize = _summarizeDets
|
1154 |
+
elif iouType in ["keypoints"]:
|
1155 |
+
summarize = _summarizeKps
|
1156 |
+
elif iouType in ["densepose"]:
|
1157 |
+
summarize = _summarizeUvs
|
1158 |
+
self.stats = summarize()
|
1159 |
+
|
1160 |
+
def __str__(self):
|
1161 |
+
self.summarize()
|
1162 |
+
|
1163 |
+
# ================ functions for dense pose ==============================
|
1164 |
+
def findAllClosestVertsUV(self, U_points, V_points, Index_points):
|
1165 |
+
ClosestVerts = np.ones(Index_points.shape) * -1
|
1166 |
+
for i in np.arange(24):
|
1167 |
+
#
|
1168 |
+
if (i + 1) in Index_points:
|
1169 |
+
UVs = np.array(
|
1170 |
+
[U_points[Index_points == (i + 1)], V_points[Index_points == (i + 1)]]
|
1171 |
+
)
|
1172 |
+
Current_Part_UVs = self.Part_UVs[i]
|
1173 |
+
Current_Part_ClosestVertInds = self.Part_ClosestVertInds[i]
|
1174 |
+
D = ssd.cdist(Current_Part_UVs.transpose(), UVs.transpose()).squeeze()
|
1175 |
+
ClosestVerts[Index_points == (i + 1)] = Current_Part_ClosestVertInds[
|
1176 |
+
np.argmin(D, axis=0)
|
1177 |
+
]
|
1178 |
+
ClosestVertsTransformed = self.PDIST_transform[ClosestVerts.astype(int) - 1]
|
1179 |
+
ClosestVertsTransformed[ClosestVerts < 0] = 0
|
1180 |
+
return ClosestVertsTransformed
|
1181 |
+
|
1182 |
+
def findClosestVertsCse(self, embedding, py, px, mask, mesh_name):
|
1183 |
+
mesh_vertex_embeddings = self.embedder(mesh_name)
|
1184 |
+
pixel_embeddings = embedding[:, py, px].t().to(device="cuda")
|
1185 |
+
mask_vals = mask[py, px]
|
1186 |
+
edm = squared_euclidean_distance_matrix(pixel_embeddings, mesh_vertex_embeddings)
|
1187 |
+
vertex_indices = edm.argmin(dim=1).cpu()
|
1188 |
+
vertex_indices[mask_vals <= 0] = -1
|
1189 |
+
return vertex_indices
|
1190 |
+
|
1191 |
+
def findAllClosestVertsGT(self, gt):
|
1192 |
+
#
|
1193 |
+
I_gt = np.array(gt["dp_I"])
|
1194 |
+
U_gt = np.array(gt["dp_U"])
|
1195 |
+
V_gt = np.array(gt["dp_V"])
|
1196 |
+
#
|
1197 |
+
# print(I_gt)
|
1198 |
+
#
|
1199 |
+
ClosestVertsGT = np.ones(I_gt.shape) * -1
|
1200 |
+
for i in np.arange(24):
|
1201 |
+
if (i + 1) in I_gt:
|
1202 |
+
UVs = np.array([U_gt[I_gt == (i + 1)], V_gt[I_gt == (i + 1)]])
|
1203 |
+
Current_Part_UVs = self.Part_UVs[i]
|
1204 |
+
Current_Part_ClosestVertInds = self.Part_ClosestVertInds[i]
|
1205 |
+
D = ssd.cdist(Current_Part_UVs.transpose(), UVs.transpose()).squeeze()
|
1206 |
+
ClosestVertsGT[I_gt == (i + 1)] = Current_Part_ClosestVertInds[np.argmin(D, axis=0)]
|
1207 |
+
#
|
1208 |
+
ClosestVertsGTTransformed = self.PDIST_transform[ClosestVertsGT.astype(int) - 1]
|
1209 |
+
ClosestVertsGTTransformed[ClosestVertsGT < 0] = 0
|
1210 |
+
return ClosestVertsGT, ClosestVertsGTTransformed
|
1211 |
+
|
1212 |
+
def getDistancesCse(self, cVertsGT, cVerts, mesh_name):
|
1213 |
+
geodists_vertices = torch.ones_like(cVertsGT) * float("inf")
|
1214 |
+
selected = (cVertsGT >= 0) * (cVerts >= 0)
|
1215 |
+
mesh = create_mesh(mesh_name, "cpu")
|
1216 |
+
geodists_vertices[selected] = mesh.geodists[cVertsGT[selected], cVerts[selected]]
|
1217 |
+
return geodists_vertices.numpy()
|
1218 |
+
|
1219 |
+
def getDistancesUV(self, cVertsGT, cVerts):
|
1220 |
+
#
|
1221 |
+
n = 27554
|
1222 |
+
dists = []
|
1223 |
+
for d in range(len(cVertsGT)):
|
1224 |
+
if cVertsGT[d] > 0:
|
1225 |
+
if cVerts[d] > 0:
|
1226 |
+
i = cVertsGT[d] - 1
|
1227 |
+
j = cVerts[d] - 1
|
1228 |
+
if j == i:
|
1229 |
+
dists.append(0)
|
1230 |
+
elif j > i:
|
1231 |
+
ccc = i
|
1232 |
+
i = j
|
1233 |
+
j = ccc
|
1234 |
+
i = n - i - 1
|
1235 |
+
j = n - j - 1
|
1236 |
+
k = (n * (n - 1) / 2) - (n - i) * ((n - i) - 1) / 2 + j - i - 1
|
1237 |
+
k = (n * n - n) / 2 - k - 1
|
1238 |
+
dists.append(self.Pdist_matrix[int(k)][0])
|
1239 |
+
else:
|
1240 |
+
i = n - i - 1
|
1241 |
+
j = n - j - 1
|
1242 |
+
k = (n * (n - 1) / 2) - (n - i) * ((n - i) - 1) / 2 + j - i - 1
|
1243 |
+
k = (n * n - n) / 2 - k - 1
|
1244 |
+
dists.append(self.Pdist_matrix[int(k)][0])
|
1245 |
+
else:
|
1246 |
+
dists.append(np.inf)
|
1247 |
+
return np.atleast_1d(np.array(dists).squeeze())
|
1248 |
+
|
1249 |
+
|
1250 |
+
class Params:
|
1251 |
+
"""
|
1252 |
+
Params for coco evaluation api
|
1253 |
+
"""
|
1254 |
+
|
1255 |
+
def setDetParams(self):
|
1256 |
+
self.imgIds = []
|
1257 |
+
self.catIds = []
|
1258 |
+
# np.arange causes trouble. the data point on arange is slightly larger than the true value
|
1259 |
+
self.iouThrs = np.linspace(0.5, 0.95, int(np.round((0.95 - 0.5) / 0.05)) + 1, endpoint=True)
|
1260 |
+
self.recThrs = np.linspace(0.0, 1.00, int(np.round((1.00 - 0.0) / 0.01)) + 1, endpoint=True)
|
1261 |
+
self.maxDets = [1, 10, 100]
|
1262 |
+
self.areaRng = [
|
1263 |
+
[0**2, 1e5**2],
|
1264 |
+
[0**2, 32**2],
|
1265 |
+
[32**2, 96**2],
|
1266 |
+
[96**2, 1e5**2],
|
1267 |
+
]
|
1268 |
+
self.areaRngLbl = ["all", "small", "medium", "large"]
|
1269 |
+
self.useCats = 1
|
1270 |
+
|
1271 |
+
def setKpParams(self):
|
1272 |
+
self.imgIds = []
|
1273 |
+
self.catIds = []
|
1274 |
+
# np.arange causes trouble. the data point on arange is slightly larger than the true value
|
1275 |
+
self.iouThrs = np.linspace(0.5, 0.95, np.round((0.95 - 0.5) / 0.05) + 1, endpoint=True)
|
1276 |
+
self.recThrs = np.linspace(0.0, 1.00, np.round((1.00 - 0.0) / 0.01) + 1, endpoint=True)
|
1277 |
+
self.maxDets = [20]
|
1278 |
+
self.areaRng = [[0**2, 1e5**2], [32**2, 96**2], [96**2, 1e5**2]]
|
1279 |
+
self.areaRngLbl = ["all", "medium", "large"]
|
1280 |
+
self.useCats = 1
|
1281 |
+
|
1282 |
+
def setUvParams(self):
|
1283 |
+
self.imgIds = []
|
1284 |
+
self.catIds = []
|
1285 |
+
self.iouThrs = np.linspace(0.5, 0.95, int(np.round((0.95 - 0.5) / 0.05)) + 1, endpoint=True)
|
1286 |
+
self.recThrs = np.linspace(0.0, 1.00, int(np.round((1.00 - 0.0) / 0.01)) + 1, endpoint=True)
|
1287 |
+
self.maxDets = [20]
|
1288 |
+
self.areaRng = [[0**2, 1e5**2], [32**2, 96**2], [96**2, 1e5**2]]
|
1289 |
+
self.areaRngLbl = ["all", "medium", "large"]
|
1290 |
+
self.useCats = 1
|
1291 |
+
|
1292 |
+
def __init__(self, iouType="segm"):
|
1293 |
+
if iouType == "segm" or iouType == "bbox":
|
1294 |
+
self.setDetParams()
|
1295 |
+
elif iouType == "keypoints":
|
1296 |
+
self.setKpParams()
|
1297 |
+
elif iouType == "densepose":
|
1298 |
+
self.setUvParams()
|
1299 |
+
else:
|
1300 |
+
raise Exception("iouType not supported")
|
1301 |
+
self.iouType = iouType
|
1302 |
+
# useSegm is deprecated
|
1303 |
+
self.useSegm = None
|