vkit.element.score_map

  1# Copyright 2022 vkit-x Administrator. All Rights Reserved.
  2#
  3# This project (vkit-x/vkit) is dual-licensed under commercial and SSPL licenses.
  4#
  5# The commercial license gives you the full rights to create and distribute software
  6# on your own terms without any SSPL license obligations. For more information,
  7# please see the "LICENSE_COMMERCIAL.txt" file.
  8#
  9# This project is also available under Server Side Public License (SSPL).
 10# The SSPL licensing is ideal for use cases such as open source projects with
 11# SSPL distribution, student/academic purposes, hobby projects, internal research
 12# projects without external distribution, or other projects where all SSPL
 13# obligations can be met. For more information, please see the "LICENSE_SSPL.txt" file.
 14from typing import Optional, Tuple, Union, Iterable, Callable, List, TypeVar
 15from contextlib import ContextDecorator
 16
 17import attrs
 18import numpy as np
 19import cv2 as cv
 20
 21from .type import Shapable, ElementSetOperationMode
 22from .opt import generate_shape_and_resized_shape
 23
 24
 25@attrs.define
 26class ScoreMapSetItemConfig:
 27    value: Union[
 28        'ScoreMap',
 29        np.ndarray,
 30        float,
 31    ] = 1.0  # yapf: disable
 32    keep_max_value: bool = False
 33    keep_min_value: bool = False
 34
 35
 36@attrs.define
 37class NpVec:
 38    x: np.ndarray
 39    y: np.ndarray
 40
 41    @classmethod
 42    def from_point(cls, point: 'Point'):
 43        return cls(
 44            x=np.asarray(point.smooth_x, dtype=np.float32),
 45            y=np.asarray(point.smooth_y, dtype=np.float32),
 46        )
 47
 48    def __add__(self, other: 'NpVec'):
 49        return NpVec(x=self.x + other.x, y=self.y + other.y)
 50
 51    def __sub__(self, other: 'NpVec'):
 52        return NpVec(x=self.x - other.x, y=self.y - other.y)  # type: ignore
 53
 54    def __mul__(self, other: 'NpVec') -> np.ndarray:
 55        return self.x * other.y - self.y * other.x  # type: ignore
 56
 57
 58class WritableScoreMapContextDecorator(ContextDecorator):
 59
 60    def __init__(self, score_map: 'ScoreMap'):
 61        super().__init__()
 62        self.score_map = score_map
 63
 64    def __enter__(self):
 65        if self.score_map.mat.flags.c_contiguous:
 66            assert not self.score_map.mat.flags.writeable
 67
 68        try:
 69            self.score_map.mat.flags.writeable = True
 70        except ValueError:
 71            # Copy on write.
 72            object.__setattr__(
 73                self.score_map,
 74                'mat',
 75                np.array(self.score_map.mat),
 76            )
 77            assert self.score_map.mat.flags.writeable
 78
 79    def __exit__(self, *exc):  # type: ignore
 80        self.score_map.mat.flags.writeable = False
 81
 82
 83_E = TypeVar('_E', 'Box', 'Polygon', 'Mask')
 84
 85
 86@attrs.define(frozen=True, eq=False)
 87class ScoreMap(Shapable):
 88    mat: np.ndarray
 89    box: Optional['Box'] = None
 90    is_prob: bool = True
 91
 92    def __attrs_post_init__(self):
 93        if self.mat.ndim != 2:
 94            raise RuntimeError('ndim should == 2.')
 95        if self.box and self.shape != self.box.shape:
 96            raise RuntimeError('self.shape != box.shape.')
 97
 98        if self.mat.dtype != np.float32:
 99            raise RuntimeError('mat.dtype != np.float32')
100
101        # For the control of write.
102        self.mat.flags.writeable = False
103
104        if self.is_prob:
105            score_min = self.mat.min()
106            score_max = self.mat.max()
107            if score_min < 0.0 or score_max > 1.0:
108                raise RuntimeError('score not in range [0.0, 1.0]')
109
110    ###############
111    # Constructor #
112    ###############
113    @classmethod
114    def from_shape(
115        cls,
116        shape: Tuple[int, int],
117        value: float = 0.0,
118        is_prob: bool = True,
119    ):
120        height, width = shape
121        if is_prob:
122            assert 0.0 <= value <= 1.0
123        mat = np.full((height, width), fill_value=value, dtype=np.float32)
124        return cls(mat=mat, is_prob=is_prob)
125
126    @classmethod
127    def from_shapable(
128        cls,
129        shapable: Shapable,
130        value: float = 0.0,
131        is_prob: bool = True,
132    ):
133        return cls.from_shape(
134            shape=shapable.shape,
135            value=value,
136            is_prob=is_prob,
137        )
138
139    @classmethod
140    def from_quad_interpolation(
141        cls,
142        point0: 'Point',
143        point1: 'Point',
144        point2: 'Point',
145        point3: 'Point',
146        func_np_uv_to_mat: Callable[[np.ndarray], np.ndarray],
147        is_prob: bool = True,
148    ):
149        '''
150        Ref: https://www.reedbeta.com/blog/quadrilateral-interpolation-part-2/
151
152        points in clockwise order:
153        point0 0.0 -- u --> 1.0 point1
154                                 0.0
155                                  ↓
156                                  v
157                                  ↓
158                                 1.0
159        point3 0.0 <-- u -- 1.0 point2
160
161        lerp(a, b, r) = a + r * (b - a)
162        pointx(u, v) = lerp(
163            lerp(point0, point1, u),
164            lerp(point3, point2, u),
165            v,
166        )
167
168        Hence,
169        u in [0.0, 1.0], and
170        u -> 0.0, x -> line (point0, point3)
171        u -> 1.0, x -> line (point1, point2)
172
173        v in [0.0, 1.0], and
174        v -> 0.0, x -> line (point0, point1)
175        v -> 1.0, x -> line (point3, point2)
176        '''
177        polygon = Polygon.create((
178            point0,
179            point1,
180            point2,
181            point3,
182        ))
183        bounding_box = polygon.bounding_box
184        self_relative_polygon = polygon.self_relative_polygon
185        np_active_mask = polygon.internals.np_mask
186
187        np_vec_0 = NpVec.from_point(self_relative_polygon.points[0])
188        np_vec_1 = NpVec.from_point(self_relative_polygon.points[1])
189        np_vec_2 = NpVec.from_point(self_relative_polygon.points[2])
190        np_vec_3 = NpVec.from_point(self_relative_polygon.points[3])
191
192        # F(x, y) -> x
193        np_pointx_x = np.repeat(
194            np.expand_dims(
195                np.arange(bounding_box.width, dtype=np.int32),
196                axis=0,
197            ),
198            bounding_box.height,
199            axis=0,
200        )
201        # F(x, y) -> y
202        np_pointx_y = np.repeat(
203            np.expand_dims(
204                np.arange(bounding_box.height, dtype=np.int32),
205                axis=1,
206            ),
207            bounding_box.width,
208            axis=1,
209        )
210        np_vec_x = NpVec(x=np_pointx_x, y=np_pointx_y)
211
212        np_vec_q = np_vec_x - np_vec_0
213        np_vec_b1 = np_vec_1 - np_vec_0
214        np_vec_b2 = np_vec_3 - np_vec_0
215        np_vec_b3 = ((np_vec_0 - np_vec_1) - np_vec_3) + np_vec_2
216
217        scale_a = float(np_vec_b2 * np_vec_b3)
218
219        np_b: np.ndarray = (np_vec_b3 * np_vec_q) - (np_vec_b1 * np_vec_b2)  # type: ignore
220        np_b = np_b.astype(np.float32)
221
222        np_c = np_vec_b1 * np_vec_q
223        np_c = np_c.astype(np.float32)
224
225        # Solve v.
226        if abs(scale_a) < 0.001:
227            np_v = -np_c / np_b
228
229        else:
230            np_discrim = np.sqrt(np.power(np_b, 2) - 4 * scale_a * np_c)
231            scale_i2a = 0.5 / scale_a
232            np_v_pos = (-np_b + np_discrim) * scale_i2a
233            np_v_neg: np.ndarray = (-np_b - np_discrim) * scale_i2a  # type: ignore
234
235            np_masked_v_pos: np.ndarray = np_v_pos[np_active_mask]
236            np_v_pos_valid = (0.0 <= np_masked_v_pos) & (np_masked_v_pos <= 1.0)
237
238            np_masked_v_neg: np.ndarray = np_v_neg[np_active_mask]
239            np_v_neg_valid = (0.0 <= np_masked_v_neg) & (np_masked_v_neg <= 1.0)
240
241            if np_v_pos_valid.sum() >= np_v_neg_valid.sum():
242                np_v = np_v_pos
243            else:
244                np_v = np_v_neg
245
246        np_v[~np_active_mask] = 0.0
247        np_v = np.clip(np_v, 0.0, 1.0)
248
249        # Solve u.
250        np_u = np.zeros_like(np_v)
251
252        np_denom_x: np.ndarray = np_vec_b1.x + np_vec_b3.x * np_v
253        np_denom_y: np.ndarray = np_vec_b1.y + np_vec_b3.y * np_v
254
255        np_denom_x_mask = (np.abs(np_denom_x) > np.abs(np_denom_y)) & (np_denom_x != 0.0)
256        if np_denom_x_mask.any():
257            np_q_x = np_vec_q.x
258            np_u[np_denom_x_mask] = (
259                (np_q_x[np_denom_x_mask] - np_vec_b2.x * np_v[np_denom_x_mask])
260                / np_denom_x[np_denom_x_mask]
261            )
262
263        np_denom_y_mask = (~np_denom_x_mask) & (np_denom_y != 0.0)
264        if np_denom_y_mask.any():
265            np_q_y = np_vec_q.y
266            np_u[np_denom_y_mask] = (
267                (np_q_y[np_denom_y_mask] - np_vec_b2.y * np_v[np_denom_y_mask])
268                / np_denom_y[np_denom_y_mask]
269            )
270
271        np_u[~np_active_mask] = 0.0
272        np_u = np.clip(np_u, 0.0, 1.0)
273
274        # Stack to (height, width, 2)
275        np_uv = np.stack((np_u, np_v), axis=-1)
276
277        # Mat.
278        mat = func_np_uv_to_mat(np_uv)
279        return cls(
280            mat=mat,
281            box=bounding_box,
282            is_prob=is_prob,
283        )
284
285    ############
286    # Property #
287    ############
288    @property
289    def height(self):
290        return self.mat.shape[0]
291
292    @property
293    def width(self):
294        return self.mat.shape[1]
295
296    @property
297    def equivalent_box(self):
298        return self.box or Box.from_shapable(self)
299
300    @property
301    def writable_context(self):
302        return WritableScoreMapContextDecorator(self)
303
304    ############
305    # Operator #
306    ############
307    def copy(self):
308        return attrs.evolve(self, mat=self.mat.copy())
309
310    def assign_mat(self, mat: np.ndarray):
311        with self.writable_context:
312            object.__setattr__(self, 'mat', mat)
313
314    @classmethod
315    def unpack_element_value_pairs(
316        cls,
317        is_prob: bool,
318        element_value_pairs: Iterable[Tuple[_E, Union['ScoreMap', np.ndarray, float]]],
319    ):
320        elements: List[_E] = []
321
322        values: List[Union[ScoreMap, np.ndarray, float]] = []
323        for box, value in element_value_pairs:
324            elements.append(box)
325
326            if is_prob:
327                if isinstance(value, ScoreMap):
328                    assert (0.0 <= value.mat).all() and (value.mat <= 1.0).all()
329                elif isinstance(value, np.ndarray):
330                    assert (0.0 <= value).all() and (value <= 1.0).all()
331                elif isinstance(value, float):
332                    assert 0.0 <= value <= 1.0
333                else:
334                    raise NotImplementedError()
335            values.append(value)
336
337        return elements, values
338
339    def fill_by_box_value_pairs(
340        self,
341        box_value_pairs: Iterable[Tuple['Box', Union['ScoreMap', np.ndarray, float]]],
342        mode: ElementSetOperationMode = ElementSetOperationMode.UNION,
343        keep_max_value: bool = False,
344        keep_min_value: bool = False,
345        skip_values_uniqueness_check: bool = False,
346    ):
347        boxes, values = self.unpack_element_value_pairs(self.is_prob, box_value_pairs)
348
349        boxes_mask = generate_fill_by_boxes_mask(self.shape, boxes, mode)
350        if boxes_mask is None:
351            for box, value in zip(boxes, values):
352                box.fill_score_map(
353                    score_map=self,
354                    value=value,
355                    keep_max_value=keep_max_value,
356                    keep_min_value=keep_min_value,
357                )
358
359        else:
360            unique = True
361            if not skip_values_uniqueness_check:
362                unique = check_elements_uniqueness(values)
363
364            if unique:
365                boxes_mask.fill_score_map(
366                    score_map=self,
367                    value=values[0],
368                    keep_max_value=keep_max_value,
369                    keep_min_value=keep_min_value,
370                )
371            else:
372                for box, value in zip(boxes, values):
373                    box_mask = box.extract_mask(boxes_mask).to_box_attached(box)
374                    box_mask.fill_score_map(
375                        score_map=self,
376                        value=value,
377                        keep_max_value=keep_max_value,
378                        keep_min_value=keep_min_value,
379                    )
380
381    def fill_by_boxes(
382        self,
383        boxes: Iterable['Box'],
384        value: Union['ScoreMap', np.ndarray, float] = 1.0,
385        mode: ElementSetOperationMode = ElementSetOperationMode.UNION,
386        keep_max_value: bool = False,
387        keep_min_value: bool = False,
388    ):
389        self.fill_by_box_value_pairs(
390            box_value_pairs=((box, value) for box in boxes),
391            mode=mode,
392            keep_max_value=keep_max_value,
393            keep_min_value=keep_min_value,
394            skip_values_uniqueness_check=True,
395        )
396
397    def fill_by_polygon_value_pairs(
398        self,
399        polygon_value_pairs: Iterable[Tuple['Polygon', Union['ScoreMap', np.ndarray, float]]],
400        mode: ElementSetOperationMode = ElementSetOperationMode.UNION,
401        keep_max_value: bool = False,
402        keep_min_value: bool = False,
403        skip_values_uniqueness_check: bool = False,
404    ):
405        polygons, values = self.unpack_element_value_pairs(self.is_prob, polygon_value_pairs)
406
407        polygons_mask = generate_fill_by_polygons_mask(self.shape, polygons, mode)
408        if polygons_mask is None:
409            for polygon, value in zip(polygons, values):
410                polygon.fill_score_map(
411                    score_map=self,
412                    value=value,
413                    keep_max_value=keep_max_value,
414                    keep_min_value=keep_min_value,
415                )
416
417        else:
418            unique = True
419            if not skip_values_uniqueness_check:
420                unique = check_elements_uniqueness(values)
421
422            if unique:
423                polygons_mask.fill_score_map(
424                    score_map=self,
425                    value=values[0],
426                    keep_max_value=keep_max_value,
427                    keep_min_value=keep_min_value,
428                )
429            else:
430                for polygon, value in zip(polygons, values):
431                    bounding_box = polygon.to_bounding_box()
432                    polygon_mask = bounding_box.extract_mask(polygons_mask)
433                    polygon_mask = polygon_mask.to_box_attached(bounding_box)
434                    polygon_mask.fill_score_map(
435                        score_map=self,
436                        value=value,
437                        keep_max_value=keep_max_value,
438                        keep_min_value=keep_min_value,
439                    )
440
441    def fill_by_polygons(
442        self,
443        polygons: Iterable['Polygon'],
444        value: Union['ScoreMap', np.ndarray, float] = 1.0,
445        mode: ElementSetOperationMode = ElementSetOperationMode.UNION,
446        keep_max_value: bool = False,
447        keep_min_value: bool = False,
448    ):
449        self.fill_by_polygon_value_pairs(
450            polygon_value_pairs=((polygon, value) for polygon in polygons),
451            mode=mode,
452            keep_max_value=keep_max_value,
453            keep_min_value=keep_min_value,
454            skip_values_uniqueness_check=True,
455        )
456
457    def fill_by_mask_value_pairs(
458        self,
459        mask_value_pairs: Iterable[Tuple['Mask', Union['ScoreMap', np.ndarray, float]]],
460        mode: ElementSetOperationMode = ElementSetOperationMode.UNION,
461        keep_max_value: bool = False,
462        keep_min_value: bool = False,
463        skip_values_uniqueness_check: bool = False,
464    ):
465        masks, values = self.unpack_element_value_pairs(self.is_prob, mask_value_pairs)
466
467        masks_mask = generate_fill_by_masks_mask(self.shape, masks, mode)
468        if masks_mask is None:
469            for mask, value in zip(masks, values):
470                mask.fill_score_map(
471                    score_map=self,
472                    value=value,
473                    keep_max_value=keep_max_value,
474                    keep_min_value=keep_min_value,
475                )
476
477        else:
478            unique = True
479            if not skip_values_uniqueness_check:
480                unique = check_elements_uniqueness(values)
481
482            if unique:
483                masks_mask.fill_score_map(
484                    score_map=self,
485                    value=values[0],
486                    keep_max_value=keep_max_value,
487                    keep_min_value=keep_min_value,
488                )
489            else:
490                for mask, value in zip(masks, values):
491                    if mask.box:
492                        boxed_mask = mask.box.extract_mask(masks_mask)
493                    else:
494                        boxed_mask = masks_mask
495
496                    boxed_mask = boxed_mask.copy()
497                    mask.to_inverted_mask().fill_mask(boxed_mask, value=0)
498                    boxed_mask.fill_score_map(
499                        score_map=self,
500                        value=value,
501                        keep_max_value=keep_max_value,
502                        keep_min_value=keep_min_value,
503                    )
504
505    def fill_by_masks(
506        self,
507        masks: Iterable['Mask'],
508        value: Union['ScoreMap', np.ndarray, float] = 1.0,
509        mode: ElementSetOperationMode = ElementSetOperationMode.UNION,
510        keep_max_value: bool = False,
511        keep_min_value: bool = False,
512    ):
513        self.fill_by_mask_value_pairs(
514            mask_value_pairs=((mask, value) for mask in masks),
515            mode=mode,
516            keep_max_value=keep_max_value,
517            keep_min_value=keep_min_value,
518            skip_values_uniqueness_check=True,
519        )
520
521    def __setitem__(
522        self,
523        element: Union[
524            'Box',
525            'Polygon',
526            'Mask',
527        ],
528        config: Union[
529            'ScoreMap',
530            np.ndarray,
531            float,
532            ScoreMapSetItemConfig,
533        ],
534    ):  # yapf: disable
535        if not isinstance(config, ScoreMapSetItemConfig):
536            value = config
537            keep_max_value = False
538            keep_min_value = False
539        else:
540            assert isinstance(config, ScoreMapSetItemConfig)
541            value = config.value
542            keep_max_value = config.keep_max_value
543            keep_min_value = config.keep_min_value
544
545        element.fill_score_map(
546            score_map=self,
547            value=value,
548            keep_min_value=keep_min_value,
549            keep_max_value=keep_max_value,
550        )
551
552    def __getitem__(
553        self,
554        element: Union[
555            'Box',
556            'Polygon',
557            'Mask',
558        ],
559    ):  # yapf: disable
560        return element.extract_score_map(self)
561
562    def fill_by_quad_interpolation(
563        self,
564        point0: 'Point',
565        point1: 'Point',
566        point2: 'Point',
567        point3: 'Point',
568        func_np_uv_to_mat: Callable[[np.ndarray], np.ndarray],
569        keep_max_value: bool = False,
570        keep_min_value: bool = False,
571    ):
572        score_map = self.from_quad_interpolation(
573            point0=point0,
574            point1=point1,
575            point2=point2,
576            point3=point3,
577            func_np_uv_to_mat=func_np_uv_to_mat,
578            is_prob=self.is_prob,
579        )
580        assert score_map.box
581        with self.writable_context:
582            score_map.box.fill_np_array(
583                mat=self.mat,
584                value=score_map.mat,
585                np_mask=(score_map.mat > 0.0),
586                keep_max_value=keep_max_value,
587                keep_min_value=keep_min_value,
588            )
589
590    def to_shifted_score_map(self, offset_y: int = 0, offset_x: int = 0):
591        assert self.box
592        shifted_box = self.box.to_shifted_box(offset_y=offset_y, offset_x=offset_x)
593        return attrs.evolve(self, box=shifted_box)
594
595    def to_conducted_resized_polygon(
596        self,
597        shapable_or_shape: Union[Shapable, Tuple[int, int]],
598        resized_height: Optional[int] = None,
599        resized_width: Optional[int] = None,
600        cv_resize_interpolation: int = cv.INTER_CUBIC,
601    ):
602        assert self.box
603        resized_box = self.box.to_conducted_resized_box(
604            shapable_or_shape=shapable_or_shape,
605            resized_height=resized_height,
606            resized_width=resized_width,
607        )
608        resized_score_map = self.to_box_detached().to_resized_score_map(
609            resized_height=resized_box.height,
610            resized_width=resized_box.width,
611            cv_resize_interpolation=cv_resize_interpolation,
612        )
613        resized_score_map = resized_score_map.to_box_attached(resized_box)
614        return resized_score_map
615
616    def to_resized_score_map(
617        self,
618        resized_height: Optional[int] = None,
619        resized_width: Optional[int] = None,
620        cv_resize_interpolation: int = cv.INTER_CUBIC,
621    ):
622        assert not self.box
623        _, _, resized_height, resized_width = generate_shape_and_resized_shape(
624            shapable_or_shape=self.shape,
625            resized_height=resized_height,
626            resized_width=resized_width,
627        )
628        mat = cv.resize(
629            self.mat,
630            (resized_width, resized_height),
631            interpolation=cv_resize_interpolation,
632        )
633        if self.is_prob:
634            # NOTE: Interpolation like bi-cubic could generate out-of-bound values.
635            mat = np.clip(mat, 0.0, 1.0)
636        return attrs.evolve(self, mat=mat)
637
638    def to_cropped_score_map(
639        self,
640        up: Optional[int] = None,
641        down: Optional[int] = None,
642        left: Optional[int] = None,
643        right: Optional[int] = None,
644    ):
645        assert not self.box
646
647        up = up or 0
648        down = down or self.height - 1
649        left = left or 0
650        right = right or self.width - 1
651
652        return attrs.evolve(self, mat=self.mat[up:down + 1, left:right + 1])
653
654    def to_box_attached(self, box: 'Box'):
655        assert self.height == box.height
656        assert self.width == box.width
657        return attrs.evolve(self, box=box)
658
659    def to_box_detached(self):
660        assert self.box
661        return attrs.evolve(self, box=None)
662
663    def fill_np_array(
664        self,
665        mat: np.ndarray,
666        value: Union[np.ndarray, Tuple[float, ...], float],
667        keep_max_value: bool = False,
668        keep_min_value: bool = False,
669    ):
670        self.equivalent_box.fill_np_array(
671            mat=mat,
672            value=value,
673            alpha=self,
674            keep_max_value=keep_max_value,
675            keep_min_value=keep_min_value,
676        )
677
678    def fill_image(
679        self,
680        image: 'Image',
681        value: Union['Image', np.ndarray, Tuple[int, ...], int],
682    ):
683        self.equivalent_box.fill_image(
684            image=image,
685            value=value,
686            alpha=self,
687        )
688
689    def to_mask(self, threshold: float = 0.0):
690        mat = (self.mat > threshold).astype(np.uint8)
691        return Mask(mat=mat, box=self.box)
692
693
694def generate_fill_by_score_maps_mask(
695    shape: Tuple[int, int],
696    score_maps: Iterable[ScoreMap],
697    mode: ElementSetOperationMode,
698):
699    if mode == ElementSetOperationMode.UNION:
700        return None
701    else:
702        return Mask.from_score_maps(shape, score_maps, mode)
703
704
705# Cyclic dependency, by design.
706from .uniqueness import check_elements_uniqueness  # noqa: E402
707from .image import Image  # noqa: E402
708from .point import Point  # noqa: E402
709from .box import Box, generate_fill_by_boxes_mask  # noqa: E402
710from .mask import Mask, generate_fill_by_masks_mask  # noqa: E402
711from .polygon import Polygon, generate_fill_by_polygons_mask  # noqa: E402