main.py

import cv2
import os
import sys
import numpy as np
import argparse

if len(sys.argv) != 2:
    print("Usage: python main.py <path_to_xray_images>")
    sys.exit(1)

image_dir = sys.argv[1]

if not os.path.isdir(image_dir):
    print(f"The directory {image_dir} does not exist.")
    sys.exit(1)

image_files = os.listdir(image_dir)
images = []

for filename in image_files:
    image_path = os.path.join(image_dir, filename)
    image = cv2.imread(image_path)
    if image is not None:
        images.append(image)

def denoise(image):
    median_filtered_image = cv2.medianBlur(image, 3)
    denoised = cv2.fastNlMeansDenoisingColored(median_filtered_image, None, 8, 8, 7, 13)
    return denoised

def sort_corners(corners):
    corners = np.array(corners)
    x_sorted = corners[np.argsort(corners[:, 0]), :]
    left_most = x_sorted[:2, :]
    right_most = x_sorted[2:, :]
    left_most = left_most[np.argsort(left_most[:, 1]), :]
    (tl, bl) = left_most
    right_most = right_most[np.argsort(right_most[:, 1]), :]
    (tr, br) = right_most
    return np.array([tl, tr, bl, br], dtype="float32")

def perspective_warp(image):
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    edges = cv2.Canny(image, 0.5*np.mean(gray), 1.5*np.mean(gray))
    contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
    hull_list = []
    for i in range(len(contours)):
        hull = cv2.convexHull(contours[i])
        hull_list.append(hull)
    max_contour = max(hull_list, key=lambda x: cv2.arcLength(x, True))
    desired_corners = np.array([
    [0, 0],  
    [image.shape[1], 0],   
    [0, image.shape[0]],
    [image.shape[1], image.shape[0]]],
    dtype=np.float32)
    rect = cv2.minAreaRect(max_contour)
    box = cv2.boxPoints(rect)
    box = np.intp(box)
    corners = np.float32(box.reshape(4,2))
    corners = sort_corners(corners)
    matrix = cv2.getPerspectiveTransform(corners, desired_corners)
    warped_image = cv2.warpPerspective(image, matrix, (image.shape[1], image.shape[0]))
    return warped_image

def apply_clahe(image, clip_limit=2.0, tile_grid_size=(8, 8)):
    lab_image = cv2.cvtColor(image, cv2.COLOR_BGR2LAB)
    l, a, b = cv2.split(lab_image)
    clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=tile_grid_size)
    l_clahe = clahe.apply(l)
    # Merge the CLAHE enhanced L channel with the original A and B channels
    lab_image_clahe = cv2.merge((l_clahe, a, b))
    # Convert the LAB image back to BGR
    clahe_image_bgr = cv2.cvtColor(lab_image_clahe, cv2.COLOR_LAB2BGR)
    return clahe_image_bgr

def inpaint_black_circle(image):
    circle_detect = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    mask = np.zeros_like(circle_detect)
    cv2.circle(mask, (190, 47), 35, (255), thickness=-1)
    inpainted_image = cv2.inpaint(image, mask, 3, cv2.INPAINT_TELEA)
    return inpainted_image
        
def main(output_dir='Results'):
    output_dir = 'Results'
    if not os.path.exists(output_dir):
        os.makedirs(output_dir)
    for i, image in enumerate(images, start=1):
        filtered = denoise(image)
        warped_fix = perspective_warp(filtered)
        inpaint = inpaint_black_circle(warped_fix)
        results = apply_clahe(inpaint)
        if i <= 50:
            output_filename = f'im{i:03}-healthy.jpg'
        elif i > 50:
            output_filename = f'im{i:03}-pneumonia.jpg'
        output_filepath = os.path.join(output_dir, output_filename)
        cv2.imwrite(output_filepath, results)

if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="Process X-ray images.")
    parser.add_argument(
        'image_directory', 
        type=str, 
        help='Directory containing X-ray images')
    args = parser.parse_args()
    main()