Spaces:
Build error
Build error
| import streamlit as st | |
| import cv2 | |
| from PIL import Image | |
| import numpy as np | |
| import io | |
| from ultralytics import YOLO | |
| # Define a class for YOLO-based crack detection | |
| class YOLOPredictor: | |
| def __init__(self, model_path): | |
| # Initialize the YOLO model with the given model path | |
| self.model = YOLO(model_path) | |
| # Define class labels | |
| self.class_labels = [ | |
| 'alligator_crack_high', 'alligator_crack_low', 'alligator_crack_medium', | |
| 'long_transverse_crack_high', 'long_transverse_crack_low', 'long_transverse_crack_medium' | |
| ] | |
| def predict(self, image): | |
| # Use the model to predict cracks in the image | |
| results = self.model.predict(image) | |
| # Debugging: Print results object to inspect its structure | |
| print("Results:", results) | |
| # Plot the results on the image | |
| result_img = results[0].plot() | |
| # Convert the image from BGR to RGB (OpenCV uses BGR by default) | |
| #result_img = cv2.cvtColor(result_img, cv2.COLOR_BGR2RGB) | |
| # Extract class labels and confidence scores from the results | |
| confidences = [] | |
| classes = [] | |
| for box in results[0].boxes: | |
| confidences.append(box.conf.item()) | |
| classes.append(int(box.cls.item())) # Ensure class index is integer | |
| # Map class indices to class labels | |
| class_labels_mapped = [self.class_labels[idx] for idx in classes] | |
| return result_img, class_labels_mapped, confidences | |
| def main(): | |
| st.title("Thermal Crack Detection Dashboard") | |
| st.write("Upload an image to detect cracks and visualize the results.") | |
| # File uploader for the user to upload an image | |
| uploaded_file = st.file_uploader("Choose an image...", type=["jpg", "jpeg", "png"]) | |
| # Initialize predictor outside the if block | |
| predictor = None | |
| # Check if a file has been uploaded | |
| if uploaded_file is not None: | |
| # Open the uploaded image | |
| image = Image.open(uploaded_file) | |
| # Convert the image to a NumPy array | |
| image = np.array(image) | |
| # Display the uploaded image | |
| st.image(image, caption='Uploaded Image', use_column_width=True) | |
| st.write("Running prediction...") | |
| # Create a YOLOPredictor object with the specified model path | |
| predictor = YOLOPredictor('model/best.pt') | |
| # Predict cracks in the uploaded image | |
| result_img, classes, confidences = predictor.predict(image) | |
| # Display the predicted image | |
| st.image(result_img, caption='Predicted Image', use_column_width=True) | |
| # Display detection classes and confidence scores | |
| if confidences and classes: | |
| st.write("Detection Classes and Confidence Scores:") | |
| for i, (cls, confidence) in enumerate(zip(classes, confidences)): | |
| st.write(f"Crack {i+1} - {cls}, Confidence: {confidence:.2f}") | |
| else: | |
| st.write("No detection classes and confidence scores available.") | |
| # Convert the result image to a format suitable for download | |
| pil_img = Image.fromarray(result_img) | |
| buf = io.BytesIO() | |
| pil_img.save(buf, format="PNG") | |
| byte_im = buf.getvalue() | |
| # Add a download button for the predicted image | |
| st.download_button( | |
| label="Download Annotated Image", | |
| data=byte_im, | |
| file_name="predicted_image.png", | |
| mime="image/png" | |
| ) | |
| # Model information section | |
| st.sidebar.write("### Model Information") | |
| st.sidebar.write("Model: YOLO") | |
| st.sidebar.write("Version: v8") | |
| st.sidebar.write("Accuracy: 95%") | |
| if predictor is not None: | |
| st.sidebar.write(f"Total Classes: {len(predictor.class_labels)}") | |
| st.sidebar.write("### Class Labels:") | |
| for idx, label in enumerate(predictor.class_labels): | |
| st.sidebar.write(f"{idx}: {label}") | |
| else: | |
| st.sidebar.write("Please upload an image to display model information.") | |
| if __name__ == "__main__": | |
| main() | |