import cv2
import numpy as np
import joblib
from skimage.measure import regionprops, label

# Fungsi untuk menghitung eccentricity
def calculate_eccentricity(region):
    if region.minor_axis_length == 0:
        return 0
    else:
        return np.sqrt(1 - (region.minor_axis_length / region.major_axis_length)**2)

# Load model dan LabelEncoder yang telah disimpan
model_path = 'D:/Skripsi/ekstraksi fitur/rf_model.joblib'
label_encoder_path = 'D:/Skripsi/ekstraksi fitur/label_encoder.joblib'
rf_model = joblib.load(model_path)
le = joblib.load(label_encoder_path)

# Fungsi untuk mengekstraksi fitur dari gambar tunggal dengan preprocessing khusus
def extract_features(image_path, label_name):
    img = cv2.imread(image_path, 1)
    if img is None:
        print(f"Gagal membaca file: {image_path}")
        return None

    blue, green, red = cv2.split(img)

    # Thresholding dan preprocessing khusus untuk setiap jenis bakteri
    if label_name == 'Corynebacterium Diphteriae':
        ret, img_bin = cv2.threshold(green, 115, 255, cv2.THRESH_BINARY_INV)
        img_bin = cv2.erode(cv2.dilate(img_bin, None, iterations=1), None, iterations=4)
    elif label_name == 'Mycobacterium Tuberculosis':
        ret, img_bin = cv2.threshold(green, 135, 255, cv2.THRESH_BINARY_INV)
        img_bin = cv2.dilate(img_bin, None, iterations=1)

        num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(img_bin, connectivity=8)
        mask = np.zeros_like(img_bin)
        min_size = 50
        max_size = 400

        for j in range(1, num_labels):
            if min_size <= stats[j, cv2.CC_STAT_AREA] <= max_size:
                mask[labels == j] = 255

        img_bin = mask
    elif label_name == 'Neiseria Gonorroea':
        ret, img_bin = cv2.threshold(green, 0, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
        img_bin = cv2.dilate(img_bin, None, iterations=1)

        num_labels, labels, stats, centroids = cv2.connectedComponentsWithStats(img_bin, connectivity=8)
        mask = np.zeros_like(img_bin)
        min_size = 200
        max_size = 2500

        for j in range(1, num_labels):
            if min_size <= stats[j, cv2.CC_STAT_AREA] <= max_size:
                mask[labels == j] = 255

        img_bin = mask
    elif label_name == 'Staphylococcus Aureus':
        ret, img1 = cv2.threshold(green, 110, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
        img_bin = cv2.erode(cv2.dilate(img1.copy(), None, iterations=1), None, iterations=1)
        img_bin = cv2.morphologyEx(img_bin.copy(), cv2.MORPH_OPEN, None)
    elif label_name == 'Streptococcus Pneumoniae':
        ret, img1 = cv2.threshold(green, 110, 255, cv2.THRESH_BINARY_INV + cv2.THRESH_OTSU)
        img_bin = cv2.erode(img1.copy(), None, iterations=1)
        img_bin = cv2.dilate(img1.copy(), None, iterations=1)

    # Labeling
    labeled_img = label(img_bin)
    regions = regionprops(labeled_img)
    num_objects = len(regions)

    if num_objects == 0:
        print(f"Tidak ditemukan objek pada gambar: {image_path}")
        return None

    # Inisialisasi variabel untuk agregasi
    total_eccentricity = 0
    total_metric = 0
    total_area = 0
    total_perimeter = 0

    for region in regions:
        # Fitur bentuk
        eccentricity = calculate_eccentricity(region)
        area = region.area
        perimeter = region.perimeter
        metric = (4 * np.pi * area) / (perimeter ** 2) if perimeter != 0 else 0

        # Agregasi fitur
        total_eccentricity += eccentricity
        total_metric += metric
        total_area += area
        total_perimeter += perimeter

    # Hitung rata-rata
    mean_eccentricity = total_eccentricity / num_objects
    mean_metric = total_metric / num_objects

    # Membuat array fitur
    features = np.array([[num_objects, mean_eccentricity, mean_metric, total_area, total_perimeter]])

    return features

# Fungsi untuk melakukan prediksi
def predict(image_path, label_name):
    features = extract_features(image_path, label_name)
    if features is None:
        return

    # Melakukan prediksi
    y_pred = rf_model.predict(features)
    label_pred = le.inverse_transform(y_pred)

    print(f"Prediksi untuk gambar {image_path}: {label_pred[0]}")

# Contoh penggunaan
image_path = 'D:/Skripsi/ekstraksi fitur/testing/Corynebacterium Diphteriae/cd100.jpg'
label_name = 'Corynebacterium Diphteriae'  # Sesuaikan dengan label yang benar untuk preprocessing yang tepat
predict(image_path, label_name)
