paddlex 3.0.0b2__py3-none-any.whl → 3.0.0rc0__py3-none-any.whl

paddlex/inference/models/instance_segmentation/result.py

@@ -0,0 +1,161 @@
+# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import os
+import cv2
+import numpy as np
+import copy
+import PIL
+from PIL import Image, ImageDraw, ImageFont
+
+from ...utils.color_map import get_colormap, font_colormap
+from ...common.result import BaseCVResult, StrMixin, JsonMixin
+from ....utils.fonts import PINGFANG_FONT_FILE_PATH
+from ..object_detection.result import draw_box
+
+
+def draw_segm(im, masks, mask_info, alpha=0.7):
+    """
+    Draw segmentation on image
+    """
+    mask_color_id = 0
+    w_ratio = 0.4
+    color_list = get_colormap(rgb=True)
+    im = np.array(im).astype("float32")
+    clsid2color = {}
+    masks = np.array(masks)
+    masks = masks.astype(np.uint8)
+    for i in range(masks.shape[0]):
+        mask, score, clsid = masks[i], mask_info[i]["score"], mask_info[i]["class_id"]
+
+        if clsid not in clsid2color:
+            color_index = i % len(color_list)
+            clsid2color[clsid] = color_list[color_index]
+        color_mask = clsid2color[clsid]
+        for c in range(3):
+            color_mask[c] = color_mask[c] * (1 - w_ratio) + w_ratio * 255
+        idx = np.nonzero(mask)
+        color_mask = np.array(color_mask)
+        idx0 = np.minimum(idx[0], im.shape[0] - 1)
+        idx1 = np.minimum(idx[1], im.shape[1] - 1)
+        im[idx0, idx1, :] *= 1.0 - alpha
+        im[idx0, idx1, :] += alpha * color_mask
+        sum_x = np.sum(mask, axis=0)
+        x = np.where(sum_x > 0.5)[0]
+        sum_y = np.sum(mask, axis=1)
+        y = np.where(sum_y > 0.5)[0]
+        x0, x1, y0, y1 = x[0], x[-1], y[0], y[-1]
+        cv2.rectangle(
+            im, (x0, y0), (x1, y1), tuple(color_mask.astype("int32").tolist()), 1
+        )
+        bbox_text = "%s %.2f" % (mask_info[i]["label"], score)
+        t_size = cv2.getTextSize(bbox_text, 0, 0.3, thickness=1)[0]
+        cv2.rectangle(
+            im,
+            (x0, y0),
+            (x0 + t_size[0], y0 - t_size[1] - 3),
+            tuple(color_mask.astype("int32").tolist()),
+            -1,
+        )
+        cv2.putText(
+            im,
+            bbox_text,
+            (x0, y0 - 2),
+            cv2.FONT_HERSHEY_SIMPLEX,
+            0.3,
+            (0, 0, 0),
+            1,
+            lineType=cv2.LINE_AA,
+        )
+    return Image.fromarray(im.astype("uint8"))
+
+
+def restore_to_draw_masks(img_size, boxes, masks):
+    """
+    Restores extracted masks to the original shape and draws them on a blank image.
+
+    """
+
+    restored_masks = []
+
+    for i, (box, mask) in enumerate(zip(boxes, masks)):
+        restored_mask = np.zeros(img_size, dtype=np.uint8)
+        x_min, y_min, x_max, y_max = map(lambda x: int(round(x)), box["coordinate"])
+        restored_mask[y_min:y_max, x_min:x_max] = mask
+        restored_masks.append(restored_mask)
+
+    return np.array(restored_masks)
+
+
+def draw_mask(im, boxes, np_masks, img_size):
+    """
+    Args:
+        im (PIL.Image.Image): PIL image
+        boxes (list): a list of dictionaries representing detection box information.
+        np_masks (np.ndarray): shape:[N, im_h, im_w]
+    Returns:
+        im (PIL.Image.Image): visualized image
+    """
+    color_list = get_colormap(rgb=True)
+    w_ratio = 0.4
+    alpha = 0.7
+    im = np.array(im).astype("float32")
+    clsid2color = {}
+    np_masks = restore_to_draw_masks(img_size, boxes, np_masks)
+    im_h, im_w = im.shape[:2]
+    np_masks = np_masks[:, :im_h, :im_w]
+    for i in range(len(np_masks)):
+        clsid, score = int(boxes[i]["cls_id"]), boxes[i]["score"]
+        mask = np_masks[i]
+        if clsid not in clsid2color:
+            color_index = i % len(color_list)
+            clsid2color[clsid] = color_list[color_index]
+        color_mask = clsid2color[clsid]
+        for c in range(3):
+            color_mask[c] = color_mask[c] * (1 - w_ratio) + w_ratio * 255
+        idx = np.nonzero(mask)
+        color_mask = np.array(color_mask)
+        im[idx[0], idx[1], :] *= 1.0 - alpha
+        im[idx[0], idx[1], :] += alpha * color_mask
+    return Image.fromarray(im.astype("uint8"))
+
+
+class InstanceSegResult(BaseCVResult):
+    """Save Result Transform"""
+
+    def _to_img(self):
+        """apply"""
+        # image = self._img_reader.read(self["input_path"])
+        image = Image.fromarray(self["input_img"])
+        ori_img_size = list(image.size)[::-1]
+        boxes = self["boxes"]
+        masks = self["masks"]
+        if next((True for item in self["boxes"] if "coordinate" in item), False):
+            image = draw_mask(image, boxes, masks, ori_img_size)
+            image = draw_box(image, boxes)
+        else:
+            image = draw_segm(image, masks, boxes)
+
+        return {"res": image}
+
+    def _to_str(self, *args, **kwargs):
+        data = copy.deepcopy(self)
+        data.pop("input_img")
+        data["masks"] = "..."
+        return JsonMixin._to_str(data, *args, **kwargs)
+
+    def _to_json(self, *args, **kwargs):
+        data = copy.deepcopy(self)
+        data.pop("input_img")
+        return JsonMixin._to_json(data, *args, **kwargs)

paddlex/inference/models/keypoint_detection/__init__.py

@@ -0,0 +1,15 @@
+# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .predictor import KptPredictor

paddlex/inference/models/keypoint_detection/predictor.py

@@ -0,0 +1,188 @@
+# copyright (c) 2024 PaddlePaddle Authors. All Rights Reserve.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Any, List, Optional, Sequence
+
+import numpy as np
+
+from ....modules.keypoint_detection.model_list import MODELS
+from ....utils import logging
+from ...common.batch_sampler import ImageBatchSampler
+
+from ..common import ToBatch
+from ..object_detection import DetPredictor
+from .processors import TopDownAffine, KptPostProcess
+from .result import KptResult
+
+
+class KptBatchSampler(ImageBatchSampler):
+    def sample(self, inputs):
+        if not isinstance(inputs, list):
+            inputs = [inputs]
+
+        batch = []
+        for input in inputs:
+            if isinstance(input, (np.ndarray, dict)):
+                batch.append(input)
+                if len(batch) == self.batch_size:
+                    yield batch
+                    batch = []
+            elif isinstance(input, str):
+                file_path = (
+                    self._download_from_url(input)
+                    if input.startswith("http")
+                    else input
+                )
+                file_list = self._get_files_list(file_path)
+                for file_path in file_list:
+                    batch.append(file_path)
+                    if len(batch) == self.batch_size:
+                        yield batch
+                        batch = []
+            else:
+                logging.warning(
+                    f"Not supported input data type! Only `numpy.ndarray` and `str` are supported! So has been ignored: {input}."
+                )
+        if len(batch) > 0:
+            yield batch
+
+
+class KptPredictor(DetPredictor):
+
+    entities = MODELS
+
+    flip_perm = [  # The left-right joints exchange order list
+        [1, 2],
+        [3, 4],
+        [5, 6],
+        [7, 8],
+        [9, 10],
+        [11, 12],
+        [13, 14],
+        [15, 16],
+    ]
+
+    def __init__(
+        self,
+        *args,
+        flip: bool = False,
+        use_udp: Optional[bool] = None,
+        **kwargs,
+    ):
+        """Keypoint Predictor
+
+        Args:
+            flip (bool): Whether to do flipping test. Default value is ``False``.
+            use_udp (Optional[bool]): Whether to use unbiased data processing. Default value is ``None``.
+
+        """
+        self.flip = flip
+        self.use_udp = use_udp
+        super().__init__(*args, **kwargs)
+        for op in self.pre_ops:
+            if isinstance(op, TopDownAffine):
+                self.input_size = op.input_size
+                break
+        if any([name in self.model_name for name in ["PP-TinyPose"]]):
+            self.shift_heatmap = True
+        else:
+            self.shift_heatmap = False
+
+    def _build_batch_sampler(self):
+        return KptBatchSampler()
+
+    def _get_result_class(self):
+        return KptResult
+
+    def _format_output(self, pred: Sequence[Any]) -> List[dict]:
+        """Transform batch outputs into a list of single image output."""
+
+        return [
+            {
+                "heatmap": res[0],
+                "masks": res[1],
+            }
+            for res in zip(*pred)
+        ]
+
+    def flip_back(self, output_flipped, matched_parts):
+        assert (
+            output_flipped.ndim == 4
+        ), "output_flipped should be [batch_size, num_joints, height, width]"
+
+        output_flipped = output_flipped[:, :, :, ::-1]
+
+        for pair in matched_parts:
+            tmp = output_flipped[:, pair[0], :, :].copy()
+            output_flipped[:, pair[0], :, :] = output_flipped[:, pair[1], :, :]
+            output_flipped[:, pair[1], :, :] = tmp
+
+        return output_flipped
+
+    def process(self, batch_data: List[dict]):
+        """
+        Process a batch of data through the preprocessing, inference, and postprocessing.
+
+        Args:
+            batch_data (List[Union[str, np.ndarray], ...]): A batch of input data (e.g., image file paths).
+
+        Returns:
+            dict: A dictionary containing the input path, raw image, class IDs, scores, and label names
+                for every instance of the batch. Keys include 'input_path', 'input_img', 'class_ids', 'scores', and 'label_names'.
+        """
+        datas = batch_data
+        # preprocess
+        for pre_op in self.pre_ops[:-1]:
+            datas = pre_op(datas)
+
+        # use `ToBatch` format batch inputs
+        batch_inputs = self.pre_ops[-1]([data["img"] for data in datas])
+
+        # do infer
+        batch_preds = self.infer(batch_inputs)
+
+        if self.flip:
+            # flip w
+            batch_inputs[0] = np.flip(batch_inputs[0], axis=3)
+            preds_flipped = self.infer(batch_inputs)
+
+            output_flipped = self.flip_back(preds_flipped[0], self.flip_perm)
+            if self.shift_heatmap:
+                output_flipped[:, :, :, 1:] = output_flipped.copy()[:, :, :, 0:-1]
+            batch_preds[0] = (batch_preds[0] + output_flipped) * 0.5
+
+        # process a batch of predictions into a list of single image result
+        preds_list = self._format_output(batch_preds)
+
+        # postprocess
+        keypoints = self.post_op(preds_list, datas)
+
+        return {
+            "input_path": [data.get("img_path", None) for data in datas],
+            "input_img": [data["ori_img"] for data in datas],
+            "kpts": keypoints,
+        }
+
+    @DetPredictor.register("TopDownEvalAffine")
+    def build_topdown_affine(self, trainsize, use_udp=False):
+        return TopDownAffine(
+            input_size=trainsize,
+            use_udp=use_udp if self.use_udp is None else self.use_udp,
+        )
+
+    def build_to_batch(self):
+        return ToBatch()
+
+    def build_postprocess(self):
+        return KptPostProcess(use_dark=True)