Update test.py for IoU in native image-space (#1439)
Browse files* Update test.py for IoU in native image-space
* remove redundant
* gn to device
* remove output scale_coords
* --img-size correction
* update
* native-space labels
* pred to predn
* remove clip_coords()
- models/common.py +14 -0
- test.py +11 -13
models/common.py
CHANGED
@@ -148,6 +148,8 @@ class autoShape(nn.Module):
|
|
148 |
batch = range(len(imgs)) # batch size
|
149 |
for i in batch:
|
150 |
imgs[i] = np.array(imgs[i]) # to numpy
|
|
|
|
|
151 |
imgs[i] = imgs[i][:, :, :3] if imgs[i].ndim == 3 else np.tile(imgs[i][:, :, None], 3) # enforce 3ch input
|
152 |
s = imgs[i].shape[:2] # HWC
|
153 |
shape0.append(s) # image shape
|
@@ -184,6 +186,7 @@ class Detections:
|
|
184 |
gn = [torch.Tensor([*[im.shape[i] for i in [1, 0, 1, 0]], 1., 1.]) for im in imgs] # normalization gains
|
185 |
self.xyxyn = [x / g for x, g in zip(self.xyxy, gn)] # xyxy normalized
|
186 |
self.xywhn = [x / g for x, g in zip(self.xywh, gn)] # xywh normalized
|
|
|
187 |
|
188 |
def display(self, pprint=False, show=False, save=False):
|
189 |
colors = color_list()
|
@@ -216,6 +219,17 @@ class Detections:
|
|
216 |
def save(self):
|
217 |
self.display(save=True) # save results
|
218 |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
219 |
|
220 |
class Flatten(nn.Module):
|
221 |
# Use after nn.AdaptiveAvgPool2d(1) to remove last 2 dimensions
|
|
|
148 |
batch = range(len(imgs)) # batch size
|
149 |
for i in batch:
|
150 |
imgs[i] = np.array(imgs[i]) # to numpy
|
151 |
+
if imgs[i].shape[0] < 5: # image in CHW
|
152 |
+
imgs[i] = imgs[i].transpose((1, 2, 0)) # reverse dataloader .transpose(2, 0, 1)
|
153 |
imgs[i] = imgs[i][:, :, :3] if imgs[i].ndim == 3 else np.tile(imgs[i][:, :, None], 3) # enforce 3ch input
|
154 |
s = imgs[i].shape[:2] # HWC
|
155 |
shape0.append(s) # image shape
|
|
|
186 |
gn = [torch.Tensor([*[im.shape[i] for i in [1, 0, 1, 0]], 1., 1.]) for im in imgs] # normalization gains
|
187 |
self.xyxyn = [x / g for x, g in zip(self.xyxy, gn)] # xyxy normalized
|
188 |
self.xywhn = [x / g for x, g in zip(self.xywh, gn)] # xywh normalized
|
189 |
+
self.n = len(self.pred)
|
190 |
|
191 |
def display(self, pprint=False, show=False, save=False):
|
192 |
colors = color_list()
|
|
|
219 |
def save(self):
|
220 |
self.display(save=True) # save results
|
221 |
|
222 |
+
def __len__(self):
|
223 |
+
return self.n
|
224 |
+
|
225 |
+
def tolist(self):
|
226 |
+
# return a list of Detections objects, i.e. 'for result in results.tolist():'
|
227 |
+
x = [Detections([self.imgs[i]], [self.pred[i]], self.names) for i in range(self.n)]
|
228 |
+
for d in x:
|
229 |
+
for k in ['imgs', 'pred', 'xyxy', 'xyxyn', 'xywh', 'xywhn']:
|
230 |
+
setattr(d, k, getattr(d, k)[0]) # pop out of list
|
231 |
+
return x
|
232 |
+
|
233 |
|
234 |
class Flatten(nn.Module):
|
235 |
# Use after nn.AdaptiveAvgPool2d(1) to remove last 2 dimensions
|
test.py
CHANGED
@@ -12,7 +12,7 @@ from tqdm import tqdm
|
|
12 |
from models.experimental import attempt_load
|
13 |
from utils.datasets import create_dataloader
|
14 |
from utils.general import coco80_to_coco91_class, check_dataset, check_file, check_img_size, box_iou, \
|
15 |
-
non_max_suppression, scale_coords, xyxy2xywh, xywh2xyxy,
|
16 |
from utils.loss import compute_loss
|
17 |
from utils.metrics import ap_per_class
|
18 |
from utils.plots import plot_images, output_to_target
|
@@ -124,6 +124,7 @@ def test(data,
|
|
124 |
labels = targets[targets[:, 0] == si, 1:]
|
125 |
nl = len(labels)
|
126 |
tcls = labels[:, 0].tolist() if nl else [] # target class
|
|
|
127 |
seen += 1
|
128 |
|
129 |
if len(pred) == 0:
|
@@ -131,13 +132,14 @@ def test(data,
|
|
131 |
stats.append((torch.zeros(0, niou, dtype=torch.bool), torch.Tensor(), torch.Tensor(), tcls))
|
132 |
continue
|
133 |
|
|
|
|
|
|
|
|
|
134 |
# Append to text file
|
135 |
-
path = Path(paths[si])
|
136 |
if save_txt:
|
137 |
gn = torch.tensor(shapes[si][0])[[1, 0, 1, 0]] # normalization gain whwh
|
138 |
-
|
139 |
-
x[:, :4] = scale_coords(img[si].shape[1:], x[:, :4], shapes[si][0], shapes[si][1]) # to original
|
140 |
-
for *xyxy, conf, cls in x:
|
141 |
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
|
142 |
line = (cls, *xywh, conf) if save_conf else (cls, *xywh) # label format
|
143 |
with open(save_dir / 'labels' / (path.stem + '.txt'), 'a') as f:
|
@@ -150,19 +152,14 @@ def test(data,
|
|
150 |
"box_caption": "%s %.3f" % (names[cls], conf),
|
151 |
"scores": {"class_score": conf},
|
152 |
"domain": "pixel"} for *xyxy, conf, cls in pred.tolist()]
|
153 |
-
boxes = {"predictions": {"box_data": box_data, "class_labels": names}}
|
154 |
wandb_images.append(wandb.Image(img[si], boxes=boxes, caption=path.name))
|
155 |
|
156 |
-
# Clip boxes to image bounds
|
157 |
-
clip_coords(pred, (height, width))
|
158 |
-
|
159 |
# Append to pycocotools JSON dictionary
|
160 |
if save_json:
|
161 |
# [{"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}, ...
|
162 |
image_id = int(path.stem) if path.stem.isnumeric() else path.stem
|
163 |
-
box =
|
164 |
-
scale_coords(img[si].shape[1:], box, shapes[si][0], shapes[si][1]) # to original shape
|
165 |
-
box = xyxy2xywh(box) # xywh
|
166 |
box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner
|
167 |
for p, b in zip(pred.tolist(), box.tolist()):
|
168 |
jdict.append({'image_id': image_id,
|
@@ -178,6 +175,7 @@ def test(data,
|
|
178 |
|
179 |
# target boxes
|
180 |
tbox = xywh2xyxy(labels[:, 1:5]) * whwh
|
|
|
181 |
|
182 |
# Per target class
|
183 |
for cls in torch.unique(tcls_tensor):
|
@@ -187,7 +185,7 @@ def test(data,
|
|
187 |
# Search for detections
|
188 |
if pi.shape[0]:
|
189 |
# Prediction to target ious
|
190 |
-
ious, i = box_iou(
|
191 |
|
192 |
# Append detections
|
193 |
detected_set = set()
|
|
|
12 |
from models.experimental import attempt_load
|
13 |
from utils.datasets import create_dataloader
|
14 |
from utils.general import coco80_to_coco91_class, check_dataset, check_file, check_img_size, box_iou, \
|
15 |
+
non_max_suppression, scale_coords, xyxy2xywh, xywh2xyxy, set_logging, increment_path
|
16 |
from utils.loss import compute_loss
|
17 |
from utils.metrics import ap_per_class
|
18 |
from utils.plots import plot_images, output_to_target
|
|
|
124 |
labels = targets[targets[:, 0] == si, 1:]
|
125 |
nl = len(labels)
|
126 |
tcls = labels[:, 0].tolist() if nl else [] # target class
|
127 |
+
path = Path(paths[si])
|
128 |
seen += 1
|
129 |
|
130 |
if len(pred) == 0:
|
|
|
132 |
stats.append((torch.zeros(0, niou, dtype=torch.bool), torch.Tensor(), torch.Tensor(), tcls))
|
133 |
continue
|
134 |
|
135 |
+
# Predictions
|
136 |
+
predn = pred.clone()
|
137 |
+
scale_coords(img[si].shape[1:], predn[:, :4], shapes[si][0], shapes[si][1]) # native-space pred
|
138 |
+
|
139 |
# Append to text file
|
|
|
140 |
if save_txt:
|
141 |
gn = torch.tensor(shapes[si][0])[[1, 0, 1, 0]] # normalization gain whwh
|
142 |
+
for *xyxy, conf, cls in predn.tolist():
|
|
|
|
|
143 |
xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh
|
144 |
line = (cls, *xywh, conf) if save_conf else (cls, *xywh) # label format
|
145 |
with open(save_dir / 'labels' / (path.stem + '.txt'), 'a') as f:
|
|
|
152 |
"box_caption": "%s %.3f" % (names[cls], conf),
|
153 |
"scores": {"class_score": conf},
|
154 |
"domain": "pixel"} for *xyxy, conf, cls in pred.tolist()]
|
155 |
+
boxes = {"predictions": {"box_data": box_data, "class_labels": names}} # inference-space
|
156 |
wandb_images.append(wandb.Image(img[si], boxes=boxes, caption=path.name))
|
157 |
|
|
|
|
|
|
|
158 |
# Append to pycocotools JSON dictionary
|
159 |
if save_json:
|
160 |
# [{"image_id": 42, "category_id": 18, "bbox": [258.15, 41.29, 348.26, 243.78], "score": 0.236}, ...
|
161 |
image_id = int(path.stem) if path.stem.isnumeric() else path.stem
|
162 |
+
box = xyxy2xywh(predn[:, :4]) # xywh
|
|
|
|
|
163 |
box[:, :2] -= box[:, 2:] / 2 # xy center to top-left corner
|
164 |
for p, b in zip(pred.tolist(), box.tolist()):
|
165 |
jdict.append({'image_id': image_id,
|
|
|
175 |
|
176 |
# target boxes
|
177 |
tbox = xywh2xyxy(labels[:, 1:5]) * whwh
|
178 |
+
scale_coords(img[si].shape[1:], tbox, shapes[si][0], shapes[si][1]) # native-space labels
|
179 |
|
180 |
# Per target class
|
181 |
for cls in torch.unique(tcls_tensor):
|
|
|
185 |
# Search for detections
|
186 |
if pi.shape[0]:
|
187 |
# Prediction to target ious
|
188 |
+
ious, i = box_iou(predn[pi, :4], tbox[ti]).max(1) # best ious, indices
|
189 |
|
190 |
# Append detections
|
191 |
detected_set = set()
|