import rawpy
import imageio
import numpy as np
import cv2
import os
 
def arw_to_fft(arw_path):
    # Get file name and directory
    base_dir = os.path.dirname(arw_path)
    base_name = os.path.splitext(os.path.basename(arw_path))[0]
 
    # Step 1: Read ARW file using rawpy
    with rawpy.imread(arw_path) as raw:
        rgb = raw.postprocess()
 
    # Step 2: Convert RGB to grayscale
    gray = cv2.cvtColor(rgb, cv2.COLOR_RGB2GRAY)
 
    # Step 3: Compute 2D FFT and shift zero frequency to center
    f = np.fft.fft2(gray)
    fshift = np.fft.fftshift(f)
 
    # Step 4: Calculate magnitude spectrum
    magnitude_spectrum = 20 * np.log(np.abs(fshift) + 1)  # Avoid log(0)
 
    # Normalize both images to 8-bit (0-255) for saving
    gray_8bit = cv2.normalize(gray, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
    mag_8bit = cv2.normalize(magnitude_spectrum, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
 
    # Step 5: Save the images
    gray_output_path = os.path.join(base_dir, f"{base_name}_gray.png")
    fft_output_path = os.path.join(base_dir, f"{base_name}_fft.png")
 
    cv2.imwrite(gray_output_path, gray_8bit)
    cv2.imwrite(fft_output_path, mag_8bit)
 
    print(f"Saved grayscale image: {gray_output_path}")
    print(f"Saved FFT spectrum image: {fft_output_path}")
 
 
 
arw_to_fft('C:\image.arw')