| import torch |
| import numpy as np |
| from huggingface_hub import hf_hub_download |
| from segmentation import SegmentationModule |
|
|
| MODEL = None |
| DEVICE = torch.device("cpu") |
|
|
| def load_model(use_box=False): |
| global MODEL, DEVICE |
| MODEL = SegmentationModule(use_box=use_box) |
|
|
| ckpt_path = hf_hub_download( |
| repo_id="phoebe777777/111", |
| filename="microscopy_matching_seg.pth", |
| token=None, |
| force_download=False |
| ) |
| MODEL.load_state_dict(torch.load(ckpt_path, map_location="cpu"), strict=False) |
| MODEL.eval() |
| if torch.cuda.is_available(): |
| DEVICE = torch.device("cuda") |
| MODEL.move_to_device(DEVICE) |
| print("✅ Model moved to CUDA") |
| else: |
| DEVICE = torch.device("cpu") |
| MODEL.move_to_device(DEVICE) |
| print("✅ Model on CPU") |
| return MODEL, DEVICE |
|
|
|
|
| @torch.no_grad() |
| def run(model, img_path, box=None, device="cpu"): |
| print("DEVICE:", device) |
| model.move_to_device(device) |
| model.eval() |
| with torch.no_grad(): |
| if box is not None: |
| use_box = True |
| else: |
| use_box = False |
| model.use_box = use_box |
| output = model(img_path, box=box) |
| mask = output |
| return mask |
|
|
