def __test_img(self, image_path, save_path):
     image_raw = ImageHelper.cv2_open_bgr(image_path)
     inputs = ImageHelper.bgr2rgb(image_raw)
     heatmap_avg = self.__get_heatmap(inputs)
     all_peaks = self.__extract_heatmap_info(heatmap_avg)
     image_save = self.__draw_key_point(all_peaks, image_raw)
     cv2.imwrite(save_path, image_save)
    def __test_img(self, image_path, save_path):
        Log.info('Image Path: {}'.format(image_path))
        image_raw = ImageHelper.cv2_open_bgr(image_path)
        inputs = ImageHelper.bgr2rgb(image_raw)
        inputs = ImageHelper.resize(inputs, tuple(self.configer.get('data', 'input_size')), Image.CUBIC)
        inputs = ToTensor()(inputs)
        inputs = Normalize(mean=self.configer.get('trans_params', 'mean'),
                           std=self.configer.get('trans_params', 'std'))(inputs)

        with torch.no_grad():
            inputs = inputs.unsqueeze(0).to(self.device)
            bbox, cls = self.det_net(inputs)

        bbox = bbox.cpu().data.squeeze(0)
        cls = F.softmax(cls.cpu().squeeze(0), dim=-1).data
        boxes, lbls, scores, has_obj = self.__decode(bbox, cls)
        if has_obj:
            boxes = boxes.cpu().numpy()
            boxes = np.clip(boxes, 0, 1)
            lbls = lbls.cpu().numpy()
            scores = scores.cpu().numpy()

            img_canvas = self.__draw_box(image_raw, boxes, lbls, scores)

        else:
            # print('None obj detected!')
            img_canvas = image_raw

        Log.info('Save Path: {}'.format(save_path))
        cv2.imwrite(save_path, img_canvas)
        # Boxes is within 0-1.
        self.__save_json(save_path, boxes, lbls, scores, image_raw)

        return image_raw, lbls, scores, boxes, has_obj
Beispiel #3
0
 def __test_img(self, image_path, save_path):
     Log.info('Image Path: {}'.format(image_path))
     image_raw = ImageHelper.cv2_open_bgr(image_path)
     inputs = ImageHelper.bgr2rgb(image_raw)
     paf_avg, heatmap_avg = self.__get_paf_and_heatmap(inputs)
     all_peaks = self.__extract_heatmap_info(heatmap_avg)
     special_k, connection_all = self.__extract_paf_info(image_raw, paf_avg, all_peaks)
     subset, candidate = self.__get_subsets(connection_all, special_k, all_peaks)
     subset, img_canvas = self.__draw_key_point(subset, all_peaks, image_raw)
     img_canvas = self.__link_key_point(img_canvas, candidate, subset)
     cv2.imwrite(save_path, img_canvas)