| import streamlit as st | |
| import requests | |
| from PIL import Image | |
| from io import BytesIO | |
| CLASS_LABELS = { | |
| 0: "airplane", | |
| 1: "bird", | |
| 2: "car", | |
| 3: "cat", | |
| 4: "deer", | |
| 5: "dog", | |
| 6: "horse", | |
| 7: "monkey", | |
| 8: "ship", | |
| 9: "truck", | |
| } | |
| def get_classification(image_bytes): | |
| response = requests.post("http://localhost:5000/classify", files={"file": image_bytes}) | |
| class_id = response.json()["classification"] | |
| return CLASS_LABELS[class_id] | |
| st.title("Image Classification") | |
| st.write("Upload an image to classify") | |
| uploaded_file = st.file_uploader("Choose an image", type=["jpg", "jpeg", "png"]) | |
| if uploaded_file is not None: | |
| image = Image.open(uploaded_file) | |
| st.image(image, caption="Uploaded Image", use_column_width=True) | |
| if st.button("Classify"): | |
| img_bytes = uploaded_file.read() | |
| label = get_classification(img_bytes) | |
| st.write("Prediction:", label) | |