Update app.py

app.py

@@ -1,120 +1,126 @@
-from fastapi import FastAPI, Query
-from pydantic import BaseModel
-from typing import List, Tuple
-from fastapi import Body
-
-import torch
-import torchxrayvision as xrv
-import torchvision
-import skimage.io
-import numpy as np
-import requests
-import cv2
-
-from io import BytesIO
-import matplotlib.pyplot as plt
-from pytorch_grad_cam import GradCAM
-from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
-from pytorch_grad_cam.utils.image import show_cam_on_image
-
-app = FastAPI()
-
-model = xrv.models.DenseNet(weights="densenet121-res224-all")
-model.eval()
-
-
-def preprocess_image_from_url(image_url: str) -> torch.Tensor:
-    response = requests.get(image_url)
-    img = skimage.io.imread(BytesIO(response.content))
-    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)
-    img_tensor = torch.from_numpy(img)
-    return img_tensor
-
-
-def get_predictions_and_bounding_box(img_tensor: torch.Tensor):
-    with torch.no_grad():
-        output = model(img_tensor[None, ...])[0]
-
-    predictions = dict(zip(model.pathologies, output.numpy()))
-    sorted_preds = sorted(predictions.items(), key=lambda x: -x[1])
-
-    top_pred_label, top_conf = sorted_preds[0]
-    top_pred_index = list(model.pathologies).index(top_pred_label)
-
-    target_layer = model.features[-1]
-    cam = GradCAM(model=model, target_layers=[target_layer])
-    grayscale_cam = cam(input_tensor=img_tensor[None, ...],
-                        targets=[ClassifierOutputTarget(top_pred_index)])[0, :]
-
-    input_img = img_tensor.numpy()[0]
-    input_img_norm = (input_img - input_img.min()) / (input_img.max() - input_img.min())
-    input_img_rgb = cv2.cvtColor((input_img_norm * 255).astype(np.uint8), cv2.COLOR_GRAY2RGB)
-
-    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)
-
-    bounding_boxes = []
-    for cnt in contours:
-        x, y, w, h = cv2.boundingRect(cnt)
-        bounding_boxes.append((int(x), int(y), int(x + w), int(y + h)))  
-
-    return sorted_preds, bounding_boxes
-
-
-class Prediction(BaseModel):
-    label: str
-    confidence: float
-
-
-class PredictionResponse(BaseModel):
-    predictions: List[Prediction]
-    top_prediction_bounding_boxes: List[Tuple[int, int, int, int]]
-
-
-@app.get("/predict", response_model=PredictionResponse)
-def predict(image_url: str = Query(..., description="URL of chest X-ray image")):
-    try:
-        img_tensor = preprocess_image_from_url(image_url)
-        preds, bboxes = get_predictions_and_bounding_box(img_tensor)
-        prediction_list = [Prediction(label=label, confidence=float(conf)) for label, conf in preds]
-
-        return PredictionResponse(
-            predictions=prediction_list,
-            top_prediction_bounding_boxes=bboxes
-        )
-    except Exception as e:
-        return {"error": str(e)}
-
-class URLRequest(BaseModel):
-    url: str
-
-@app.post("/predict", response_model=PredictionResponse)
-def predict_from_url(body: URLRequest):
-    try:
-        img_tensor = preprocess_image_from_url(body.url)
-        preds, bboxes = get_predictions_and_bounding_box(img_tensor)
-        prediction_list = [Prediction(label=label, confidence=float(conf)) for label, conf in preds]
-
-        return PredictionResponse(
-            predictions=prediction_list,
-            top_prediction_bounding_boxes=bboxes
-        )
-    except Exception as e:
-        return {"error": str(e)}
-    
-
-
-# uvicorn app:app --reload
+from fastapi import FastAPI, Query
+from pydantic import BaseModel
+from typing import List, Tuple
+from fastapi import Body
+
+import torch
+import torchxrayvision as xrv
+import torchvision
+import skimage.io
+import numpy as np
+import requests
+import cv2
+
+from io import BytesIO
+import matplotlib.pyplot as plt
+from pytorch_grad_cam import GradCAM
+from pytorch_grad_cam.utils.model_targets import ClassifierOutputTarget
+from pytorch_grad_cam.utils.image import show_cam_on_image
+
+app = FastAPI()
+
+model = xrv.models.DenseNet(weights="densenet121-res224-all")
+model.eval()
+
+
+import os
+
+# Set cache directory to a writable location
+os.environ['TORCHXRAYVISION_CACHE'] = './torchxrayvision_cache'
+os.environ['MPLCONFIGDIR'] = './.config/matplotlib'  # For matplotlib font issues
+
+def preprocess_image_from_url(image_url: str) -> torch.Tensor:
+    response = requests.get(image_url)
+    img = skimage.io.imread(BytesIO(response.content))
+    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)
+    img_tensor = torch.from_numpy(img)
+    return img_tensor
+
+
+def get_predictions_and_bounding_box(img_tensor: torch.Tensor):
+    with torch.no_grad():
+        output = model(img_tensor[None, ...])[0]
+
+    predictions = dict(zip(model.pathologies, output.numpy()))
+    sorted_preds = sorted(predictions.items(), key=lambda x: -x[1])
+
+    top_pred_label, top_conf = sorted_preds[0]
+    top_pred_index = list(model.pathologies).index(top_pred_label)
+
+    target_layer = model.features[-1]
+    cam = GradCAM(model=model, target_layers=[target_layer])
+    grayscale_cam = cam(input_tensor=img_tensor[None, ...],
+                        targets=[ClassifierOutputTarget(top_pred_index)])[0, :]
+
+    input_img = img_tensor.numpy()[0]
+    input_img_norm = (input_img - input_img.min()) / (input_img.max() - input_img.min())
+    input_img_rgb = cv2.cvtColor((input_img_norm * 255).astype(np.uint8), cv2.COLOR_GRAY2RGB)
+
+    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)
+
+    bounding_boxes = []
+    for cnt in contours:
+        x, y, w, h = cv2.boundingRect(cnt)
+        bounding_boxes.append((int(x), int(y), int(x + w), int(y + h)))  
+
+    return sorted_preds, bounding_boxes
+
+
+class Prediction(BaseModel):
+    label: str
+    confidence: float
+
+
+class PredictionResponse(BaseModel):
+    predictions: List[Prediction]
+    top_prediction_bounding_boxes: List[Tuple[int, int, int, int]]
+
+
+@app.get("/predict", response_model=PredictionResponse)
+def predict(image_url: str = Query(..., description="URL of chest X-ray image")):
+    try:
+        img_tensor = preprocess_image_from_url(image_url)
+        preds, bboxes = get_predictions_and_bounding_box(img_tensor)
+        prediction_list = [Prediction(label=label, confidence=float(conf)) for label, conf in preds]
+
+        return PredictionResponse(
+            predictions=prediction_list,
+            top_prediction_bounding_boxes=bboxes
+        )
+    except Exception as e:
+        return {"error": str(e)}
+
+class URLRequest(BaseModel):
+    url: str
+
+@app.post("/predict", response_model=PredictionResponse)
+def predict_from_url(body: URLRequest):
+    try:
+        img_tensor = preprocess_image_from_url(body.url)
+        preds, bboxes = get_predictions_and_bounding_box(img_tensor)
+        prediction_list = [Prediction(label=label, confidence=float(conf)) for label, conf in preds]
+
+        return PredictionResponse(
+            predictions=prediction_list,
+            top_prediction_bounding_boxes=bboxes
+        )
+    except Exception as e:
+        return {"error": str(e)}
+    
+
+
+# uvicorn app:app --reload