| 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) | |
| mask = res[0].masks.data[0] * (res[0].boxes.cls[0] + 1) + res[0].masks.data[1] * (res[0].boxes.cls[1] + 1) | |
| mask = mask.cpu().numpy() | |
| return mask | |