Skip to content

improve script #1

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
lucgerrits wants to merge 1 commit into
base: main
Choose a base branch
from
+81 −38
Open

improve script #1

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
119 changes: 81 additions & 38 deletions image_enhancer.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,73 +13,116 @@
import os
import cv2
import numpy as np
from PIL import Image, ImageEnhance
from PIL import Image, ImageEnhance, ImageOps, ImageChops
from tqdm import tqdm

# Function to enhance image sharpness, contrast and apply Gaussian blur
def enhance_image(image_path, output_path, sharpness=4, contrast=1.3, blur=3):
"""Enhance image sharpness, contrast, and blur.
import argparse

def is_bw(img):
"""Check if the image is black and white (grayscale)."""
if img.mode == 'L':
return True # The image is in single-channel grayscale
elif img.mode == 'RGB':
r, g, b = img.split()
if ImageChops.difference(r, g).getbbox() is None and ImageChops.difference(g, b).getbbox() is None:
return True # All channels are identical
return False

def apply_clahe(pil_img, clip_limit=1.0):
"""Apply Contrast Limited Adaptive Histogram Equalization (CLAHE) with a specified clipLimit."""
if pil_img.mode == 'RGB':
img = np.array(pil_img)
img_lab = cv2.cvtColor(img, cv2.COLOR_RGB2LAB)
l_channel, a_channel, b_channel = cv2.split(img_lab)
clahe = cv2.createCLAHE(clipLimit=clip_limit, tileGridSize=(8, 8))
l_channel = clahe.apply(l_channel)
img_lab = cv2.merge((l_channel, a_channel, b_channel))
img_rgb = cv2.cvtColor(img_lab, cv2.COLOR_LAB2RGB)
return Image.fromarray(img_rgb)
else:
return pil_img

def apply_equalization(pil_img):
"""Apply histogram equalization to the luminance channel of the image."""
if pil_img.mode == 'RGB':
ycbcr_img = pil_img.convert('YCbCr')
channels = list(ycbcr_img.split())
channels[0] = ImageOps.equalize(channels[0])
ycbcr_img = Image.merge('YCbCr', channels)
return ycbcr_img.convert('RGB')
else:
return ImageOps.equalize(pil_img)

def enhance_image(image_path, output_path, sharpness=4, contrast=1.3, blur=3, equalize=False, equalize_color=True, clip_limit=1.0, use_equalization=False):
"""Enhance image sharpness, contrast, blur, and optionally equalize.

Args:
image_path (str): Path to the input image.
output_path (str): Path to save the enhanced image.
sharpness (float, optional): Sharpness level. Defaults to 4.
contrast (float, optional): Contrast level. Defaults to 1.3.
blur (int, optional): Blur level. Defaults to 3.
equalize (bool, optional): Apply histogram equalization. Defaults to False.
equalize_color (bool, optional): Apply equalization to color images as well. Defaults to True.
clip_limit (float, optional): Clip limit for CLAHE. Defaults to 1.0.
use_equalization (bool, optional): Use equalization instead of CLAHE. Defaults to False.
"""

# Load the image
img = cv2.imread(image_path)

# Convert the image to RGB
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
if img is None:
print(f"Error: Could not read the image file {image_path}. Skipping.")
return

# Convert the image to PIL Image
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
pil_img = Image.fromarray(img)

# Enhance the sharpness
if equalize:
if use_equalization:
pil_img = apply_equalization(pil_img)
elif equalize_color or is_bw(pil_img):
pil_img = apply_clahe(pil_img, clip_limit=clip_limit)

enhancer = ImageEnhance.Sharpness(pil_img)
img_enhanced = enhancer.enhance(sharpness)

# Enhance the contrast
enhancer = ImageEnhance.Contrast(img_enhanced)
img_enhanced = enhancer.enhance(contrast)

# Convert back to OpenCV image (numpy array)
img_enhanced = np.array(img_enhanced)

# Apply a small amount of Gaussian blur
img_enhanced = cv2.GaussianBlur(img_enhanced, (blur, blur), 0)

# Convert back to PIL Image and save
img_enhanced = Image.fromarray(img_enhanced)
img_enhanced.save(output_path)


def process_directory(input_dir, output_dir_name, sharpness=1.2, contrast=1.1):
"""Process all images in a directory, enhancing them.
img_format = os.path.splitext(output_path)[1][1:].upper()
if img_format == 'JPEG' or img_format == 'JPG':
img_enhanced.save(output_path, quality=100)
else:
img_enhanced.save(output_path)

Args:
input_dir (str): Path to the input directory with images.
output_dir_name (str): Name of output directory to save enhanced images.
sharpness (float, optional): Sharpness level. Defaults to 1.2.
contrast (float, optional): Contrast level. Defaults to 1.1.
"""

# Create the output directory inside the input directory
output_dir = os.path.join(input_dir, output_dir_name)
def process_directory(input_dir, output_dir_name, sharpness=1.2, contrast=1.1, blur=3, equalize=False, equalize_color=True, clip_limit=1.0, use_equalization=False):
script_dir = os.getcwd()
output_dir = os.path.join(script_dir, output_dir_name)
if not os.path.exists(output_dir):
os.makedirs(output_dir)

# Get a list of all images in the input directory
image_files = [f for f in os.listdir(input_dir) if f.endswith(".jpg") or f.endswith(".png")]

# Process all images in the input directory
for filename in tqdm(image_files, desc="Processing images"):
input_path = os.path.join(input_dir, filename)
output_path = os.path.join(output_dir, filename)
enhance_image(input_path, output_path, sharpness, contrast)
enhance_image(input_path, output_path, sharpness, contrast, blur, equalize, equalize_color, clip_limit, use_equalization)

def main():
parser = argparse.ArgumentParser(description="Enhance images using OpenCV and PIL.")
parser.add_argument('input_dir', type=str, help='Path to the input directory with images.')
parser.add_argument('output_dir_name', type=str, help='Name of the output directory for enhanced images.')
parser.add_argument('--sharpness', type=float, default=1.2, help='Sharpness level. Default is 1.2.')
parser.add_argument('--contrast', type=float, default=1.1, help='Contrast level. Default is 1.1.')
parser.add_argument('--blur', type=int, default=3, help='Blur level. Default is 3.')
parser.add_argument('--equalize', action='store_true', help='Apply histogram equalization. Default is False.')
parser.add_argument('--equalize_color', action='store_true', help='Apply equalization to color images as well. Default is True.')
parser.add_argument('--clip_limit', type=float, default=1.0, help='Clip limit for CLAHE. Default is 1.0.')
parser.add_argument('--use_equalization', action='store_true', help='Use equalization instead of CLAHE. Default is False.')

args = parser.parse_args()

process_directory(args.input_dir, args.output_dir_name, args.sharpness, args.contrast, args.blur, args.equalize, args.equalize_color, args.clip_limit, args.use_equalization)

if __name__ == "__main__":
main()

# Example usage
process_directory('example/path', 'example/path/newfolder')