Update app.py - set the categories using learn

app.py

@@ -1,79 +1,79 @@
-# Bismillahir Rahmaanir Raheem 
-# Almadadh Ya Gause Radi Allahu Ta'alah Anh - Ameen 
-
-# Import necessary libraries from fastai and gradio
-from fastai.vision.all import *
-import gradio as gr
-
-
-# Function to determine if an image contains pneumonia
-# Checks if the filename contains 'virus' or 'bacteria'
-def is_pneumonia(x): 
-    return (x.find('virus')!=-1 or x.find('bacteria')!=-1)
-
-
-# Load the trained fastai model for predictions
-learn = load_learner('pneumonia_model.pkl')
-
-
-# Define the possible categories for prediction
-categories = ('Pneumonia', 'Normal') 
-
-
-# Function to make a prediction on an input image
-def predict(img):
-    pred, idx, probs = learn.predict(img)  # Get the prediction, index, and probabilities
-    return dict(zip(categories, map(float, probs)))  # Return the probabilities mapped to categories
-
-
-# Title of the Gradio interface
-title = "Pediatric Pneumonia Chest X-Ray Predictor"
-
-
-# Description of the interface, including model and dataset information
-description = """
-A pediatric pneumonia chest x-ray predictor model trained on the chest-xray-pneumonia dataset using ResNet34 via 
-<a href='http://www.fast.ai/' target='_blank'>fast.ai</a>. The dataset is from: 
-<a href='http://www.kaggle.com/datasets/paultimothymooney/chest-xray-pneumonia' target='_blank'>Chest X-Ray Images (Pneumonia)</a> 
-and the associated scientific journal paper is 
-<a href='http://www.cell.com/cell/fulltext/S0092-8674(18)30154-5' target='_blank'>Identifying Medical Diagnoses and Treatable 
-Diseases by Image-Based Deep Learning</a>. The accuracy of the model is: 87.50%
-"""
-
-
-# Article or additional information to be displayed
-article = """
-<p style='text-align: center'><span style='font-size: 15pt;'>Pediatric Pneumonia Chest X-Ray Predictor. Dr Zakia Salod. 2024. </span></p>
-"""
-
-
-# Gradio input component for image upload
-image = gr.Image(height=512, width=512)
-
-# Gradio output component for displaying the label
-label = gr.Label()
-
-
-# Example images to demonstrate the model's predictions
-examples = [
-    ['NORMAL2-IM-0222-0001.jpeg'],
-    ['person159_bacteria_747.jpeg'],
-    ['person1618_virus_2805.jpeg'],
-]
-
-
-# Create the Gradio interface
-iface = gr.Interface(
-    fn=predict,              # Function to call for predictions
-    title=title,             # Title of the interface
-    description=description, # Description of the interface
-    article=article,         # Additional article or information
-    inputs=image,            # Input component
-    outputs=label,           # Output component
-    theme="default",         # Theme of the interface
-    examples=examples        # Example images
-)
-
-
-# Launch the Gradio interface
-iface.launch(inline=False)
+# Bismillahir Rahmaanir Raheem 
+# Almadadh Ya Gause Radi Allahu Ta'alah Anh - Ameen 
+
+# Import necessary libraries from fastai and gradio
+from fastai.vision.all import *
+import gradio as gr
+
+
+# Function to determine if an image contains pneumonia
+# Checks if the filename contains 'virus' or 'bacteria'
+def is_pneumonia(x): 
+    return (x.find('virus')!=-1 or x.find('bacteria')!=-1)
+
+
+# Load the trained fastai model for predictions
+learn = load_learner('pneumonia_model.pkl')
+
+
+# Define the possible categories for prediction
+categories = learn.dls.vocab 
+
+
+# Function to make a prediction on an input image
+def predict(img):
+    pred, idx, probs = learn.predict(img)  # Get the prediction, index, and probabilities
+    return dict(zip(categories, map(float, probs)))  # Return the probabilities mapped to categories
+
+
+# Title of the Gradio interface
+title = "Pediatric Pneumonia Chest X-Ray Predictor"
+
+
+# Description of the interface, including model and dataset information
+description = """
+A pediatric pneumonia chest x-ray predictor model trained on the chest-xray-pneumonia dataset using ResNet34 via 
+<a href='http://www.fast.ai/' target='_blank'>fast.ai</a>. The dataset is from: 
+<a href='http://www.kaggle.com/datasets/paultimothymooney/chest-xray-pneumonia' target='_blank'>Chest X-Ray Images (Pneumonia)</a> 
+and the associated scientific journal paper is 
+<a href='http://www.cell.com/cell/fulltext/S0092-8674(18)30154-5' target='_blank'>Identifying Medical Diagnoses and Treatable 
+Diseases by Image-Based Deep Learning</a>. The accuracy of the model is: 87.50%
+"""
+
+
+# Article or additional information to be displayed
+article = """
+<p style='text-align: center'><span style='font-size: 15pt;'>Pediatric Pneumonia Chest X-Ray Predictor. Dr Zakia Salod. 2024. </span></p>
+"""
+
+
+# Gradio input component for image upload
+image = gr.Image(height=512, width=512)
+
+# Gradio output component for displaying the label
+label = gr.Label()
+
+
+# Example images to demonstrate the model's predictions
+examples = [
+    ['NORMAL2-IM-0222-0001.jpeg'],
+    ['person159_bacteria_747.jpeg'],
+    ['person1618_virus_2805.jpeg'],
+]
+
+
+# Create the Gradio interface
+iface = gr.Interface(
+    fn=predict,              # Function to call for predictions
+    title=title,             # Title of the interface
+    description=description, # Description of the interface
+    article=article,         # Additional article or information
+    inputs=image,            # Input component
+    outputs=label,           # Output component
+    theme="default",         # Theme of the interface
+    examples=examples        # Example images
+)
+
+
+# Launch the Gradio interface
+iface.launch(inline=False)