Ejemplo n.º 1
0
    def compute_inputs(self, image_group, annotations_group):
        """
        Compute inputs for the network using an image_group.
        """
        # construct an image batch object
        batch_images = np.zeros((len(image_group), self.input_size, self.input_size, 3), dtype=np.float32)

        batch_hms = np.zeros((len(image_group), self.output_size, self.output_size, self.num_classes()),
                             dtype=np.float32)
        batch_hms_2 = np.zeros((len(image_group), self.output_size, self.output_size, self.num_classes()),
                               dtype=np.float32)
        batch_whs = np.zeros((len(image_group), self.max_objects, 2), dtype=np.float32)
        batch_regs = np.zeros((len(image_group), self.max_objects, 2), dtype=np.float32)
        batch_reg_masks = np.zeros((len(image_group), self.max_objects), dtype=np.float32)
        batch_indices = np.zeros((len(image_group), self.max_objects), dtype=np.float32)

        # copy all images to the upper left part of the image batch object
        for b, (image, annotations) in enumerate(zip(image_group, annotations_group)):
            c = np.array([image.shape[1] / 2., image.shape[0] / 2.], dtype=np.float32)
            s = max(image.shape[0], image.shape[1]) * 1.0
            trans_input = get_affine_transform(c, s, self.input_size)

            # inputs
            image = self.preprocess_image(image, c, s, tgt_w=self.input_size, tgt_h=self.input_size)
            batch_images[b] = image

            # outputs
            bboxes = annotations['bboxes']
            #assert bboxes.shape[0] != 0
            class_ids = annotations['labels']
            #assert class_ids.shape[0] != 0

            trans_output = get_affine_transform(c, s, self.output_size)
            for i in range(bboxes.shape[0]):
                bbox = bboxes[i].copy()
                cls_id = (int)(class_ids[i])
                # (x1, y1)
                bbox[:2] = affine_transform(bbox[:2], trans_output)
                # (x2, y2)
                bbox[2:] = affine_transform(bbox[2:], trans_output)
                bbox[[0, 2]] = np.clip(bbox[[0, 2]], 0, self.output_size - 1)
                bbox[[1, 3]] = np.clip(bbox[[1, 3]], 0, self.output_size - 1)
                h, w = bbox[3] - bbox[1], bbox[2] - bbox[0]
                if h > 0 and w > 0:
                    radius_h, radius_w = gaussian_radius((math.ceil(h), math.ceil(w)))
                    radius_h = max(0, int(radius_h))
                    radius_w = max(0, int(radius_w))

                    radius = gaussian_radius_2((math.ceil(h), math.ceil(w)))
                    radius = max(0, int(radius))
                    ct = np.array([(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2], dtype=np.float32)
                    ct_int = ct.astype(np.int32)
                    draw_gaussian(batch_hms[b, :, :, cls_id], ct_int, radius_h, radius_w)
                    draw_gaussian_2(batch_hms_2[b, :, :, cls_id], ct_int, radius)
                    batch_whs[b, i] = 1. * w, 1. * h
                    batch_indices[b, i] = ct_int[1] * self.output_size + ct_int[0]
                    batch_regs[b, i] = ct - ct_int
                    batch_reg_masks[b, i] = 1

                    #open this code to show generated images and labels

                    # hm = batch_hms[b, :, :, cls_id]
                    # hm = np.round(hm * 255).astype(np.uint8)
                    # hm = cv2.cvtColor(hm, cv2.COLOR_GRAY2BGR)
                    # hm_2 = batch_hms_2[b, :, :, cls_id]
                    # hm_2 = np.round(hm_2 * 255).astype(np.uint8)
                    # hm_2 = cv2.cvtColor(hm_2, cv2.COLOR_GRAY2BGR)
                    # cv2.rectangle(hm, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])), (0, 255, 0), 1)
                    # cv2.rectangle(hm_2, (int(bbox[0]), int(bbox[1])), (int(bbox[2]), int(bbox[3])), (0, 255, 0), 1)
                    # cv2.namedWindow('hm', cv2.WINDOW_NORMAL)
                    # cv2.imshow('hm', np.hstack([hm, hm_2]))
                    # cv2.waitKey()
            # print(np.sum(batch_reg_masks[b]))
            # for i in range(self.num_classes()):
            #     plt.subplot(4, 5, i + 1)
            #     hm = batch_hms[b, :, :, i]
            #     plt.imshow(hm, cmap='gray')
            #     plt.axis('off')
            # plt.show()
            # hm = np.sum(batch_hms[0], axis=-1)
            # hm = np.round(hm * 255).astype(np.uint8)
            # hm = cv2.cvtColor(hm, cv2.COLOR_GRAY2BGR)
            # hm_2 = np.sum(batch_hms_2[0], axis=-1)
            # hm_2 = np.round(hm_2 * 255).astype(np.uint8)
            # hm_2 = cv2.cvtColor(hm_2, cv2.COLOR_GRAY2BGR)
            # for i in range(bboxes.shape[0]):
            #     x1, y1 = np.round(affine_transform(bboxes[i, :2], trans_input)).astype(np.int32)
            #     x2, y2 = np.round(affine_transform(bboxes[i, 2:], trans_input)).astype(np.int32)
            #     x1_, y1_ = np.round(affine_transform(bboxes[i, :2], trans_output)).astype(np.int32)
            #     x2_, y2_ = np.round(affine_transform(bboxes[i, 2:], trans_output)).astype(np.int32)
            #     class_id = class_ids[i]
            #     cv2.rectangle(image, (x1, y1), (x2, y2), (0, 255, 0), 1)
            #     cv2.rectangle(hm, (x1_, y1_), (x2_, y2_), (0, 255, 0), 1)
            #     cv2.rectangle(hm_2, (x1_, y1_), (x2_, y2_), (0, 255, 0), 1)
            # cv2.namedWindow('hm', cv2.WINDOW_NORMAL)
            # cv2.imshow('hm', np.hstack([hm, hm_2]))
            # cv2.namedWindow('image', cv2.WINDOW_NORMAL)
            # cv2.imshow('image', image)
            # cv2.waitKey()

        return [batch_images, batch_hms_2, batch_whs, batch_regs, batch_reg_masks, batch_indices]
Ejemplo n.º 2
0
    def compute_inputs(self, image_group, annotations_group):
        """
        Compute inputs for the network using an image_group.
        """
        # construct an image batch object
        batch_images = np.zeros((len(image_group), self.input_size, self.input_size, 3), dtype=np.float32)

        batch_hms = np.zeros((len(image_group), self.output_size, self.output_size, self.num_classes()),
                             dtype=np.float32)
        batch_hms_2 = np.zeros((len(image_group), self.output_size, self.output_size, self.num_classes()),
                               dtype=np.float32)
        batch_whs = np.zeros((len(image_group), self.max_objects, 2), dtype=np.float32)
        batch_regs = np.zeros((len(image_group), self.max_objects, 2), dtype=np.float32)
        batch_reg_masks = np.zeros((len(image_group), self.max_objects), dtype=np.float32)
        batch_indices = np.zeros((len(image_group), self.max_objects), dtype=np.float32)
        
        # copy all images to the upper left part of the image batch object
        for b, (image, annotations) in enumerate(zip(image_group, annotations_group)):
            c = np.array([image.shape[1] / 2., image.shape[0] / 2.], dtype=np.float32)
            s = max(image.shape[0], image.shape[1]) * 1.0
            trans_input = get_affine_transform(c, s, self.input_size)

            # inputs
            image = self.preprocess_image(image, c, s, tgt_w=self.input_size, tgt_h=self.input_size)
            batch_images[b] = image

            # outputs
            bboxes = annotations['bboxes']
            
            assert bboxes.shape[0] != 0
            class_ids = annotations['labels']
            assert class_ids.shape[0] != 0

            trans_output = get_affine_transform(c, s, self.output_size)
            for i in range(bboxes.shape[0]):
                bbox = bboxes[i].copy()
                cls_id = class_ids[i]
                # (x1, y1)
                bbox[:2] = affine_transform(bbox[:2], trans_output)
                # (x2, y2)
                bbox[2:] = affine_transform(bbox[2:], trans_output)
                bbox[[0, 2]] = np.clip(bbox[[0, 2]], 0, self.output_size - 1)
                bbox[[1, 3]] = np.clip(bbox[[1, 3]], 0, self.output_size - 1)
                h, w = bbox[3] - bbox[1], bbox[2] - bbox[0]
                if h > 0 and w > 0:
                    
                    radius_h, radius_w = gaussian_radius((math.ceil(h), math.ceil(w)))
                    radius_h = max(0, int(radius_h))
                    radius_w = max(0, int(radius_w))

                    radius = gaussian_radius_2((math.ceil(h), math.ceil(w)))
                    radius = max(0, int(radius))
                    ct = np.array([(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2], dtype=np.float32)
                    ct_int = ct.astype(np.int32)
                    draw_gaussian(batch_hms[b, :, :, cls_id], ct_int, (radius_h, radius_w))
                    draw_gaussian(batch_hms_2[b, :, :, cls_id], ct_int, radius)
                    batch_whs[b, i] = 1. * w, 1. * h
                    batch_indices[b, i] = ct_int[1] * self.output_size + ct_int[0]
                    batch_regs[b, i] = ct - ct_int
                    batch_reg_masks[b, i] = 1

        return [batch_images, batch_hms_2, batch_whs, batch_regs, batch_reg_masks, batch_indices]