backend/xiaci.py

import cv2
import numpy as np
from tkinter import Tk, Button, filedialog, Label, messagebox, Label, messagebox, StringVar, OptionMenu
from PIL import Image, ImageTk
import os
import cv2
import numpy as np
from sklearn.mixture import GaussianMixture
import matplotlib.pyplot as plt

def preprocess_image(image_path):
    image = cv2.imread(image_path)
    if image is None:
        raise FileNotFoundError(f"Image not found or unable to read: {image_path}")
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    blurred = cv2.GaussianBlur(gray, (5, 5), 0)
    return blurred, image

def segment_crystal(blurred_image):
    _, binary = cv2.threshold(blurred_image, 30, 255, cv2.THRESH_BINARY)
    contours, _ = cv2.findContours(binary, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    if len(contours) == 0:
        raise ValueError("No crystal detected in the image.")
    crystal_contour = max(contours, key=cv2.contourArea)
    mask = np.zeros_like(blurred_image)
    cv2.drawContours(mask, [crystal_contour], -1, 255, thickness=cv2.FILLED)
    segmented_crystal = cv2.bitwise_and(blurred_image, blurred_image, mask=mask)
    return segmented_crystal, crystal_contour, binary

def detect_defects_GMM(segmented_crystal, crystal_contour):
   # Extract the region of interest (ROI) using the crystal contour
    mask = np.zeros_like(segmented_crystal)
    cv2.drawContours(mask, [crystal_contour], -1, 255, thickness=cv2.FILLED)
    roi = cv2.bitwise_and(segmented_crystal, segmented_crystal, mask=mask)
    
    # Flatten the ROI to a 1D array of pixel intensities
    pixel_intensities = roi.flatten()
    pixel_intensities = pixel_intensities[pixel_intensities > 0]  # Remove background pixels
    
    if len(pixel_intensities) < 10:
        raise ValueError("Not enough data points to fit Gaussian Mixture Model.")
    
    # Reshape for GMM
    X = pixel_intensities.reshape(-1, 1)
    
    # Fit a Gaussian Mixture Model with two components
    gmm = GaussianMixture(n_components=2, random_state=0).fit(X)
    
    # Get the means and covariances of the fitted Gaussians
    means = gmm.means_.flatten()
    covars = gmm.covariances_.flatten()
    
    # Determine which component corresponds to high brightness
    high_brightness_mean_index = np.argmax(means)
    high_brightness_mean = means[high_brightness_mean_index]
    high_brightness_covar = covars[high_brightness_mean_index]
    
    # Calculate the probability density function (PDF) values for each pixel intensity
    pdf_values = gmm.score_samples(X)
    
    # Set a threshold to identify high brightness regions
    threshold = np.percentile(pdf_values, 98)  # Adjust this threshold as needed
    
    # Identify high brightness pixels
    high_brightness_pixels = X[pdf_values >= threshold].flatten()
    
    # Find contours corresponding to high brightness regions
    _, binary_high_brightness = cv2.threshold(roi, int(high_brightness_mean), 255, cv2.THRESH_BINARY)
    contours, _ = cv2.findContours(binary_high_brightness, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    
    defects = []
    for contour in contours:
        perimeter = cv2.arcLength(contour, True)
        if perimeter > 5:
            defects.append(contour)
    
    # Create a black image with the same shape as the original image
    binary_defects = np.zeros_like(segmented_crystal, dtype=np.uint8)
    
    # Draw high brightness regions on the binary defects image
    for y in range(segmented_crystal.shape[0]):
        for x in range(segmented_crystal.shape[1]):
            if mask[y, x] != 0 and segmented_crystal[y, x] >= high_brightness_mean:
                binary_defects[y, x] = 255

     # Generate a GMM plot image
    fig, ax = plt.subplots(figsize=(4.8, 3.6))
    ax.hist(pixel_intensities, bins=50, density=True, alpha=0.5, color='gray', edgecolor='black')
    
    x_vals = np.linspace(0, 255, 1000).reshape(-1, 1)
    log_prob = gmm.score_samples(x_vals)
    responsibilities = gmm.predict_proba(x_vals)
  
    # Convert the plot to an image
    fig.canvas.draw()
    gmm_plot_img = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
    gmm_plot_img = gmm_plot_img.reshape(fig.canvas.get_width_height()[::-1] + (3,))
    plt.close(fig)

    cv2.imshow("gmm",gmm_plot_img)
    cv2.waitKey(10)
    
    return defects, binary_defects

def detect_defects(segmented_crystal, crystal_contour):
    edges = cv2.Canny(segmented_crystal, 30, 90)
    mask = np.zeros_like(edges)
    cv2.drawContours(mask, [crystal_contour], -1, 255, thickness=cv2.FILLED)
    inverted_mask = mask#cv2.bitwise_not(mask)
    defects_edges = cv2.bitwise_and(edges, inverted_mask)
    contours, _ = cv2.findContours(defects_edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    defects = []
    for contour in contours:
        perimeter = cv2.arcLength(contour, True)
        if perimeter > 5:
            defects.append(contour)
    return defects,defects_edges

def calculate_total_area(crystal_contour):
    total_area = cv2.contourArea(crystal_contour)
    return total_area

def score_defects(defects, total_area):
    defect_area = sum(cv2.contourArea(defect) for defect in defects)
    score = (defect_area / total_area) * 100
    return score

def resize_image_for_display(image, width=480, height=360):
    return cv2.resize(image, (width, height))

def process_image():
    global original_image, binary_image, image_with_defects, segmented_crystal, crystal_contour, defects, total_area, score
    
    # Get the directory of the current script
    script_dir = os.path.dirname(os.path.abspath(__file__))
    
    # Open file dialog starting from the script's directory
    image_path = filedialog.askopenfilename(initialdir=script_dir, title="Select an Image",
                                           filetypes=(("JPEG files", "*.jpg *.jpeg"), ("PNG files", "*.png"), ("All files", "*.*")))
    
    if not image_path:
        return
    
    try:
        blurred_image, original_image = preprocess_image(image_path)
        segmented_crystal, crystal_contour, binary_image = segment_crystal(blurred_image)
        selected_algorithm = algorithm_var.get()
        if selected_algorithm == '1':
             defects,binary_image = detect_defects(segmented_crystal, crystal_contour)
        elif selected_algorithm == '2':
             defects,binary_image = detect_defects_GMM(segmented_crystal, crystal_contour)
       
        total_area = calculate_total_area(crystal_contour)
        score = score_defects(defects, total_area)
        
        result_text = (
            f"Detected {len(defects)} defects.\n"
            f"Total Defect Area: {sum(cv2.contourArea(defect) for defect in defects):.2f} pixels\n"
            f"Total Crystal Area: {total_area} pixels\n"
            f"Defect Score (% of Total Area): {score:.2f}%"
        )
        result_label.config(text=result_text)
        
        print(result_text)

        image_with_defects = cv2.drawContours(original_image.copy(), defects, -1, (0, 255, 0), 2)

        update_images()
    except Exception as e:
        messagebox.showerror("Error", str(e))

def update_images():
    images = [
        ("Original Image", resize_image_for_display(original_image)),
        ("Binary Defect Image", resize_image_for_display(binary_image)),
        ("Image with Defects", resize_image_for_display(image_with_defects)),
        ("Segmented Crystal", resize_image_for_display(segmented_crystal))
    ]
    
    for i, (label_text, img) in enumerate(images):
        label = labels[i]
        label.config(text=label_text)
        
        img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
        img_pil = Image.fromarray(img_rgb)
        img_tk = ImageTk.PhotoImage(img_pil)
        
        image_labels[i].config(image=img_tk)
        image_labels[i].image = img_tk

root = Tk()
root.title("Crystal Defect Detection")

original_image = None
binary_image = None
image_with_defects = None
segmented_crystal = None
crystal_contour = None
defects = []
total_area = 0
score = 0

labels = [Label(root, text=f"Image {i+1}") for i in range(4)]
image_labels = [Label(root) for i in range(4)]

# Add a label to display the detection results
result_label = Label(root, text="", justify='left')

button_open = Button(root, text="Open Image", command=process_image)

# Algorithm selection dropdown menu
algorithm_var = StringVar(value='1')  # Default to algorithm 1
algorithm_menu = OptionMenu(root, algorithm_var, '1', '2')
algorithm_menu.config(width=10)

# Arrange labels and images in a 2x2 grid
for i in range(4):
    row = i // 2
    col = i % 2
    labels[i].grid(row=row*2, column=col*2, padx=10, pady=5)
    image_labels[i].grid(row=row*2+1, column=col*2, padx=10, pady=5)

# Place the result label below the images
result_label.grid(row=4, column=0, columnspan=2, padx=10, pady=10)

button_open.grid(row=5, column=0, columnspan=2, padx=10, pady=10)
algorithm_menu.grid(row=5, column=1, columnspan=2, padx=30, pady=10)

root.mainloop()