add class names

ui/app.py

@@ -26,6 +26,7 @@ class PlantDiseaseApp:
         self.current_modelName = "CNN from Scratch"
         self.model = self.model_loader.loadModel(self.current_modelName)
         self.flagged_predictions = []
+        self.class_names = utils.get_class_names()
 
     def predict(self, image, modelName, confidence_threshold):
         """

ui/classNames.txt

@@ -0,0 +1,39 @@
+Apple___Apple_scab
+Apple___Black_rot
+Apple___Cedar_apple_rust
+Apple___healthy
+Background_without_leaves
+Blueberry___healthy
+Cherry_(including_sour)_Powdery_mildew
+Cherry_(including_sour)_healthy
+Corn___Cercospora_leaf_spot Gray_leaf_spot
+Corn___Common_rust
+Corn___Northern_Leaf_Blight
+Corn___healthy
+Grape___Black_rot
+Grape__Esca(Black_Measles)
+Grape__Leaf_blight(Isariopsis_Leaf_Spot)
+Grape___healthy
+Orange__Haunglongbing(Citrus_greening)
+Peach___Bacterial_spot
+Peach___healthy
+Pepper,bell__Bacterial_spot
+Pepper,bell__healthy
+Potato___Early_blight
+Potato___Late_blight
+Potato___healthy
+Raspberry___healthy
+Soybean___healthy
+Squash___Powdery_mildew
+Strawberry___Leaf_scorch
+Strawberry___healthy
+Tomato___Bacterial_spot
+Tomato___Early_blight
+Tomato___Late_blight
+Tomato___Leaf_Mold
+Tomato___Septoria_leaf_spot
+Tomato__Spider_mites(Two-spotted_spider_mite)
+Tomato___Target_Spot
+Tomato___Tomato_Yellow_Leaf_Curl_Virus
+Tomato___Tomato_mosaic_virus
+Tomato___healthy

ui/utils.py

@@ -12,11 +12,14 @@ IMAGE_SIZE = (256, 256)
 NORMALIZE_MEAN = [0.485, 0.456, 0.406]
 NORMALIZE_STD = [0.229, 0.224, 0.225]
 
-CLASS_NAMES = []
+CLASS_NAMES_FILE = "classNames.txt" 
 TOP_K_PREDICTIONS = 5
 CONFIDENCE_THRESHOLD = 0.01
 
 
+with open(CLASS_NAMES_FILE, "r") as f:
+    CLASS_NAMES = [line.strip() for line in f.readlines() if line.strip()]
+
 def preprocess_image(image):
     """
     Preprocess image for model input
@@ -37,10 +40,13 @@ def preprocess_image(image):
     return tensor.unsqueeze(0)
 
 
-def postprocess_predictions(logits, class_names=CLASS_NAMES, top_k=TOP_K_PREDICTIONS):
+def postprocess_predictions(logits, class_names=None, top_k=TOP_K_PREDICTIONS):
     """
     Convert logits to formatted predictions
     """
+    if class_names is None:
+        class_names = CLASS_NAMES
+
     probs = torch.nn.functional.softmax(logits, dim=1)
     probs = probs.cpu().detach().numpy()[0]
 
@@ -107,6 +113,10 @@ def create_confidence_label(predictions, top_k=5):
     return "\n".join(lines)
 
 
+def get_class_names():
+    """Return the loaded class names from the txt file."""
+    return CLASS_NAMES
+
 if __name__ == "__main__":
     print("Testing utility functions...")