Update app.py

app.py

@@ -1,65 +1,19 @@
-from fastapi import FastAPI, UploadFile, File
-from PIL import Image
-import torch
-import torchvision.transforms as transforms
-import torchvision.models as models
 from huggingface_hub import hf_hub_download
-import io
-
-app = FastAPI()
+from safetensors.torch import load_file
 
 MODEL_REPO = "itsomk/chexpert-densenet121"
-MODEL_FILE = "chexpert_densenet121.pth"
+MODEL_FILE = "chexpert_pytorch.safetensors"
 
+# Download weights file
 model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
 
-model = models.densenet121(pretrained=False)
-model.classifier = torch.nn.Linear(model.classifier.in_features, 14)
-
-checkpoint = torch.load(model_path, map_location="cpu")
-
-if "state_dict" in checkpoint:
-    checkpoint = checkpoint["state_dict"]
-
-state_dict = {k.replace("module.", ""): v for k, v in checkpoint.items()}
-model.load_state_dict(state_dict, strict=False)
+# Load the safetensors weights
+state = load_file(model_path)
 
+# Build the DenseNet121 model and load weights
+import torchvision.models as models
+model = models.densenet121(pretrained=False)
+num_features = model.classifier.in_features
+model.classifier = torch.nn.Linear(num_features, 14)
+model.load_state_dict(state, strict=False)
 model.eval()
-
-
-# Image preprocessing (CheXpert standard)
-transform = transforms.Compose([
-    transforms.Resize((224, 224)),
-    transforms.ToTensor(),
-    transforms.Normalize(
-        mean=[0.485, 0.456, 0.406],
-        std=[0.229, 0.224, 0.225]
-    )
-])
-
-LABELS = [
-    "Atelectasis", "Cardiomegaly", "Consolidation", "Edema",
-    "Enlarged Cardiomediastinum", "Fracture", "Lung Lesion",
-    "Lung Opacity", "No Finding", "Pleural Effusion",
-    "Pleural Other", "Pneumonia", "Pneumothorax", "Support Devices"
-]
-
-@app.get("/")
-def health():
-    return {"status": "CheXpert DenseNet121 is running"}
-
-@app.post("/predict")
-async def predict(file: UploadFile = File(...)):
-    image = Image.open(io.BytesIO(await file.read())).convert("RGB")
-    image = transform(image).unsqueeze(0)
-
-    with torch.no_grad():
-        logits = model(image)
-        probs = torch.sigmoid(logits)[0]
-
-    results = {
-        LABELS[i]: float(probs[i])
-        for i in range(len(LABELS))
-    }
-
-    return results