Update README.md

README.md

@@ -1,47 +1,29 @@
-import torch
-from PIL import Image
-from transformers import AutoModel, AutoTokenizer, AutoImageProcessor
-
-device = "cuda" if torch.cuda.is_available() else "cpu"
-
-# ===== 加载模型 =====
-# model = AutoModel.from_pretrained(
-#     "StanfordAIMI/CheXficient",
-#     trust_remote_code=True
-# ).to(device)
-
-model = AutoModel.from_pretrained(
-    "/mnt/d/torch/CheXficient/hf_model",
-    trust_remote_code=True
-).to(device)
-
-
-# ===== 加载tokenizer =====
-tokenizer = AutoTokenizer.from_pretrained(
-    "emilyalsentzer/Bio_ClinicalBERT"
-)
-
-# ===== 加载image processor =====
-image_processor = AutoImageProcessor.from_pretrained(
-    "facebook/dinov2-base"
-)
-
-# ===== 准备数据 =====
-image = Image.open("xray.jpg").convert("RGB")
-text = ["pneumonia", "no acute cardiopulmonary abnormality"]
-
-image_inputs = image_processor(images=image, return_tensors="pt").to(device)
-text_inputs = tokenizer(text, padding=True, return_tensors="pt").to(device)
-
-# ===== 推理 =====
-with torch.no_grad():
-    outputs = model(
-        pixel_values=image_inputs["pixel_values"],
-        input_ids=text_inputs["input_ids"],
-        attention_mask=text_inputs["attention_mask"],
-    )
-
-logits = outputs["logits_per_image"]
-probs = logits.softmax(dim=-1)
-
-print(probs)
+import torch
+from PIL import Image
+from transformers import AutoModel, AutoTokenizer, AutoImageProcessor
+
+repo_id = "StanfordAIMI/CheXficient"
+device = "cuda" if torch.cuda.is_available() else "cpu"
+
+model = AutoModel.from_pretrained(
+    repo_id,
+    trust_remote_code=True
+).to(device)
+
+tokenizer = AutoTokenizer.from_pretrained(repo_id, trust_remote_code=True)
+image_processor = AutoImageProcessor.from_pretrained(repo_id, trust_remote_code=True)
+
+model.eval()
+
+image = Image.open("./CXR/images/5AF3BB6C1BCC83C.png").convert("RGB")
+text = ["Pneumonia", "no Pneumonia"]
+
+image_inputs = image_processor(images=image, return_tensors="pt").to(device)
+text_inputs = tokenizer(text, padding=True, return_tensors="pt").to(device)
+
+with torch.no_grad():
+    outputs = model(
+        pixel_values=image_inputs["pixel_values"],
+        text_tokens=text_inputs,
+    )
+print(outputs)