| from ultralytics import YOLO |
| import cv2 |
| import torch |
|
|
| BEST_WEIGHT = './segmentation/weights/oai_s_best4.pt' |
| DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu") |
|
|
| def create_model(weight_path=None): |
| if weight_path: |
| return YOLO(weight_path) |
| else: |
| return YOLO(BEST_WEIGHT) |
|
|
| class Segmenter(): |
| def __init__(self, weight_path=None): |
| self.model = create_model(weight_path).to(DEVICE) |
| |
| def segment(self, img): |
| """ |
| input: image (H, W, C) |
| output: mask (H, W) with femur is 1 and tibia is 2 |
| """ |
| img = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) |
|
|
| eimg = cv2.equalizeHist(img) |
|
|
| clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)) |
| eimg = clahe.apply(eimg) |
|
|
| eimg = cv2.cvtColor(eimg, cv2.COLOR_GRAY2RGB) |
|
|
| res = self.model(eimg, verbose=False) |
| result = res[0] |
|
|
| if result.masks is None or result.boxes is None: |
| raise ValueError("Segmentation failed: no mask/object detected.") |
|
|
| masks = result.masks.data |
| classes = result.boxes.cls |
|
|
| if masks is None or len(masks) < 2 or len(classes) < 2: |
| raise ValueError("Segmentation failed: need at least 2 detected objects (femur and tibia).") |
|
|
| mask = torch.zeros_like(masks[0]) |
| n = min(len(masks), len(classes)) |
| for i in range(n): |
| class_value = int(classes[i].item()) + 1 |
| mask = torch.where(masks[i] > 0.5, torch.tensor(float(class_value), device=mask.device), mask) |
|
|
| mask = mask.cpu().numpy() |
|
|
| return mask |
|
|
