Spaces:
Running
on
Zero
Running
on
Zero
File size: 2,095 Bytes
a3a3ae4 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 |
# -*- coding: utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
from abc import ABCMeta
import torch
import torch.nn as nn
from scepter.modules.annotator.registry import ANNOTATORS
from scepter.modules.model.base_model import BaseModel
from scepter.modules.utils.config import dict_to_yaml
@ANNOTATORS.register_class()
class BaseAnnotator(BaseModel, metaclass=ABCMeta):
para_dict = {}
def __init__(self, cfg, logger=None):
super().__init__(cfg, logger=logger)
@torch.no_grad()
@torch.inference_mode
def forward(self, *args, **kwargs):
raise NotImplementedError
@staticmethod
def get_config_template():
return dict_to_yaml('ANNOTATORS',
__class__.__name__,
BaseAnnotator.para_dict,
set_name=True)
@ANNOTATORS.register_class()
class GeneralAnnotator(BaseAnnotator, metaclass=ABCMeta):
def __init__(self, cfg, logger=None):
super().__init__(cfg, logger=logger)
anno_models = cfg.get('ANNOTATORS', [])
self.annotators = nn.ModuleList()
for n, anno_config in enumerate(anno_models):
annotator = ANNOTATORS.build(anno_config, logger=logger)
annotator.input_keys = anno_config.get('INPUT_KEYS', [])
if isinstance(annotator.input_keys, str):
annotator.input_keys = [annotator.input_keys]
annotator.output_keys = anno_config.get('OUTPUT_KEYS', [])
if isinstance(annotator.output_keys, str):
annotator.output_keys = [annotator.output_keys]
assert len(annotator.input_keys) == len(annotator.output_keys)
self.annotators.append(annotator)
def forward(self, input_dict):
output_dict = {}
for annotator in self.annotators:
for idx, in_key in enumerate(annotator.input_keys):
if in_key in input_dict:
image = annotator(input_dict[in_key])
output_dict[annotator.output_keys[idx]] = image
return output_dict
|