示例#1
0
    def on_batch_end(self, state: RunnerState):
        names = state.input[self.outpath_key]
        images = state.input[self.input_key].cpu()
        images = tensor_to_ndimage(images, dtype=np.uint8)

        for image, name in zip(images, names):
            fname = self.get_image_path(state, name, self.filename_suffix)
            imageio.imwrite(fname, image)
    def apply(self, img: torch.Tensor, **params) -> np.ndarray:
        """Apply the transform to the image"""
        if len(img.shape) == 2:
            img = img.unsqueeze(0)

        return tensor_to_ndimage(img,
                                 denormalize=self.denormalize,
                                 move_channels_dim=self.move_channels_dim)
示例#3
0
    def on_batch_end(self, state: RunnerState):
        if self.valid_only and state.need_backward:
            return

        names = state.input[self.input_image_name_key]
        images = tensor_to_ndimage(state.input[self.input_image_key].cpu(),
                                   dtype=np.uint8)
        in_masks = state.input[self.input_mask_key].cpu().numpy()
        out_masks = state.output[self.output_mask_key]
        in_objects = state.input[self.output_objects_key]
        out_objects = state.output[self.output_objects_key]

        image_metrics = state.output[self.output_image_metrics_key]
        for image_idx, image in enumerate(images):
            name = names[image_idx]
            metrics = image_metrics[image_idx]

            assert self.iou_threshold in metrics
            vis_folders = self._get_categories(metrics[self.iou_threshold])

            if not vis_folders:
                continue

            # original
            self.save_everywhere(state, vis_folders, name, "img", image)
            # masks
            for suffix, mask in zip(
                ("m_in", "m_out"),
                (in_masks[image_idx], out_masks[image_idx])):
                for d_cls_id, d_tm_id, d_name, d_classified in zip(
                        self.target_map_info.d_class_ids,
                        self.target_map_info.d_tm_ids,
                        self.target_map_info.d_class_names,
                        self.target_map_info.d_classified):
                    save_img = Visualizer.draw_hitmap(mask[d_tm_id], image)
                    self.save_everywhere(state, vis_folders, name,
                                         f"{d_name}.{suffix}", save_img)
                    if d_classified and not self.detection_only:
                        for c_name, c_tm_id in zip(
                                self.target_map_info.get_c_names(d_cls_id),
                                self.target_map_info.get_c_tm_ids(d_cls_id)):
                            save_img = Visualizer.draw_hitmap(
                                mask[c_tm_id], image)
                            self.save_everywhere(state, vis_folders, name,
                                                 f"{c_name}.{suffix}",
                                                 save_img)
            # objects
            for suffix, objects in zip(
                ("o_in", "o_out"),
                (in_objects[image_idx], out_objects[image_idx])):
                objects_mask = self.converter.build_target_map(
                    objects,
                    image_size=image.shape[:2],
                    for_visualization=True)
                save_img = Visualizer.draw_target_on_image(objects_mask, image)
                self.save_everywhere(state, vis_folders, name, suffix,
                                     save_img)
示例#4
0
    def on_batch_end(self, state: RunnerState):
        names = state.input[self.outpath_key]
        images = tensor_to_ndimage(state.input[self.input_key].cpu())
        masks = state.output[self.output_key]

        for name, image, mask in zip(names, images, masks):
            image = mask_to_overlay_image(image, mask, self.mask_strength)
            fname = self.get_image_path(state, name, self.filename_suffix)
            imageio.imwrite(fname, image)
示例#5
0
def test_tensor_to_ndimage():
    orig_images = np.random.randint(0, 255, (2, 20, 10, 3), np.uint8)

    torch_images = torch.stack([
        normalize(to_tensor(im), _IMAGENET_MEAN, _IMAGENET_STD)
        for im in orig_images
    ],
                               dim=0)

    byte_images = tensor_to_ndimage(torch_images, dtype=np.uint8)
    float_images = tensor_to_ndimage(torch_images, dtype=np.float32)

    assert np.allclose(byte_images, orig_images)
    assert np.allclose(float_images, orig_images / 255, atol=1e-3, rtol=1e-3)

    assert np.allclose(tensor_to_ndimage(torch_images[0]),
                       orig_images[0] / 255,
                       atol=1e-3,
                       rtol=1e-3)
示例#6
0
    def on_batch_end(self, state: RunnerState):
        names = state.input[self.outpath_key]
        images = tensor_to_ndimage(state.input[self.input_key].cpu())
        if self.mask_in_input:
            masks = state.input[self.mask_key].cpu().squeeze(1).numpy()
        else:
            masks = state.output[self.mask_key]

        for name, image, mask in zip(names, images, masks):
            image = Visualizer.draw_target_on_image(mask, image,
                                                    self.mask_strength)
            fname = self.get_image_path(state, name, self.filename_suffix)
            imageio.imwrite(fname, image)