Made labels global

app.py

@@ -8,17 +8,14 @@ from timm.data.transforms_factory import create_transform
 
 
 class ImageClassifier(object):
-    def __init__(self, model_name):
-        self.model = timm.create_model(
-            model_name,
-            pretrained=True
-        ).eval()
+    def __init__(self, model, labels):
+        self.model = model
+        self.labels = labels
 
     def get_top_5_predictions(self, image):        
         values, indices = torch.topk(self.get_output_probabilities(image), 5)
-        labels = self.get_labels()
         return [
-            {'label': labels[i], 'score': v.item()}
+            {'label': self.labels[i], 'score': v.item()}
             for i, v in zip(indices, values)
         ]
     
@@ -27,6 +24,7 @@ class ImageClassifier(object):
         return torch.nn.functional.softmax(output[0], dim=0)
 
     def classify_image(self, image):
+        self.model.eval()
         transform = self.create_image_transform()
         return self.model(transform(image).unsqueeze(0))
     
@@ -34,13 +32,11 @@ class ImageClassifier(object):
         return create_transform(**resolve_data_config(
             self.model.pretrained_cfg, model=self.model))
 
-    def get_labels(self):
-        return self.model.pretrained_cfg['label_names']
 
 class ImageClassificationApp(object):
-    def __init__(self, title, model_name):
+    def __init__(self, title, classifier):
         self.title = title
-        self.classifier = ImageClassifier(model_name)
+        self.classifier = classifier
 
     def render(self):
         st.title(self.title)
@@ -71,7 +67,13 @@ class ImageClassificationApp(object):
 
 
 if __name__ == '__main__':
+    model = timm.create_model(
+        'hf-hub:nateraw/resnet50-oxford-iiit-pet',
+        pretrained=True
+    )
+    labels = model.pretrained_cfg['label_names']
+    classifier = ImageClassifier(model, labels)
     ImageClassificationApp(
         'Pet Image Classification App', 
-        'hf-hub:nateraw/resnet50-oxford-iiit-pet'
+        classifier
     ).render()