import cv2 import torch from transformers import AutoImageProcessor, Swinv2ForImageClassification from cam import ClassActivationMap class GlaucomaModel(object): def __init__(self, cls_model_path="pamixsun/swinv2_tiny_for_glaucoma_classification", device=torch.device('cpu')): # where to load the model, gpu or cpu ? self.device = device # classification model for nails disease self.cls_extractor = AutoImageProcessor.from_pretrained(cls_model_path) self.cls_model = Swinv2ForImageClassification.from_pretrained(cls_model_path).to(device).eval() # class activation map self.cam = ClassActivationMap(self.cls_model, self.cls_extractor) # classification id to label self.id2label = self.cls_model.config.id2label # number of classes for nails disease self.num_diseases = len(self.id2label) def glaucoma_pred(self, image): """ Args: image: image array in RGB order. """ inputs = self.cls_extractor(images=image.copy(), return_tensors="pt") with torch.no_grad(): inputs.to(self.device) outputs = self.cls_model(**inputs).logits disease_idx = outputs.cpu()[0, :].detach().numpy().argmax() return disease_idx def process(self, image): """ Args: image: image array in RGB order. """ image_shape = image.shape[:2] disease_idx = self.glaucoma_pred(image) cam = self.cam.get_cam(image, disease_idx) cam = cv2.resize(cam, image_shape[::-1]) return disease_idx, cam