Optimize requirements for Hugging Face Space build

app.py

@@ -13,25 +13,23 @@ with open("class_names.txt", "r") as f:
 
 ### 2. Model and transforms preparation ###
 
-# Create EffNetB2 model
+# Create EffNetB4 model
 effnetb4, effnetb4_transforms = create_effnetb4_model(
-    num_classes=101, # len(class_names) would also work
+    num_classes=101,
 )
 
 # Load saved weights
 effnetb4.load_state_dict(
     torch.load(
         f="model_weights.pth",
-        map_location=torch.device("cpu"),  # load to CPU
+        map_location=torch.device("cpu"),
     )
 )
 
 ### 3. Predict function ###
 
-# Create predict function
 def predict(img) -> Tuple[Dict, float]:
-    """Transforms and performs a prediction on img and returns prediction and time taken.
-    """
+    """Transforms and performs a prediction on img and returns prediction and time taken."""
     # Start the timer
     start_time = timer()
     
@@ -44,7 +42,7 @@ def predict(img) -> Tuple[Dict, float]:
         # Pass the transformed image through the model and turn the prediction logits into prediction probabilities
         pred_probs = torch.softmax(effnetb4(img), dim=1)
     
-    # Create a prediction label and prediction probability dictionary for each prediction class (this is the required format for Gradio's output parameter)
+    # 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
@@ -56,23 +54,27 @@ def predict(img) -> Tuple[Dict, float]:
 ### 4. Gradio app ###
 
 # Create title, description and article strings
-title = "FoodVision"
-description="An EfficientNet-B4 model trained on **Food-101** to classify food images into 101 categories.\n\nCreated by [Rohan Jain](https://www.linkedin.com/in/jaroh23/)"
+title = "FoodVision Big"
+description = "An EfficientNet-B4 model trained on Food-101 to classify food images into 101 categories."
 article = "Created by Rohan Jain - https://www.linkedin.com/in/jaroh23/"
 
 # Create examples list from "examples/" directory
 example_list = [["examples/" + example] for example in os.listdir("examples")]
 
 # Create the Gradio demo
-demo = gr.Interface(fn=predict, # mapping function from input to output
-                    inputs=gr.Image(type="pil"), # what are the inputs?
-                    outputs=[gr.Label(num_top_classes=3, label="Predictions"), # what are the outputs?
-                             gr.Number(label="Prediction time (s)")], # our fn has two outputs, therefore we have two outputs
-                    # Create examples list from "examples/" directory
-                    examples=example_list, 
-                    title=title,
-                    description=description,
-                    article=article)
+demo = gr.Interface(
+    fn=predict,
+    inputs=gr.Image(type="pil"),
+    outputs=[
+        gr.Label(num_top_classes=5, label="Top 5 Predictions"),
+        gr.Number(label="Prediction time (s)")
+    ],
+    examples=example_list, 
+    title=title,
+    description=description,
+    article=article
+)
 
 # Launch the demo
-demo.launch(share=True)
+if __name__ == "__main__":
+    demo.launch()

requirements.txt

@@ -1,3 +1,11 @@
-torch==2.9.0
-torchvision==0.24.0
-gradio==4.36.0
+# Core dependencies
+torch==2.2.2
+torchvision==0.17.2
+gradio==3.50.2
+
+# Optional performance & stability dependencies
+numpy>=1.26.0
+pillow>=9.0.0
+
+# (Optional) helpful for torchvision transforms
+requests>=2.31.0