Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| import torch | |
| from torchvision import transforms | |
| import pickle | |
| from PIL import Image | |
| import io | |
| def load_image(): | |
| uploaded_image = st.file_uploader(label="Pick a bird image") | |
| if uploaded_image is not None: | |
| image_data = uploaded_image.getvalue() | |
| st.image(image_data) | |
| return Image.open(io.BytesIO(image_data)) | |
| else: | |
| return None | |
| def load_classes(): | |
| with open('bird_classes_dictionary.pkl', 'rb') as f: | |
| class_dict = pickle.load(f) | |
| return class_dict | |
| def load_model(): | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| model = torch.load("efficient1_re", map_location=device) | |
| model.eval() | |
| return model | |
| def predict(model, categories, image): | |
| device = "cuda" if torch.cuda.is_available() else "cpu" | |
| preprocess = transforms.Compose([ | |
| transforms.Resize(224), | |
| transforms.ToTensor() | |
| ]) | |
| image_tensor = preprocess(image).to(device) | |
| image_batch = image_tensor.unsqueeze(0) | |
| with torch.inference_mode(): | |
| output = model(image_batch) | |
| probabilities = torch.nn.functional.softmax(output, dim=1) | |
| probs, categs = torch.topk(probabilities, 1) | |
| st.write(list(categories.keys())[list(categories.values()).index(categs)]) | |
| st.write(f"Prediction probability: {int(probs)}") | |
| def main(): | |
| st.title("Which Bird is this?") | |
| st.markdown("Enter the image of a bird to know its species.") | |
| model = load_model() | |
| categories = load_classes() | |
| image = load_image() | |
| result = st.button('Run on image') | |
| if result: | |
| #st.write('Calculating results...') | |
| predict(model, categories, image) | |
| if __name__ == "__main__": | |
| main() |