| 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) |
| if torch.cuda.is_available(): |
| cfg.MODEL.DEVICE = 'cuda' |
| else: |
| cfg.MODEL.DEVICE = 'cpu' |
| 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") |
| 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 |
|
|
