import cv2
import numpy as np
import pytesseract
import re

# Path to the Tesseract executable
pytesseract.pytesseract.tesseract_cmd = r'C:\Program Files\Tesseract-OCR\tesseract.exe'

# Load the image
image_path = "./kavanur.jpg"
image = cv2.imread(image_path)
if image is None:
    print(f"Error: Image at path {image_path} not found.")
    exit(1)

# Convert the image to HSV color space
hsv_image = cv2.cvtColor(image, cv2.COLOR_BGR2HSV)

# Define color ranges in HSV


colors = {
    'White': {'lower': np.array([0, 0, 200]), 'upper': np.array([180, 25, 255])}
}

# Function to increase image contrast
def increase_contrast(image):
    lab = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
    l, a, b = cv2.split(lab)
    clahe = cv2.createCLAHE(clipLimit=3.0, tileGridSize=(8, 8))
    cl = clahe.apply(l)
    limg = cv2.merge((cl, a, b))
    enhanced_image = cv2.cvtColor(limg, cv2.COLOR_LAB2BGR)
    return enhanced_image

# Initialize an empty list to store extracted numbers
all_numbers = []

for color_name, ranges in colors.items():
    try:
        # Create a mask for the color
        mask = cv2.inRange(hsv_image, ranges['lower'], ranges['upper'])

        # Apply the mask to the image
        masked_image = cv2.bitwise_and(image, image, mask=mask)

        # Increase contrast
        contrasted_image = increase_contrast(masked_image)

        # Convert the masked image to grayscale
        gray_image = cv2.cvtColor(contrasted_image, cv2.COLOR_BGR2GRAY)

        # Enhance the image for better OCR results
        gray_image = cv2.GaussianBlur(gray_image, (5, 5), 0)
        _, thresh_image = cv2.threshold(gray_image, 0, 255, cv2.THRESH_BINARY + cv2.THRESH_OTSU)

        # Find contours in the thresholded image
        contours, _ = cv2.findContours(thresh_image, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)

        for contour in contours:
            x, y, w, h = cv2.boundingRect(contour)
            if w > 5 and h > 5:  # Filter out small contours
                # Focus on the center of the rectangle
                center_x, center_y = x + w // 4, y + h // 4
                center_w, center_h = w // 2, h // 2
                roi = thresh_image[center_y:center_y+center_h, center_x:center_x+center_w]

                # Apply OCR on the region of interest (ROI)
                custom_config = r'--oem 3 --psm 6'
                text = pytesseract.image_to_string(roi, config=custom_config)

                # Extract numbers using regex
                numbers = [int(s) for s in re.findall(r'\b\d+\b', text)]

                if numbers:
                    all_numbers.extend(numbers)
                    # Draw a rectangle around the detected area
                    cv2.rectangle(image, (center_x, center_y), (center_x+center_w, center_y+center_h), (0, 255, 0), 2)
                    cv2.putText(image, str(numbers), (center_x, center_y-10), cv2.FONT_HERSHEY_SIMPLEX, 0.9, (36, 255, 12), 2)

    except Exception as e:
        print(f"Error processing {color_name} color: {e}")

# Save the final image with rectangles
output_path = './output_center_numbers.png'
cv2.imwrite(output_path, image)

# Display the numbers
print("Plot numbers found in the image:")
print(sorted(set(all_numbers)))