Laeeeq/ECG_Signals_Classification

1import streamlit as st2from ultralytics import YOLO3from PIL import Image4import numpy as np5import matplotlib.pyplot as plt6 7# Load your YOLOv8m model (ensure the path to your model is correct)8model = YOLO('YOLOV8m.pt') 9 10# Streamlit app title11st.title("ECG Signal Classification:")12 13# Upload an image (assumes your ECG signals are visualized as images)14uploaded_file = st.file_uploader("Upload an ECG signal image", type=["jpg", "jpeg", "png"])15 16if uploaded_file is not None:17    image = Image.open(uploaded_file)18    st.image(image, caption="Uploaded ECG Signal", use_column_width=True)19 20    # Perform inference21    results = model(image)  # Make sure the image is compatible with the model (resize if needed)22 23    # Assuming the model provides results with bounding boxes and classes24    if results and len(results) > 0:25        result = results[0]  # Get the first result from the list26 27        # Extract information from the result28        boxes = result.boxes.xyxy.cpu().numpy()  # Bounding box coordinates29        class_ids = result.boxes.cls.cpu().numpy()  # Class IDs (0: ECG HB, 1: History_MI, 2: MI-ECG, 3: Normal-ECG)30        confidences = result.boxes.conf.cpu().numpy()  # Confidence scores31 32        # Class names from your data.yaml file33        class_names = ['ECG HB', 'History_MI', 'MI-ECG', 'Normal-ECG']34 35        # Display results for each detected ECG classification36        for i, box in enumerate(boxes):37            predicted_class = int(class_ids[i])38            confidence_score = confidences[i]39 40            # Check if the predicted class is "Normal-ECG"41            if class_names[predicted_class] == "Normal-ECG":42                classification = "Normal"43            else:44                classification = "Abnormal"45 46            # Display the classification and confidence score47            st.write(f"Prediction: **{classification}**")48            st.write(f"Confidence Score: **{confidence_score:.2f}**")49 50            # Optionally draw bounding boxes on the image using Matplotlib51            fig, ax = plt.subplots()52            ax.imshow(image)53            # Draw bounding box (x1, y1, x2, y2)54            rect = plt.Rectangle(55                (box[0], box[1]), box[2] - box[0], box[3] - box[1],56                linewidth=2, edgecolor='r', facecolor='none'57            )58            ax.add_patch(rect)59            ax.text(box[0], box[1] - 10, f"{classification}: {confidence_score:.2f}",60                    color='white', fontsize=12, backgroundcolor='r')61            st.pyplot(fig)62    else:63        st.write("No valid predictions found.")64else:65    st.write("Please upload an ECG image for classification.")66