created classes

app.py

@@ -7,26 +7,71 @@ from timm.data import resolve_data_config
 from timm.data.transforms_factory import create_transform
 
 
-model = timm.create_model(
-    'hf-hub:nateraw/resnet50-oxford-iiit-pet',
-    pretrained=True
-)
-model.eval()
-transform = create_transform(**resolve_data_config(model.pretrained_cfg, model=model))
-labels = model.pretrained_cfg['label_names']
-st.title("Pet Image Classification App")
-uploaded_image = st.file_uploader("Choose an image...", type=["jpg", "png", "jpeg"])
-if uploaded_image is not None:
-    st.image(uploaded_image, caption="Uploaded Image", use_column_width=True)
-    st.subheader("Classification Results:")
-    image_data = uploaded_image.read()
-    image = Image.open(io.BytesIO(image_data))
-    output = model(transform(image).unsqueeze(0))
-    probabilities = torch.nn.functional.softmax(output[0], dim=0)
-    values, indices = torch.topk(probabilities, 5)
-    predictions = [
-        {'label': labels[i], 'score': v.item()}
-        for i, v in zip(indices, values)
-    ]
-    for prediction in predictions:
-        st.write(f"- {prediction['label']}: {prediction['score']:.4f}")
+class ImageClassifier(object):
+    def __init__(self, model_name):
+        self.model = timm.create_model(
+            model_name,
+            pretrained=True
+        ).eval()
+
+    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()}
+            for i, v in zip(indices, values)
+        ]
+    
+    def get_output_probabilities(self, image):
+        output = self.classify_image(image)
+        return torch.nn.functional.softmax(output[0], dim=0)
+
+    def classify_image(self, image):
+        transform = self.create_image_transform()
+        return self.model(transform(image).unsqueeze(0))
+    
+    def create_image_transform(self):
+        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):
+        self.title = title
+        self.classifier = ImageClassifier(model_name)
+
+    def render(self):
+        st.title(self.title)
+        uploaded_image = self.get_uploaded_image()
+        if uploaded_image is not None:
+            self.show_image_and_results(uploaded_image)
+
+    def get_uploaded_image(self):
+        return st.file_uploader('Choose an image...', type=['jpg', 'png', 'jpeg'])
+    
+    def show_image_and_results(self, uploaded_image):
+        self.show_uploaded_image(uploaded_image)
+        self.show_classification_results(self.get_image(uploaded_image.read()))
+
+    def show_uploaded_image(self, uploaded_image):
+        st.image(uploaded_image, caption='Uploaded Image', use_column_width=True)
+
+    def show_classification_results(self, image):
+        st.subheader('Classification Results:')
+        self.write_top_5_predictions(image)        
+
+    def write_top_5_predictions(self, image):        
+        for prediction in self.classifier.get_top_5_predictions(image):
+            st.write(f"- {prediction['label']}: {prediction['score']:.4f}")
+
+    def get_image(self, image_data):
+        return Image.open(io.BytesIO(image_data))
+
+
+if __name__ == '__main__':
+    ImageClassificationApp(
+        'Pet Image Classification App', 
+        'hf-hub:nateraw/resnet50-oxford-iiit-pet'
+    ).render()