File size: 1,203 Bytes
5718512 3f2d37f 5718512 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 |
import warnings
import torch
from transformers import AutoImageProcessor, AutoModel, AutoTokenizer
from utils import model_inference
# Suppress specific warnings for cleaner logs
warnings.filterwarnings("ignore", category=UserWarning)
def load_model(device, dtype):
tokenizer = AutoTokenizer.from_pretrained("Deepnoid/RadZero")
image_processor = AutoImageProcessor.from_pretrained("Deepnoid/RadZero")
model = AutoModel.from_pretrained(
"Deepnoid/RadZero",
trust_remote_code=True,
torch_dtype=dtype,
device_map=device,
)
model.eval()
models = {
"tokenizer": tokenizer,
"image_processor": image_processor,
"model": model,
}
return models
if __name__ == "__main__":
# Setup constant
device = torch.device("cuda")
dtype = torch.float32
# load models
models = load_model(device, dtype)
# load image
image_path = "cxr_image.jpg"
# inference
similarity_prob, similarity_map = model_inference(
image_path, "There is fibrosis", **models
)
print(similarity_prob)
print(similarity_map.min())
print(similarity_map.max())
print(similarity_map.shape)
|