Update README.md

README.md

@@ -58,22 +58,29 @@ Use the code below to get started with the model.
 import cv2
 import torch
 
-from transformers import AutoImageProcessor, Swinv2ForImageClassification
+from transformers import AutoImageProcessor, SegformerForSemanticSegmentation
 
 image = cv2.imread('./example.jpg')
 image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
 
-processor = AutoImageProcessor.from_pretrained("pamixsun/swinv2_tiny_for_glaucoma_classification")
-model = Swinv2ForImageClassification.from_pretrained("pamixsun/swinv2_tiny_for_glaucoma_classification")
+processor = AutoImageProcessor.from_pretrained("pamixsun/segformer_for_optic_disc_cup_segmentation")
+model = SegformerForSemanticSegmentation.from_pretrained("pamixsun/segformer_for_optic_disc_cup_segmentation")
 
 inputs = processor(image, return_tensors="pt")
 
 with torch.no_grad():
-    logits = model(**inputs).logits
-
-# model predicts either glaucoma or non-glaucoma.
-predicted_label = logits.argmax(-1).item()
-print(model.config.id2label[predicted_label])
+    inputs.to(self.device)
+    outputs = self.seg_model(**inputs)
+    logits = outputs.logits.cpu()
+
+upsampled_logits = nn.functional.interpolate(
+    logits,
+    size=image.shape[:2],
+    mode="bilinear",
+    align_corners=False,
+)
+
+pred_disc_cup = upsampled_logits.argmax(dim=1)[0].numpy().astype(np.uint8)
 
 ```