Road_Classifier / app.py
Tzetha's picture
Update app.py
b5676cd verified
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 = "Dataset/Train/_annotations.coco.json"
train_image_directory = "Dataset/Train/Images"
test_json_file = "Dataset/Test/_annotations.coco.json"
test_image_directory = "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.")