| 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) |
|
|
|
|
