File size: 4,174 Bytes
605bd6c |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 |
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()
|