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