Update app.py

app.py

@@ -2,19 +2,28 @@ from fastapi import FastAPI, UploadFile, File
 from PIL import Image
 import torch
 import torchvision.transforms as transforms
-from transformers import AutoModelForImageClassification
+import torchvision.models as models
+from huggingface_hub import hf_hub_download
 import io
 
 app = FastAPI()
 
-MODEL_NAME = "itsomk/chexpert-densenet121"
+MODEL_REPO = "itsomk/chexpert-densenet121"
+MODEL_FILE = "pytorch_model.bin"
 
-# Load model
-model = AutoModelForImageClassification.from_pretrained(MODEL_NAME)
+# Download weights
+model_path = hf_hub_download(repo_id=MODEL_REPO, filename=MODEL_FILE)
+
+# Load DenseNet121
+model = models.densenet121(pretrained=False)
+num_classes = 14
+model.classifier = torch.nn.Linear(model.classifier.in_features, num_classes)
+
+state_dict = torch.load(model_path, map_location="cpu")
+model.load_state_dict(state_dict)
 model.eval()
-model.to("cpu")
 
-# Image preprocessing (manual, because no image processor exists)
+# Image preprocessing (CheXpert standard)
 transform = transforms.Compose([
     transforms.Resize((224, 224)),
     transforms.ToTensor(),
@@ -24,20 +33,29 @@ transform = transforms.Compose([
     )
 ])
 
+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_bytes = await file.read()
-    image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-
-    image_tensor = transform(image).unsqueeze(0)
+    image = Image.open(io.BytesIO(await file.read())).convert("RGB")
+    image = transform(image).unsqueeze(0)
 
     with torch.no_grad():
-        outputs = model(image_tensor)
-        probs = torch.sigmoid(outputs.logits)[0]
+        logits = model(image)
+        probs = torch.sigmoid(logits)[0]
 
-    predictions = {
-        model.config.id2label[i]: float(probs[i])
-        for i in range(len(probs))
+    results = {
+        LABELS[i]: float(probs[i])
+        for i in range(len(LABELS))
     }
 
-    return {"predictions": predictions}
+    return results