Spaces:
Sleeping
Sleeping
| import streamlit as st | |
| from PIL import Image | |
| import numpy as np | |
| import matplotlib.pyplot as plt | |
| import pandas as pd | |
| import torch | |
| from collections import Counter | |
| from torchvision.models.detection import fasterrcnn_resnet50_fpn_v2, FasterRCNN_ResNet50_FPN_V2_Weights | |
| from torchvision.utils import draw_bounding_boxes | |
| import io | |
| # Streamlit page configuration | |
| st.set_page_config( | |
| page_title="Object Detector Dashboard", | |
| page_icon="🔍", | |
| layout="wide", | |
| initial_sidebar_state="expanded" | |
| ) | |
| st.sidebar.title("Traffic Monitoring AI") | |
| weights = FasterRCNN_ResNet50_FPN_V2_Weights.DEFAULT | |
| categories = weights.meta["categories"] | |
| img_preprocess = weights.transforms() | |
| def load_model(threshold): | |
| model = fasterrcnn_resnet50_fpn_v2(weights=weights, box_score_thresh=threshold) | |
| model.eval() | |
| return model | |
| def make_prediction(img, model): | |
| img_processed = img_preprocess(img) | |
| prediction = model(img_processed.unsqueeze(0)) | |
| prediction = prediction[0] | |
| prediction["labels"] = [categories[label] for label in prediction["labels"]] | |
| return prediction | |
| def create_image_with_bboxes(img, prediction): | |
| img_tensor = torch.tensor(img) | |
| img_with_bboxes = draw_bounding_boxes( | |
| img_tensor, | |
| boxes=prediction["boxes"], | |
| labels=prediction["labels"], | |
| colors=["Green" if label == "person" else "red" for label in prediction["labels"]], | |
| width=1 | |
| ) | |
| img_with_bboxes_np = img_with_bboxes.detach().numpy().transpose(1, 2, 0) | |
| return img_with_bboxes_np | |
| threshold = st.sidebar.slider("Confidence Threshold", min_value=0.0, max_value=1.0, value=0.1, step=0.01) | |
| st.title("Vehicle Detection") | |
| st.markdown("Upload your images for object detection:") | |
| # Allow users to upload multiple images | |
| uploaded_files = st.file_uploader("Choose images...", type=["jpg", "jpeg", "png"], accept_multiple_files=True) | |
| if uploaded_files: | |
| st.markdown("**Processing uploaded images...**") | |
| # Load the model once | |
| model = load_model(threshold) | |
| all_predictions = [] | |
| cols = st.columns(min(len(uploaded_files), 4)) # Create columns for displaying images | |
| with st.spinner("Processing images, please wait..."): | |
| for i, uploaded_file in enumerate(uploaded_files): | |
| try: | |
| img = Image.open(uploaded_file) | |
| img = img.convert("RGB") # Ensure the image is in RGB format | |
| prediction = make_prediction(img, model) | |
| img_with_bbox = create_image_with_bboxes(np.array(img).transpose(2, 0, 1), prediction) | |
| with cols[i % 4]: | |
| st.header(f"Image {i + 1}: Object Detection Results") | |
| fig = plt.figure(figsize=(5, 5)) | |
| ax = fig.add_subplot(111) | |
| plt.imshow(img_with_bbox) | |
| plt.xticks([], []) | |
| plt.yticks([], []) | |
| ax.spines[["top", "bottom", "right", "left"]].set_visible(True) | |
| st.pyplot(fig, use_container_width=True) | |
| for label in prediction["labels"]: | |
| all_predictions.append({"Image": f"Image {i + 1}", "Label": label}) | |
| except Exception as e: | |
| st.error(f"Error processing image {i + 1}: {e}") | |
| if all_predictions: | |
| image_object_counts = [] | |
| for i in range(1, len(uploaded_files) + 1): | |
| current_image_preds = [pred['Label'] for pred in all_predictions if pred['Image'] == f"Image {i}"] | |
| object_count = Counter(current_image_preds) | |
| for label, count in object_count.items(): | |
| image_object_counts.append({"Image": f"Image {i}", "Label": label, "Count": count}) | |
| df_summary = pd.DataFrame(image_object_counts) | |
| vehicle_categories = ['car', 'bus', 'motorcycle', 'truck', 'train', 'bicycle', 'scooter'] | |
| df_vehicles = df_summary[df_summary['Label'].isin(vehicle_categories)] | |
| if not df_vehicles.empty: | |
| st.header("Combined Vehicle Detection Table for All Images") | |
| st.dataframe(df_vehicles, use_container_width=True) | |