123456789101112131415161718192021222324252627282930313233343536373839 |
- import os
- os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"
- import torch
- from annotator.oneformer.detectron2.config import get_cfg
- from annotator.oneformer.detectron2.projects.deeplab import add_deeplab_config
- from annotator.oneformer.detectron2.data import MetadataCatalog
- from annotator.oneformer.oneformer import (
- add_oneformer_config,
- add_common_config,
- add_swin_config,
- add_dinat_config,
- )
- from annotator.oneformer.oneformer.demo.defaults import DefaultPredictor
- from annotator.oneformer.oneformer.demo.visualizer import Visualizer, ColorMode
- def make_detectron2_model(config_path, ckpt_path):
- cfg = get_cfg()
- add_deeplab_config(cfg)
- add_common_config(cfg)
- add_swin_config(cfg)
- add_oneformer_config(cfg)
- add_dinat_config(cfg)
- cfg.merge_from_file(config_path)
- cfg.MODEL.WEIGHTS = ckpt_path
- cfg.freeze()
- metadata = MetadataCatalog.get(cfg.DATASETS.TEST_PANOPTIC[0] if len(cfg.DATASETS.TEST_PANOPTIC) else "__unused")
- return DefaultPredictor(cfg), metadata
- def semantic_run(img, predictor, metadata):
- predictions = predictor(img[:, :, ::-1], "semantic") # Predictor of OneFormer must use BGR image !!!
- visualizer_map = Visualizer(img, is_img=False, metadata=metadata, instance_mode=ColorMode.IMAGE)
- out_map = visualizer_map.draw_sem_seg(predictions["sem_seg"].argmax(dim=0).cpu(), alpha=1, is_text=False).get_image()
- return out_map
|