diff --git a/test.py b/test.py index d72a792..9c652ca 100644 --- a/test.py +++ b/test.py @@ -128,11 +128,11 @@ def test(data, # Append to text file if save_txt: gn = torch.tensor(shapes[si][0])[[1, 0, 1, 0]] # normalization gain whwh - txt_path = str(out / Path(paths[si]).stem) - pred[:, :4] = scale_coords(img[si].shape[1:], pred[:, :4], shapes[si][0], shapes[si][1]) # to original - for *xyxy, conf, cls in pred: + x = pred.clone() + x[:, :4] = scale_coords(img[si].shape[1:], x[:, :4], shapes[si][0], shapes[si][1]) # to original + for *xyxy, conf, cls in x: xywh = (xyxy2xywh(torch.tensor(xyxy).view(1, 4)) / gn).view(-1).tolist() # normalized xywh - with open(txt_path + '.txt', 'a') as f: + with open(str(out / Path(paths[si]).stem) + '.txt', 'a') as f: f.write(('%g ' * 5 + '\n') % (cls, *xywh)) # label format # Clip boxes to image bounds