coreml-community/ControlNet-v1-1-Annotators-cpu

1import os2os.environ["KMP_DUPLICATE_LIB_OK"] = "TRUE"3 4import torch5 6from annotator.oneformer.detectron2.config import get_cfg7from annotator.oneformer.detectron2.projects.deeplab import add_deeplab_config8from annotator.oneformer.detectron2.data import MetadataCatalog9 10from annotator.oneformer.oneformer import (11    add_oneformer_config,12    add_common_config,13    add_swin_config,14    add_dinat_config,15)16 17from annotator.oneformer.oneformer.demo.defaults import DefaultPredictor18from annotator.oneformer.oneformer.demo.visualizer import Visualizer, ColorMode19 20 21def make_detectron2_model(config_path, ckpt_path):22    cfg = get_cfg()23    add_deeplab_config(cfg)24    add_common_config(cfg)25    add_swin_config(cfg)26    add_oneformer_config(cfg)27    add_dinat_config(cfg)28    cfg.merge_from_file(config_path)29    if torch.cuda.is_available():30        cfg.MODEL.DEVICE = 'cuda'31    else:32        cfg.MODEL.DEVICE = 'cpu'33    cfg.MODEL.WEIGHTS = ckpt_path34    cfg.freeze()35    metadata = MetadataCatalog.get(cfg.DATASETS.TEST_PANOPTIC[0] if len(cfg.DATASETS.TEST_PANOPTIC) else "__unused")36    return DefaultPredictor(cfg), metadata37 38 39def semantic_run(img, predictor, metadata):40    predictions = predictor(img[:, :, ::-1], "semantic")  # Predictor of OneFormer must use BGR image !!!41    visualizer_map = Visualizer(img, is_img=False, metadata=metadata, instance_mode=ColorMode.IMAGE)42    out_map = visualizer_map.draw_sem_seg(predictions["sem_seg"].argmax(dim=0).cpu(), alpha=1, is_text=False).get_image()43    return out_map44