app.py

import os
import cv2
import numpy as np
import pandas as pd
import tensorflow as tf
from tensorflow.keras.models import load_model
from flask import Flask, request, render_template
from werkzeug.utils import secure_filename
import biosppy.signals.ecg as ecg

app = Flask(__name__)

UPLOAD_FOLDER = 'uploads'
ALLOWED_EXTENSIONS = {'csv', 'png', 'jpg', 'jpeg'}
app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER

# Load the pre-trained model
model = load_model('ecgScratchEpoch2.hdf5')

def allowed_file(filename):
    return '.' in filename and filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

def image_to_signal(image_path):
    # Read and preprocess the image
    img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
    if img is None:
        raise ValueError("Failed to load image")
    
    # Resize to a standard height for consistency (e.g., 500 pixels)
    img = cv2.resize(img, (1000, 500))
    
    # Apply thresholding to isolate the waveform
    _, binary = cv2.threshold(img, 200, 255, cv2.THRESH_BINARY_INV)
    
    # Find contours of the waveform
    contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if not contours:
        raise ValueError("No waveform detected in the image")
    
    # Assume the largest contour is the ECG waveform
    contour = max(contours, key=cv2.contourArea)
    
    # Extract y-coordinates (signal amplitude) along x-axis
    signal = []
    width = img.shape[1]
    for x in range(width):
        column = contour[contour[:, :, 0] == x]
        if len(column) > 0:
            # Take the average y-coordinate if multiple points exist
            y = np.mean(column[:, :, 1])
            signal.append(y)
        else:
            # Interpolate if no point is found
            signal.append(signal[-1] if signal else 0)
    
    # Normalize signal to match expected amplitude range
    signal = np.array(signal)
    signal = (signal - np.min(signal)) / (np.max(signal) - np.min(signal)) * 1000
    
    # Save to CSV
    csv_path = os.path.join(app.config['UPLOAD_FOLDER'], 'converted_signal.csv')
    df = pd.DataFrame(signal, columns=[' Sample Value'])
    df.to_csv(csv_path, index=False)
    
    return csv_path

def model_predict(uploaded_files, model):
    output = []
    for path in uploaded_files:
        APC, NORMAL, LBB, PVC, PAB, RBB, VEB = [], [], [], [], [], [], []
        output.append(str(path))
        result = {"APC": APC, "Normal": NORMAL, "LBB": LBB, "PAB": PAB, "PVC": PVC, "RBB": RBB, "VEB": VEB}
        
        kernel = np.ones((4,4), np.uint8)
        csv = pd.read_csv(path)
        csv_data = csv[' Sample Value']
        data = np.array(csv_data)
        signals = []
        count = 1
        peaks = ecg.christov_segmenter(signal=data, sampling_rate=200)[0]
        indices = []
        
        for i in peaks[1:-1]:
            diff1 = abs(peaks[count - 1] - i)
            diff2 = abs(peaks[count + 1] - i)
            x = peaks[count - 1] + diff1 // 2
            y = peaks[count + 1] - diff2 // 2
            signal = data[x:y]
            signals.append(signal)
            count += 1
            indices.append((x, y))
        
        for signal, index in zip(signals, indices):
            if len(signal) > 10:
                img = np.zeros((128, 128))
                for i in range(len(signal)):
                    img[i, int(signal[i] / 10)] = 255
                img = cv2.dilate(img, kernel, iterations=1)
                img = img.reshape(128, 128, 1)
                prediction = model.predict(np.array([img])).argmax()
                classes = ['Normal', 'APC', 'LBB', 'PAB', 'PVC', 'RBB', 'VEB']
                result[classes[prediction]].append(index)
        
        output.append(result)
    
    return output

@app.route('/', methods=['GET'])
def index():
    return render_template('index.html')

@app.route('/', methods=['POST'])
def upload_file():
    if 'files[]' not in request.files:
        return render_template('index.html', message='No file part')
    
    files = request.files.getlist('files[]')
    file_paths = []
    
    for file in files:
        if file and allowed_file(file.filename):
            filename = secure_filename(file.filename)
            file_path = os.path.join(app.config['UPLOAD_FOLDER'], filename)
            file.save(file_path)
            
            # If the file is an image, convert to CSV
            if filename.rsplit('.', 1)[1].lower() in {'png', 'jpg', 'jpeg'}:
                try:
                    csv_path = image_to_signal(file_path)
                    file_paths.append(csv_path)
                except Exception as e:
                    return render_template('index.html', message=f'Error processing image: {str(e)}')
            else:
                file_paths.append(file_path)
    
    if not file_paths:
        return render_template('index.html', message='No valid files uploaded')
    
    results = model_predict(file_paths, model)
    return render_template('index.html', prediction=results)

if __name__ == '__main__':
    if not os.path.exists(UPLOAD_FOLDER):
        os.makedirs(UPLOAD_FOLDER)
    app.run(debug=True, host='0.0.0.0', port=5000)