bug fixes

.gitignore

@@ -0,0 +1,3 @@
+/.DS_Store
+/.venv
+/__pycache__

app.py

@@ -4,6 +4,7 @@ import os
 import torch
 
 from model import createVITModel
+from pathlib import Path
 from timeit import default_timer as timer
 from typing import Tuple, Dict
 
@@ -11,10 +12,15 @@ from typing import Tuple, Dict
 with open("classes.txt", "r") as f: # reading them in from class_names.txt
     class_names = [food_name.strip() for food_name in  f.readlines()]
 
-model, vit_transform = createVITModel()
+model, vit_transform = createVITModel(out_features=len(class_names))
 model.load_state_dict(torch.load('VIT_32_20_003.pth'))
 model = model.to('cpu')
 
+# Create a list of example inputs to our Gradio demo
+examples_source_dir = Path("examples")
+examples_source_paths = list(examples_source_dir.glob("*.jpg"))
+example_list = [str(filepath) for filepath in examples_source_paths]
+
 def predict(img) -> Tuple[Dict, float]:
     """Transforms and performs a prediction on img and returns prediction and time taken.
     """
@@ -56,4 +62,4 @@ demo = gr.Interface(fn=predict, # mapping function from input to output
                     article=article)
 
 # Launch the demo!
-demo.launch(debug=False) # generate a publically shareable URL?
+demo.launch(debug=False) # generate a publically shareable URL?