Update app.py

app.py

@@ -1,48 +1,34 @@
 import streamlit as st
-import torch
+from fastai.vision.all import *
 from PIL import Image
 import numpy as np
 import pathlib
+import plotly as px
 
-# Load the PyTorch model
-model_path = "model.pkl"
-model = torch.load(model_path, map_location=torch.device('cpu'))
-model.eval()
+# Temporary fix for pathlib.PosixPath on Windows
+temp = pathlib.PosixPath
+pathlib.PosixPath = pathlib.WindowsPath
 
-# Load the Fastai X-ray model
-xray_model_path = "xraydet.pkl"
-xray_model = torch.load(xray_model_path, map_location=torch.device('cpu'))
-xray_model.eval()
+# Load the Fastai Learner model
+learn = load_learner('model.pkl')
+learn_xray = load_learner("xraydet.pkl")
 
 def process_image(image):
     try: 
         img = image.resize((224, 224))
         img_fastai = Image.fromarray(np.array(img))
-        
-        # Preprocess image for the Fastai model
-        tfms = setup_aug_tfms([Resize(460)], mode='center')
-        img_fastai = PILImage.create(img_fastai).apply_tfms(tfms).resize((224, 224)).convert('RGB')
-        img_fastai = ToTensor()(img_fastai).unsqueeze(0)
-        
-        # Predict using the X-ray model
-        pred_xray, _, _ = xray_model.predict(img_fastai)
+        pred_xray, _, _ = learn_xray.predict(image)
         
         if pred_xray == '1':
-            # Predict using the main model
-            with torch.no_grad():
-                pred = model(img_fastai)
-                pred_id = torch.argmax(pred, dim=1).item()
-                probs = torch.softmax(pred, dim=1).numpy()[0]
-                labels = ['NORMAL', 'PNEUMONIA']
-                pred_label = labels[pred_id]
+            pred, pred_id, probs = learn.predict(img_fastai)
             
-            return pred_label, pred_id, probs
+            return pred, pred_id, probs
         else:
-            pred_label, pred_id, probs = -1, -1, -1
-            return pred_label, pred_id, probs
+            pred, pred_id, probs = -1,-1,-1
+            return pred, pred_id, probs
     except Exception:
-        pred_label, pred_id, probs = -1, -1, -1
-        return pred_label, pred_id, probs
+        pred, pred_id, probs = -1,-1,-1
+        return pred, pred_id, probs
 
 # Main function to run the Streamlit app
 def main():