import os
import cv2
import easyocr

def perform_ocr(image_path):
    # Read the image using OpenCV
    img = cv2.imread(image_path)
    
    # Initialize the OCR reader
    reader = easyocr.Reader(['en'])  # Use English language
    
    # Perform text extraction
    result = reader.readtext(img)
    
    # Extract the text from the result
    text = "\n".join([entry[1] for entry in result])
    
    return text

def main():
    screenshot_dir = 'screenshots'
    output_dir = 'ocr_results'
    
    # Create output directory if it doesn't exist
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    
    # Process all PNG files in the screenshots directory
    for filename in os.listdir(screenshot_dir):
        if filename.endswith('.png'):
            image_path = os.path.join(screenshot_dir, filename)
            
            # Perform OCR
            extracted_text = perform_ocr(image_path)
            
            # Save the extracted text
            output_filename = os.path.splitext(filename)[0] + '.txt'
            output_path = os.path.join(output_dir, output_filename)
            
            with open(output_path, 'w', encoding='utf-8') as f:
                f.write(extracted_text)
            
            print(f"Processed {filename}, saved result to {output_filename}")

if __name__ == "__main__":
    main()