finish app

Dockerfile

@@ -1,4 +1,4 @@
-FROM python:3.9-slim
+FROM python:3.11-slim
 
 WORKDIR /app
 
@@ -9,6 +9,7 @@ RUN apt-get update && apt-get install -y \
     git \
     && rm -rf /var/lib/apt/lists/*
 
+RUN pip3 install wheel
 COPY requirements.txt ./
 COPY src/ ./src/
 

src/modelOps.py

@@ -4,11 +4,12 @@ import torch
 from torch import nn
 from PIL import Image
 from pathlib import Path
+from torchvision.models import MobileNet_V2_Weights
 
 # 1. Load Model
 @st.cache_resource
-def load_model(MODEL_PATH: Path = "src/outputs/p2_e29_best_model.pth", device: str = "cpu"):
-    model = models.mobilenet_v2(pretrained=True)
+def load_model(MODEL_PATH: Path = Path("src/outputs/p2_e29_best_model.pth"), device: str = "cpu"):
+    model = models.mobilenet_v2(weights=MobileNet_V2_Weights.DEFAULT)
     model.classifier[1] = nn.Linear(model.classifier[1].in_features, 10)
     
     state_dict = torch.load(MODEL_PATH, map_location=device)

src/streamlit_app.py

@@ -11,7 +11,7 @@ def main():
     st.write("Upload an X-ray image to classify potential fractures.")
 
     device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    model = load_model(device)
+    model = load_model(device=device)
 
     uploaded_file = st.file_uploader("Upload an X-ray image", type=["jpg", "jpeg", "png"])