update v1.4.1

2025-04-02 19:27:43 +08:00
parent 467652bd69
commit 6c8af8d7a8
31 changed files with 2134 additions and 58 deletions
@@ -0,0 +1,58 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import random
+import numpy as np
+
+from scepter.modules.annotator.base_annotator import BaseAnnotator
+from scepter.modules.annotator.registry import ANNOTATORS
+from scepter.modules.utils.config import Config
+
+
+@ANNOTATORS.register_class()
+class FrameReferenceAnnotator(BaseAnnotator):
+    para_dict = {}
+
+    def __init__(self, cfg, logger=None):
+        super().__init__(cfg, logger=logger)
+        # first / last / firstlast / random
+        self.ref_cfg = cfg.get('REF_CFG', [{"mode": "first", "proba": 0.1},
+                                           {"mode": "last", "proba": 0.1},
+                                           {"mode": "firstlast", "proba": 0.1},
+                                           {"mode": "random", "proba": 0.1}])
+        self.ref_num = cfg.get('REF_NUM', 1)
+        self.ref_cfg = Config.get_dict(self.ref_cfg) if isinstance(
+            self.ref_cfg, Config) else self.ref_cfg
+        self.ref_color = cfg.get('REF_COLOR', 127.5)
+
+    def forward(self, frames, ref_cfg=None, ref_num=None):
+        ref_cfg = ref_cfg if ref_cfg is not None else self.ref_cfg
+        ref_cfg = [ref_cfg] if not isinstance(ref_cfg, list) else ref_cfg
+        probas = [item['proba'] if 'proba' in item else 1.0 / len(ref_cfg) for item in ref_cfg]
+        sel_ref_cfg = random.choices(ref_cfg, weights=probas, k=1)[0]
+        mode = sel_ref_cfg['mode'] if 'mode' in sel_ref_cfg else 'original'
+        ref_num = int(ref_num) if ref_num is not None else self.ref_num
+
+        frame_num = len(frames)
+        frame_num_range = list(range(frame_num))
+        if mode == "first":
+            sel_idx = frame_num_range[:ref_num]
+        elif mode == "last":
+            sel_idx = frame_num_range[-ref_num:]
+        elif mode == "firstlast":
+            sel_idx = frame_num_range[:ref_num] + frame_num_range[-ref_num:]
+        elif mode == "random":
+            sel_idx = random.sample(frame_num_range, ref_num)
+        else:
+            raise NotImplementedError
+
+        out_frames, out_masks = [], []
+        for i in range(frame_num):
+            if i in sel_idx:
+                out_frame = frames[i]
+                out_mask = np.zeros_like(frames[i][:, :, 0])
+            else:
+                out_frame = np.ones_like(frames[i]) * self.ref_color
+                out_mask = np.ones_like(frames[i][:, :, 0]) * 255
+            out_frames.append(out_frame)
+            out_masks.append(out_mask)
+        return out_frames, out_masks