| 40 | |
| 41 | |
| 42 | class DetectionModels: |
| 43 | def __init__(self) -> None: |
| 44 | # Preloading: Loading Models takes ~9 seconds |
| 45 | set_num_threads(5) |
| 46 | self.executor = ThreadPoolExecutor(max_workers=5) |
| 47 | self.loading_futures: List[Future[Callable[..., None]]] = [] |
| 48 | |
| 49 | try: |
| 50 | self.loading_futures.append(self.executor.submit(self._load_yolo_detector)) |
| 51 | self.loading_futures.append(self.executor.submit(self._load_vit_model)) |
| 52 | self.loading_futures.append(self.executor.submit(self._load_vit_processor)) |
| 53 | self.loading_futures.append(self.executor.submit(self._load_seg_model)) |
| 54 | self.loading_futures.append(self.executor.submit(self._load_seg_processor)) |
| 55 | except Exception as e: |
| 56 | if sys.version_info.minor >= 9: |
| 57 | self.executor.shutdown(wait=True, cancel_futures=True) |
| 58 | else: |
| 59 | self.executor.shutdown(wait=True) |
| 60 | raise e |
| 61 | |
| 62 | def _load_yolo_detector(self): |
| 63 | from ultralytics import YOLO |
| 64 | |
| 65 | self.yolo_model = YOLO("yolo11m-seg.pt") |
| 66 | |
| 67 | def _load_vit_model(self): |
| 68 | self.vit_model = CLIPModel.from_pretrained("flavour/CLIP-ViT-B-16-DataComp.XL-s13B-b90K") |
| 69 | |
| 70 | def _load_vit_processor(self): |
| 71 | self.vit_processor = CLIPProcessor.from_pretrained("flavour/CLIP-ViT-B-16-DataComp.XL-s13B-b90K") |
| 72 | |
| 73 | def _load_seg_model(self): |
| 74 | self.seg_model = CLIPSegForImageSegmentation.from_pretrained("CIDAS/clipseg-rd64-refined") |
| 75 | |
| 76 | def _load_seg_processor(self): |
| 77 | self.seg_processor = CLIPSegProcessor.from_pretrained("CIDAS/clipseg-rd64-refined") |
| 78 | |
| 79 | def check_loaded(self): |
| 80 | try: |
| 81 | if not all([future.done() for future in self.loading_futures]): |
| 82 | for future in self.loading_futures: |
| 83 | future.result() |
| 84 | |
| 85 | assert self.yolo_model |
| 86 | assert self.seg_model |
| 87 | assert self.vit_model |
| 88 | except Exception as e: |
| 89 | if sys.version_info.minor >= 9: |
| 90 | self.executor.shutdown(wait=True, cancel_futures=True) |
| 91 | else: |
| 92 | self.executor.shutdown(wait=True) |
| 93 | raise e |
| 94 | |
| 95 | |
| 96 | detection_models = DetectionModels() |