diff --git a/test.py b/test.py index 0c06f37..901d2dd 100644 --- a/test.py +++ b/test.py @@ -209,7 +209,7 @@ def test(data, f = save_dir / f'test_batch{batch_i}_labels.jpg' # filename plot_images(img, targets, paths, f, names) # labels f = save_dir / f'test_batch{batch_i}_pred.jpg' - plot_images(img, output_to_target(output, width, height), paths, f, names) # predictions + plot_images(img, output_to_target(output), paths, f, names) # predictions # Compute statistics stats = [np.concatenate(x, 0) for x in zip(*stats)] # to numpy diff --git a/utils/datasets.py b/utils/datasets.py index 1cd03ba..520f0fb 100755 --- a/utils/datasets.py +++ b/utils/datasets.py @@ -443,7 +443,6 @@ class LoadImagesAndLabels(Dataset): # for training/testing assert (shape[0] > 9) & (shape[1] > 9), 'image size <10 pixels' # verify labels - l = [] if os.path.isfile(lb_file): nf += 1 # label found with open(lb_file, 'r') as f: @@ -458,6 +457,7 @@ class LoadImagesAndLabels(Dataset): # for training/testing l = np.zeros((0, 5), dtype=np.float32) else: nm += 1 # label missing + l = np.zeros((0, 5), dtype=np.float32) x[im_file] = [l, shape] except Exception as e: nc += 1 @@ -470,7 +470,7 @@ class LoadImagesAndLabels(Dataset): # for training/testing print(f'WARNING: No labels found in {path}. See {help_url}') x['hash'] = get_hash(self.label_files + self.img_files) - x['results'] = [nf, nm, ne, nc, i] + x['results'] = [nf, nm, ne, nc, i + 1] torch.save(x, path) # save for next time logging.info(f"New cache created: {path}") return x diff --git a/utils/plots.py b/utils/plots.py index 4d375cf..c49d106 100644 --- a/utils/plots.py +++ b/utils/plots.py @@ -86,25 +86,12 @@ def plot_wh_methods(): # from utils.plots import *; plot_wh_methods() fig.savefig('comparison.png', dpi=200) -def output_to_target(output, width, height): +def output_to_target(output): # Convert model output to target format [batch_id, class_id, x, y, w, h, conf] - if isinstance(output, torch.Tensor): - output = output.cpu().numpy() - targets = [] for i, o in enumerate(output): - if o is not None: - for pred in o: - box = pred[:4] - w = (box[2] - box[0]) / width - h = (box[3] - box[1]) / height - x = box[0] / width + w / 2 - y = box[1] / height + h / 2 - conf = pred[4] - cls = int(pred[5]) - - targets.append([i, cls, x, y, w, h, conf]) - + for *box, conf, cls in o.cpu().numpy(): + targets.append([i, cls, *list(*xyxy2xywh(np.array(box)[None])), conf]) return np.array(targets) @@ -153,9 +140,7 @@ def plot_images(images, targets, paths=None, fname='images.jpg', names=None, max labels = image_targets.shape[1] == 6 # labels if no conf column conf = None if labels else image_targets[:, 6] # check for confidence presence (label vs pred) - boxes[[0, 2]] *= w boxes[[0, 2]] += block_x - boxes[[1, 3]] *= h boxes[[1, 3]] += block_y for j, box in enumerate(boxes.T): cls = int(classes[j])