File size: 1,262 Bytes
0720d3d |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 |
import cv2
import numpy as np
from PIL import Image
from transformers import BitImageProcessor
class CustomBitImageProcessor(BitImageProcessor):
def __init__(self, do_hist_equalization=False, **kwargs):
super().__init__(**kwargs)
self.do_hist_equalization = do_hist_equalization
def apply_histogram_equalization(self, image):
"""
Apply histogram equalization to an image.
"""
image_array = np.array(image)
if len(image_array.shape) == 2: # Grayscale
equalized = cv2.equalizeHist(image_array)
else: # RGB
channels = cv2.split(image_array)
equalized_channels = [cv2.equalizeHist(channel) for channel in channels]
equalized = cv2.merge(equalized_channels)
return Image.fromarray(equalized)
def preprocess(self, images, **kwargs):
"""
Apply custom preprocessing, including histogram equalization.
"""
if self.do_hist_equalization:
if isinstance(images, list):
images = [self.apply_histogram_equalization(image) for image in images]
else:
images = self.apply_histogram_equalization(images)
return super().preprocess(images, **kwargs)
|