| 116 | |
| 117 | |
| 118 | class SAM2Seg(BaseSeg): |
| 119 | RATIO_MAP = [[512, 1], [1280, 0.6], [1920, 0.4], [3840, 0.2]] |
| 120 | |
| 121 | def tocpu(self): |
| 122 | self.box_prior.cpu() |
| 123 | self.image_predictor.model.cpu() |
| 124 | torch.cuda.empty_cache() |
| 125 | |
| 126 | def tocuda(self): |
| 127 | self.box_prior.cuda() |
| 128 | self.image_predictor.model.cuda() |
| 129 | |
| 130 | def __init__( |
| 131 | self, |
| 132 | config="sam2.1_hiera_l.yaml", |
| 133 | matting_config="resnet50", |
| 134 | background=(1.0, 1.0, 1.0), |
| 135 | wo_supres=False, |
| 136 | ): |
| 137 | super().__init__() |
| 138 | |
| 139 | self.device = avaliable_device() |
| 140 | |
| 141 | try: |
| 142 | sam2_image_model = build_sam2(config, SAM2_WEIGHT) |
| 143 | except: |
| 144 | config = os.path.join("./configs/sam2.1/", config) # sam2.1 case |
| 145 | sam2_image_model = build_sam2(config, SAM2_WEIGHT) |
| 146 | |
| 147 | self.image_predictor = SAM2ImagePredictor(sam2_image_model) |
| 148 | |
| 149 | self.box_prior = None |
| 150 | |
| 151 | # Robust-Human-Matting |
| 152 | |
| 153 | # self.matting_predictor = MattingNetwork(matting_config).eval().cuda() |
| 154 | # self.matting_predictor.load_state_dict(torch.load(MATTING_WEIGHT)) |
| 155 | |
| 156 | self.background = background |
| 157 | self.wo_supers = wo_supres |
| 158 | |
| 159 | def clean_up(self): |
| 160 | self.tmp.cleanup() |
| 161 | |
| 162 | def collect_inputs(self, inputs): |
| 163 | return dict( |
| 164 | img_path=inputs["img_path"], |
| 165 | bbox=inputs["bbox"], |
| 166 | ) |
| 167 | |
| 168 | def _super_resolution(self, input_path): |
| 169 | |
| 170 | low = os.path.abspath(input_path) |
| 171 | high = self.tmp.name |
| 172 | |
| 173 | super_weights = os.path.abspath("./pretrained_models/RealESRGAN_x4plus.pth") |
| 174 | hander = os.path.join(SUPRES_PATH, "inference_realesrgan.py") |
| 175 |
no outgoing calls
no test coverage detected