(self)
| 181 | return (thresholded.to(device=image.device), mask.to(device=image.device),) |
| 182 | |
| 183 | def load_model(self): |
| 184 | global cached_clipseg_model |
| 185 | if cached_clipseg_model == None: |
| 186 | ensure_package("clipseg", "clipseg@git+https://github.com/timojl/clipseg.git@bbc86cfbb7e6a47fb6dae47ba01d3e1c2d6158b0") |
| 187 | from clipseg.clipseg import CLIPDensePredT |
| 188 | model = CLIPDensePredT(version='ViT-B/16', reduce_dim=64, complex_trans_conv=True) |
| 189 | model.eval() |
| 190 | |
| 191 | d64_file = self.download_and_cache('rd64-uni-refined.pth', 'https://owncloud.gwdg.de/index.php/s/ioHbRzFx6th32hn/download?path=%2F&files=rd64-uni-refined.pth') |
| 192 | d16_file = self.download_and_cache('rd16-uni.pth', 'https://owncloud.gwdg.de/index.php/s/ioHbRzFx6th32hn/download?path=%2F&files=rd16-uni.pth') |
| 193 | # Use CUDA if it's available |
| 194 | device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
| 195 | model.load_state_dict(torch.load(d64_file, map_location=device), strict=False) |
| 196 | model = model.eval().to(device=device) |
| 197 | cached_clipseg_model = model |
| 198 | return cached_clipseg_model |
| 199 | |
| 200 | def download_and_cache(self, cache_name, url): |
| 201 | cache_dir = os.path.join(os.path.dirname(os.path.abspath(__file__)), 'download_cache') |
no test coverage detected