Spaces:
Sleeping
Sleeping
File size: 4,664 Bytes
1560e05 |
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 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 |
from fastapi import FastAPI, Query
from fastapi.responses import JSONResponse
import torch
import torchvision
import numpy as np
import requests
import skimage.io
import cv2
import tempfile
import os
from PIL import Image
from transformers import AutoImageProcessor, AutoModel
import joblib
from pytorch_grad_cam import GradCAM
from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
import torchxrayvision as xrv
import requests
from io import BytesIO
import logging
logging.getLogger("uvicorn").setLevel(logging.WARNING)
import os
# Set a custom path for Matplotlib to use
os.environ["MPLCONFIGDIR"] = "/app/.cache/matplotlib"
# Set a custom path for TorchXRayVision to use for model weights
os.environ["TORCHXrayVISION_CACHE"] = "/app/.cache/torchxrayvision"
# Create these directories if they do not exist
os.makedirs("/app/.cache/matplotlib", exist_ok=True)
os.makedirs("/app/.cache/torchxrayvision", exist_ok=True)
app = FastAPI()
cxr_model = xrv.models.DenseNet(weights="densenet121-res224-all")
cxr_model.eval()
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
tb_processor = AutoImageProcessor.from_pretrained("StanfordAIMI/dinov2-base-xray-224")
tb_model = AutoModel.from_pretrained("StanfordAIMI/dinov2-base-xray-224").to(device)
logreg = joblib.load("logreg_model.joblib")
def preprocess_image(image_path):
img = skimage.io.imread(image_path)
img = xrv.datasets.normalize(img, 255)
if img.ndim == 3:
img = img.mean(2)[None, ...]
elif img.ndim == 2:
img = img[None, ...]
transform = torchvision.transforms.Compose([
xrv.datasets.XRayCenterCrop(),
xrv.datasets.XRayResizer(224)
])
img = transform(img)
return torch.from_numpy(img)
def get_predictions(img_tensor, model):
with torch.no_grad():
outputs = model(img_tensor[None, ...])
preds = dict(zip(model.pathologies, outputs[0].detach().numpy()))
return preds, outputs
def get_top_preds(preds, tolerance=0.01, topk=5):
sorted_preds = sorted(preds.items(), key=lambda x: -x[1])
top_conf = sorted_preds[0][1]
similar_preds = [(i, p, conf) for i, (p, conf) in enumerate(sorted_preds)
if abs(conf - top_conf) <= tolerance][:topk]
return sorted_preds, similar_preds
def get_bounding_boxes(img_tensor, model, similar_preds):
boxes = {}
target_layer = model.features[-1]
for idx, pathology, conf in similar_preds:
cam = GradCAM(model=model, target_layers=[target_layer])
pred_index = list(model.pathologies).index(pathology)
grayscale_cam = cam(input_tensor=img_tensor[None, ...],
targets=[ClassifierOutputTarget(pred_index)])[0]
cam_resized = cv2.resize(grayscale_cam, (224, 224))
cam_uint8 = (cam_resized * 255).astype(np.uint8)
_, thresh = cv2.threshold(cam_uint8, 100, 255, cv2.THRESH_BINARY)
contours, _ = cv2.findContours(thresh, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
if contours:
x, y, w, h = cv2.boundingRect(contours[0])
boxes[pathology] = [[x, y], [x + w, y + h]]
return boxes
def predict_tb(image_path):
image = Image.open(image_path)
inputs = tb_processor(images=image, return_tensors="pt").to(device)
with torch.no_grad():
outputs = tb_model(**inputs)
embeddings = outputs.pooler_output.cpu().numpy()
prediction = logreg.predict(embeddings)
return int(prediction[0] == "tb")
@app.get("/predict")
async def predict_cxr(image_url: str = Query(..., description="URL to a chest X-ray image")):
try:
response = requests.get(image_url)
if response.status_code != 200:
return JSONResponse(content={"error": "Failed to download image"}, status_code=400)
with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp:
tmp.write(response.content)
tmp_path = tmp.name
img_tensor = preprocess_image(tmp_path)
preds, _ = get_predictions(img_tensor, cxr_model)
sorted_preds, similar_preds = get_top_preds(preds)
prediction_result = {k: float(f"{v:.2f}") for k, v in preds.items()}
bounding_boxes = get_bounding_boxes(img_tensor, cxr_model, similar_preds)
tb_result = predict_tb(tmp_path)
os.remove(tmp_path)
return JSONResponse(content={
"prediction_result": prediction_result,
"bounding_box": bounding_boxes, # top-left , bottom-right coordinates
"tb_finding": tb_result
})
except Exception as e:
return JSONResponse(content={"error": str(e)}, status_code=500)
|