Road_Classifier / app.py
Tzetha's picture
Update app.py
9c940e7 verified
raw
history blame
3.42 kB
import streamlit as st
import numpy as np
import json
import os
import cv2
from sklearn.preprocessing import StandardScaler
from sklearn.neighbors import KNeighborsClassifier
# Load COCO JSON dataset
def load_coco_dataset(json_file, image_dir):
with open(json_file, 'r') as f:
data = json.load(f)
images, labels = [], []
label_dict = {}
for annotation in data['annotations']:
image_id = annotation['image_id']
category_id = annotation['category_id']
image_info = next((img for img in data['images'] if img['id'] == image_id), None)
if image_info is None:
continue
image_path = os.path.join(image_dir, image_info['file_name'])
if not os.path.exists(image_path):
continue
image = cv2.imread(image_path)
image = cv2.resize(image, (64, 64))
image = image / 255.0 # Normalize to [0,1]
images.append(image)
labels.append(category_id)
if category_id not in label_dict:
label_dict[category_id] = next(cat['name'] for cat in data['categories'] if cat['id'] == category_id)
return np.array(images), np.array(labels), label_dict
# Load dataset paths
train_json_file = os.path.join("Dataset", "Train", "_annotations.coco.json")
train_image_directory = os.path.join("Dataset", "Train", "images")
test_json_file = os.path.join("Dataset", "Test", "_annotations.coco.json")
test_image_directory = os.path.join("Dataset", "Test", "images")
# Load datasets
X_train, y_train, _ = load_coco_dataset(train_json_file, train_image_directory)
X_test, y_test, label_dict = load_coco_dataset(test_json_file, test_image_directory)
# Flatten images for KNN
X_train = X_train.reshape(X_train.shape[0], -1)
X_test = X_test.reshape(X_test.shape[0], -1)
# Scale the data
scaler = StandardScaler()
X_train = scaler.fit_transform(X_train)
X_test = scaler.transform(X_test)
# Train KNN model
model = KNeighborsClassifier(n_neighbors=5)
model.fit(X_train, y_train)
# Streamlit UI
st.set_page_config(page_title="Road Object Classifier", layout="wide")
st.title("πŸš—πŸšΆπŸΎ KNN Road Object Classifier App")
# Tabs for UI
tabs = st.tabs(["Classifier", "How It Works"])
with tabs[0]:
st.header("Upload an Image for Classification")
uploaded_file = st.file_uploader("Upload an image (Car, People, or Animal)", type=["jpg", "png", "jpeg"])
if uploaded_file:
image = cv2.imdecode(np.frombuffer(uploaded_file.read(), np.uint8), cv2.IMREAD_COLOR)
image = cv2.resize(image, (64, 64))
image_display = cv2.cvtColor(image, cv2.COLOR_BGR2RGB)
image = image / 255.0
image = image.reshape(1, -1)
image = scaler.transform(image)
prediction = model.predict(image)[0]
predicted_label = label_dict.get(prediction, "Unknown")
col1, col2 = st.columns(2)
with col1:
st.image(image_display, caption="Uploaded Image", use_column_width=True)
with col2:
st.markdown(f"## **Predicted Object:** {predicted_label}")
with tabs[1]:
st.header("How the App Works")
st.write("This app uses a K-Nearest Neighbors (KNN) model to classify objects on roads. The model is trained on a dataset labeled using COCO annotations and processes uploaded images to predict whether they contain a car, a person, or an animal.")