first commit

.gitignore

@@ -0,0 +1 @@
+__pycache__/

app.py

@@ -0,0 +1,74 @@
+import io
+import torch
+import streamlit as st
+from PIL import Image
+from transformers import AutoImageProcessor, SwinForImageClassification
+
+
+class ImageClassifier(object):
+    def __init__(self, model):
+        self.model = model
+
+    def get_top_5_predictions(self, image):       
+        values, indices = torch.topk(self.get_output_probabilities(image), 5)
+        return [
+            {'label': self.model.config.id2label[i.item()], '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.logits[0], dim=0)
+
+    def classify_image(self, image):
+        image_processor = self.create_image_processor()
+        inputs = image_processor(image, return_tensors='pt')
+        with torch.no_grad():
+            return self.model(**inputs)
+    
+    def create_image_processor(self):
+        return AutoImageProcessor.from_pretrained(self.model.name_or_path)
+
+
+class ImageClassificationApp(object):
+    def __init__(self, title, classifier):
+        self.title = title
+        self.classifier = classifier
+
+    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__':
+    model = SwinForImageClassification.from_pretrained(
+        'microsoft/swin-tiny-patch4-window7-224'
+    )
+    classifier = ImageClassifier(model)
+    ImageClassificationApp(
+        'Swin Image Classification App', 
+        classifier
+    ).render()

requirements.txt

@@ -0,0 +1,3 @@
+torch
+streamlit
+transformers