Update app.py

app.py

@@ -6,82 +6,58 @@ from model import create_effnetb2_model
 from timeit import default_timer as timer
 from typing import Tuple, Dict
 
-# Setup class names
+# Load class names
 try:
-    with open("class_names.txt", "r") as f:  # reading them in from class_names.txt
+    with open("class_names.txt", "r") as f:
         class_names = [food_name.strip() for food_name in f.readlines()]
 except FileNotFoundError:
-    raise FileNotFoundError("class_names.txt not found. Ensure it exists in the root directory.")
+    raise FileNotFoundError("class_names.txt not found.")
 
 ### 2. Model and transforms preparation ###
-
-# Create model
 try:
-    effnetb2, effnetb2_transforms = create_effnetb2_model(
-        num_classes=101,  # could also use len(class_names)
-    )
+    effnetb2, effnetb2_transforms = create_effnetb2_model(num_classes=101)
 except Exception as e:
     raise Exception(f"Error creating model: {str(e)}")
 
-# Load saved weights
+# Load weights
 try:
     effnetb2.load_state_dict(
         torch.load(
-            f="09_pretrained_effnetb2_feature_extractor_food101.pth",
-            map_location=torch.device("cpu"),  # load to CPU
+            "09_pretrained_effnetb2_feature_extractor_food101.pth",
+            map_location=torch.device("cpu"),
         )
     )
 except FileNotFoundError:
-    raise FileNotFoundError("Model weights file '09_pretrained_effnetb2_feature_extractor_food101.pth' not found.")
+    raise FileNotFoundError("Model weights file not found.")
 except Exception as e:
-    raise Exception(f"Error loading model weights: {str(e)}")
+    raise Exception(f"Error loading weights: {str(e)}")
 
 ### 3. Predict function ###
-
 def predict(img) -> Tuple[Dict, float]:
-    """Transforms and performs a prediction on img and returns prediction and time taken."""
     try:
-        # Start the timer
         start_time = timer()
-
-        # Transform the target image and add a batch dimension
         if img is None:
-            raise ValueError("Input image is None. Please provide a valid image.")
+            raise ValueError("Input image is None.")
         img = effnetb2_transforms(img).unsqueeze(0)
-
-        # Put model into evaluation mode and turn on inference mode
         effnetb2.eval()
         with torch.inference_mode():
-            # Pass the transformed image through the model and turn the prediction logits into prediction probabilities
             pred_probs = torch.softmax(effnetb2(img), dim=1)
-
-        # Create a prediction label and prediction probability dictionary for each prediction class
-        pred_labels_and_probs = {
-            class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))
-        }
-
-        # Calculate the prediction time
+        pred_labels_and_probs = {class_names[i]: float(pred_probs[0][i]) for i in range(len(class_names))}
         pred_time = round(timer() - start_time, 5)
-
-        # Return the prediction dictionary and prediction time
         return pred_labels_and_probs, pred_time
     except Exception as e:
         return {"error": f"Prediction failed: {str(e)}"}, 0.0
 
 ### 4. Gradio app ###
-
-# Create title, description
 title = "FoodVision 101 🍔👁"
-description = "An EfficientNetB2 feature extractor computer vision model to classify images of food into 101 different classes."
+description = "An EfficientNetB2 feature extractor to classify 101 food classes."
 
-# Create examples list from "examples/" directory
 try:
     example_list = [["examples/" + example] for example in os.listdir("examples")]
 except FileNotFoundError:
     example_list = []
-    print("Warning: 'examples/' directory not found. No example images will be loaded.")
+    print("Warning: 'examples/' directory not found.")
 
-# Create Gradio interface
 demo = gr.Interface(
     fn=predict,
     inputs=gr.Image(type="pil"),
@@ -94,5 +70,5 @@ demo = gr.Interface(
     description=description,
 )
 
-# Launch the app with share=True for Hugging Face Spaces
-demo.launch(share=True)
+# Launch without share=True for Hugging Face Spaces
+demo.launch()