NEW1.py

import numpy as np
import cv2
from scipy.ndimage import gaussian_filter
from sklearn.cluster import DBSCAN
import matplotlib.pyplot as plt
import os
import re
import csv

def extract_image_info(image_path):
    """
    Extract RPM and image number from the image path.
    """
    rpm_match = re.search(r'\\(\d+)rpm\\', image_path)
    img_num_match = re.search(r'\\(\d+)\.png$', image_path)
    rpm = rpm_match.group(1) if rpm_match else "unknown"
    img_num = img_num_match.group(1) if img_num_match else "unknown"
    return f"{rpm}rpm {img_num}", rpm, img_num

def preprocess_image(image_path, sharpen_sigma=3, sharpen_strength=1.5, final_blur_sigma=0.5):
    """
    Preprocess the image for better arrow detection with controlled sharpening.
    """
    img = cv2.imread(image_path)
    if img is None:
        raise FileNotFoundError(f"Could not load image from {image_path}")
    img_rgb = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    clahe = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8))
    enhanced = clahe.apply(gray)
    gaussian = cv2.GaussianBlur(enhanced, (0, 0), sharpen_sigma)
    sharpened = cv2.addWeighted(enhanced, sharpen_strength, gaussian, -(sharpen_strength - 1), 0)
    final = gaussian_filter(sharpened, sigma=final_blur_sigma)
    return img_rgb, final

def detect_arrowheads(img_rgb, blurred, threshold=15, min_area=3, max_area=80):
    """
    Detect arrowheads using blob detection and intensity analysis.
    """
    params = cv2.SimpleBlobDetector_Params()
    params.filterByArea = True
    params.minArea = min_area
    params.maxArea = max_area
    params.filterByCircularity = True
    params.minCircularity = 0.1
    params.filterByConvexity = True
    params.minConvexity = 0.1
    params.filterByInertia = True
    params.minInertiaRatio = 0.01
    detector = cv2.SimpleBlobDetector_create(params)
    keypoints = detector.detect(blurred)
    arrow_data = []
    for kp in keypoints:
        x, y = int(kp.pt[0]), int(kp.pt[1])
        if y < blurred.shape[0] and x < blurred.shape[1]:
            intensity = blurred[y, x]
            if intensity > threshold:
                r = int(img_rgb[y, x, 0])
                g = int(img_rgb[y, x, 1])
                b = int(img_rgb[y, x, 2])
                arrow_data.append({
                    'x': x,
                    'y': y,
                    'r': r,
                    'g': g,
                    'b': b,
                    'size': kp.size,
                    'intensity': intensity
                })
    return arrow_data

def cluster_arrows(arrow_data, eps=15, min_samples=1):
    """
    Cluster nearby arrowheads to remove duplicates.
    """
    if not arrow_data:
        return []
    points = np.array([[d['x'], d['y']] for d in arrow_data])
    clustering = DBSCAN(eps=eps, min_samples=min_samples).fit(points)
    clustered_data = []
    for cluster_id in set(clustering.labels_):
        if cluster_id == -1:
            continue
        cluster_indices = np.where(clustering.labels_ == cluster_id)[0]
        cluster_points = [arrow_data[i] for i in cluster_indices]
        total_weight = sum(p.get('size', 1) for p in cluster_points)
        avg_x = sum(p['x'] * p.get('size', 1) for p in cluster_points) / total_weight
        avg_y = sum(p['y'] * p.get('size', 1) for p in cluster_points) / total_weight
        avg_r = sum(p['r'] * p.get('size', 1) for p in cluster_points) / total_weight
        avg_g = sum(p['g'] * p.get('size', 1) for p in cluster_points) / total_weight
        avg_b = sum(p['b'] * p.get('size', 1) for p in cluster_points) / total_weight
        avg_intensity = sum(p['intensity'] * p.get('size', 1) for p in cluster_points) / total_weight
        clustered_data.append({
            'x': int(avg_x),
            'y': int(avg_y),
            'r': int(avg_r),
            'g': int(avg_g),
            'b': int(avg_b),
            'intensity': avg_intensity
        })
    return clustered_data

def visualize_results(img_rgb, arrow_data):
    """
    Visualize detected arrowheads on the image.
    """
    viz_img = img_rgb.copy()
    for point in arrow_data:
        cv2.circle(viz_img, (point['x'], point['y']), radius=3, color=(255, 0, 0), thickness=1)
        size = 2
        cv2.line(viz_img, (point['x'] - size, point['y']), (point['x'] + size, point['y']), (255, 0, 0), 1)
        cv2.line(viz_img, (point['x'], point['y'] - size), (point['x'], point['y'] + size), (255, 0, 0), 1)
    return viz_img

def process_images_in_folder(input_folders, output_dir):
    """
    Process all images in the specified folders and save results with added metadata.
    """
    os.makedirs(output_dir, exist_ok=True)
    for folder in input_folders:
        rotation = os.path.basename(folder)
        for height, filename in enumerate(sorted(os.listdir(folder)), start=1):
            if filename.lower().endswith('.png'):
                image_path = os.path.join(folder, filename)
                try:
                    print(f"Processing {image_path}...")
                    img_rgb, preprocessed = preprocess_image(image_path)
                    arrow_data = detect_arrowheads(img_rgb, preprocessed)
                    clustered_data = cluster_arrows(arrow_data)
                    for arrow in clustered_data:
                        arrow['rotation'] = rotation
                        arrow['height'] = height
                    csv_filename = f"{rotation}_height_{height}.csv"
                    csv_file_path = os.path.join(output_dir, csv_filename)
                    with open(csv_file_path, 'w', newline='') as f:
                        fieldnames = ['x', 'y', 'r', 'g', 'b', 'intensity', 'rotation', 'height']
                        writer = csv.DictWriter(f, fieldnames=fieldnames)
                        writer.writeheader()
                        writer.writerows(clustered_data)
                    result_img = visualize_results(img_rgb, clustered_data)
                    viz_filename = f"{rotation}_height_{height}.png"
                    plt.imsave(os.path.join(output_dir, viz_filename), result_img)
                except Exception as e:
                    print(f"Error processing {image_path}: {e}")

def main():
    """
    Main function to process images from multiple folders.
    """
    input_folders = [
        r"C:\Users\Bernardo\Desktop\Project 4\images\50rpm",
        r"C:\Users\Bernardo\Desktop\Project 4\images\75rpm",
        r"C:\Users\Bernardo\Desktop\Project 4\images\100rpm"
    ]
    output_dir = r"C:\Users\Bernardo\Desktop\Project 4\results"
    process_images_in_folder(input_folders, output_dir)
    print(f"Processing complete. Results saved to {output_dir}")

if __name__ == "__main__":
    main()