Spaces:
Sleeping
Sleeping
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.") | |