glenn-jocher commited on
Commit
225845e
·
unverified ·
1 Parent(s): df0e408

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()

Files changed (2) hide show
  1. models/common.py +14 -0
  2. 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, clip_coords, 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,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
- x = pred.clone()
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 = pred[:, :4].clone() # xyxy
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(pred[pi, :4], tbox[ti]).max(1) # best ious, indices
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()