Diff of /utils/classification.py [000000] .. [addb71]

Switch to unified view

a b/utils/classification.py
1
import cv2
2
import numpy as np
3
from tensorflow.keras.models import load_model
4
import os
5
6
# Load model with input shape validation
7
CLASSIFIER_PATH = os.path.join(os.path.dirname(__file__), '..', 'app', 'model', 'classifier_model.h5')
8
classifier_model = load_model(CLASSIFIER_PATH)
9
10
# Get model's expected input shape
11
if len(classifier_model.input_shape) == 4:  # CNN (e.g., (None, 120, 120, 1))
12
    TARGET_SIZE = classifier_model.input_shape[1:3]
13
    GRAYSCALE = classifier_model.input_shape[-1] == 1
14
else:  # Flattened input (e.g., (None, 57600))
15
    TARGET_SIZE = (int(np.sqrt(classifier_model.input_shape[1]/3)), )*2 if not GRAYSCALE else (int(np.sqrt(classifier_model.input_shape[1])),)*2
16
    GRAYSCALE = classifier_model.input_shape[1] % 3 != 0
17
18
print(f"Model expects: {TARGET_SIZE} {'grayscale' if GRAYSCALE else 'RGB'}")
19
20
def preprocess_image(image_path):
21
    """Universal preprocessor adapting to model requirements"""
22
    # Read image
23
    img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE if GRAYSCALE else cv2.IMREAD_COLOR)
24
    if img is None:
25
        raise ValueError("Invalid image file")
26
27
    # Resize and normalize
28
    img = cv2.resize(img, TARGET_SIZE)
29
    img = img.astype(np.float32) / 255.0
30
31
    # Format for model
32
    if len(classifier_model.input_shape) == 4:
33
        return np.expand_dims(img, axis=(0, -1)) if GRAYSCALE else np.expand_dims(img, axis=0)
34
    else:
35
        return img.reshape(1, -1)  # Flatten
36
37
def classify_disease(image_path):
38
    try:
39
        # Preprocess
40
        input_img = preprocess_image(image_path)
41
        
42
        # Predict
43
        preds = classifier_model.predict(input_img)
44
        classes = ['Normal', 'COVID-19', 'Pneumonia', 'Tuberculosis']
45
        
46
        # Interpret results
47
        predicted_idx = np.argmax(preds)
48
        confidence = float(np.max(preds))
49
        severity = "Severe" if confidence > 0.9 else "Moderate" if confidence > 0.7 else "Mild"
50
        
51
        return classes[predicted_idx], confidence, severity
52
        
53
    except Exception as e:
54
        raise ValueError(f"Classification failed: {str(e)}")
55
56
57