pytorchyolo/utils/transforms.py

import torch
import torch.nn.functional as F
import numpy as np

import imgaug.augmenters as iaa
from imgaug.augmentables.bbs import BoundingBox, BoundingBoxesOnImage

from .utils import xywh2xyxy_np
import torchvision.transforms as transforms

import cv2
from PIL import Image

class ImgAug(object):
    def __init__(self, augmentations=[]):
        self.augmentations = augmentations

    def __call__(self, data):
        # Unpack data
        img, boxes = data

        # Convert xywh to xyxy
        boxes = np.array(boxes)
        boxes[:, 1:] = xywh2xyxy_np(boxes[:, 1:])

        # Convert bounding boxes to imgaug
        bounding_boxes = BoundingBoxesOnImage(
            [BoundingBox(*box[1:], label=box[0]) for box in boxes],
            shape=img.shape)

        # Apply augmentations
        img, bounding_boxes = self.augmentations(
            image=img,
            bounding_boxes=bounding_boxes)

        # Clip out of image boxes
        bounding_boxes = bounding_boxes.clip_out_of_image()

        # Convert bounding boxes back to numpy
        boxes = np.zeros((len(bounding_boxes), 5))
        for box_idx, box in enumerate(bounding_boxes):
            # Extract coordinates for unpadded + unscaled image
            x1 = box.x1
            y1 = box.y1
            x2 = box.x2
            y2 = box.y2

            # Returns (x, y, w, h)
            boxes[box_idx, 0] = box.label
            boxes[box_idx, 1] = ((x1 + x2) / 2)
            boxes[box_idx, 2] = ((y1 + y2) / 2)
            boxes[box_idx, 3] = (x2 - x1)
            boxes[box_idx, 4] = (y2 - y1)

        return img, boxes


class RelativeLabels(object):
    def __init__(self, ):
        pass

    def __call__(self, data):
        img, boxes = data
        h, w, _ = img.shape
        boxes[:, [1, 3]] /= w
        boxes[:, [2, 4]] /= h
        return img, boxes


class AbsoluteLabels(object):
    def __init__(self, ):
        pass

    def __call__(self, data):
        img, boxes = data
        h, w, _ = img.shape
        boxes[:, [1, 3]] *= w
        boxes[:, [2, 4]] *= h
        return img, boxes


class PadSquare(ImgAug):
    def __init__(self, ):
        self.augmentations = iaa.Sequential([
            iaa.PadToAspectRatio(
                1.0,
                position="center-center").to_deterministic()
        ])


class ToTensor(object):
    def __init__(self, ):
        pass

    def __call__(self, data):
        img, boxes = data
        # Extract image as PyTorch tensor
        img = transforms.ToTensor()(img)

        bb_targets = torch.zeros((len(boxes), 6))
        bb_targets[:, 1:] = transforms.ToTensor()(boxes)

        return img, bb_targets


class Resize(object):
    def __init__(self, size):
        self.size = size

    def __call__(self, data):
        img, boxes = data
        img = F.interpolate(img.unsqueeze(0), size=self.size, mode="nearest").squeeze(0)
        return img, boxes

# Adjust color brightness strategy 
class adjustGrassColor(object):
    def __init__(self, ):
        self.saturation = 1.25
        self.brightness = 1.15

    def rgb_to_hsv(self, rgb_img):
        # Extract RGB channels
        r, g, b = rgb_img.unbind(0)

        # Get the max and min values across RGB
        max_val, _ = torch.max(rgb_img, dim=0)
        min_val, _ = torch.min(rgb_img, dim=0)
        diff = max_val - min_val

        # Calculate HUE
        h = torch.zeros_like(r)
        mask = max_val == min_val
        h[~mask] = 60.0 * ((g[~mask] - b[~mask]) / diff[~mask] % 6)
        mask = max_val == b
        h[mask] = 60.0 * ((r[mask] - g[mask]) / diff[mask] + 4)
        mask = max_val == g
        h[mask] = 60.0 * ((b[mask] - r[mask]) / diff[mask] + 2)

        # Calculate SATURATION
        s = torch.zeros_like(r)
        mask = max_val != 0
        s[mask] = (diff[mask] / max_val[mask])

        # Calculate VALUE
        v = max_val

        return torch.stack([h, s, v])

    def hsv_to_rgb(self, hsv_img):
        h, s, v = hsv_img.unbind(0)
        c = v * s
        hh = h / 60.0
        x = c * (1 - torch.abs(hh % 2 - 1))
        m = v - c

        segments = hh.to(torch.int32)
        r = c * (segments == 0) + x * (segments == 1) + m * (segments == 4) + m * (segments == 5)
        g = x * (segments == 0) + c * (segments == 1) + c * (segments == 2) + x * (segments == 3)
        b = m * (segments == 0) + m * (segments == 1) + x * (segments == 2) + c * (segments == 3)

        return torch.stack([r, g, b])

    def adjust_grass_color(self, rgb_img):
        hsv_img = self.rgb_to_hsv(rgb_img)
        
        # Adjust saturation
        hsv_img[1] = torch.clamp(hsv_img[1] * self.saturation, 0, 1)
        
        # Adjust brightness = 1.15
        hsv_img[2] = torch.clamp(hsv_img[2] * self.brightness, 0, 1)
    
        return self.hsv_to_rgb(hsv_img)


    def __call__(self, data):
        img, boxes = data

        img = self.adjust_grass_color(img)

        return img, boxes



# Normalize the data to Image Net if weight were trained this way, need to explore darknet code

class Normalize(object):
    def __init__(self, ):
        pass

    def __call__(self, data):
        img, boxes = data

        img = transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]) (img) # Normalize using ImageNet statistics

        return img, boxes




def load_image(path, device):
    image = Image.open(path).convert('RGB')
    transform = transforms.Compose([
        transforms.ToTensor()
    ])
    return transform(image).unsqueeze(0).to(device)







def compute_cumulative_histogram(image):
    bins = torch.linspace(0, 1, 256)
    hist = torch.histc(image, bins=256, min=0, max=1)
    cdf = torch.cumsum(hist, dim=1)
    cdf_normalized = cdf / cdf[:, -1:]
    return cdf_normalized


def match_histogram(source, reference):
    reference_cdf = compute_cumulative_histogram(reference)
    source_cdf = compute_cumulative_histogram(source)


    matched_image = torch.zeros_like(source)

    for b in range(source.size(0)):
        for c in range(source.size(1)):
            for i in range(256):
                source_val = (i + 0.5) / 256
                ref_idx = torch.searchsorted(reference_cdf[b, c], source_cdf[b, c, i])
                ref_val = (ref_idx + 0.5) / 256
                mask = (source[b, c] >= source_val - 0.5/256) & (source[b, c] < source_val + 0.5/256)
                matched_image[b, c, mask] = ref_val

    return matched_image   


# Convert back to PIL Image and save

class balance_image:
    def __init__(self):
        self.imageName =  "C:\\Users\\stevf\\OneDrive\\Documents\\Projects\\PyTorch-YOLOv3\\data\\turfgrass_VOC\\images\\YOLODataset\\images\\20230210_152530.png"
        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
        self.reference_image = load_image(self.imageName , self.device)

    def __call__(self, data):
        matched_image, boxes = data
        transform = transforms.ToPILImage()
        matched_image = match_histogram(matched_image, self.reference_image)
        # matched_image = transform(matched_image.squeeze())
        # matched_image_pil.save('matched.jpg')
        return matched_image , boxes
    


class WhiteBalanceTransform:
    def __call__(self, data):
        img, boxes = data


        # img_np = np.array(img) # Convert PIL Image to numpy array
        # print("TYPE  ",img.shape)
        # Convert to BGR format for OpenCV
        img_bgr = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)

        # Split the channels
        b, g, r = cv2.split(img_bgr)
        
        # Compute the mean of each channel
        r_avg = cv2.mean(r)[0]
        g_avg = cv2.mean(g)[0]
        b_avg = cv2.mean(b)[0]

        # Calculate scaling factors
        k = (r_avg + g_avg + b_avg) / 3
        kr = k / r_avg
        kg = k / g_avg
        kb = k / b_avg

        # White balance correction
        r = cv2.normalize(r * kr, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
        g = cv2.normalize(g * kg, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)
        b = cv2.normalize(b * kb, None, 0, 255, cv2.NORM_MINMAX).astype(np.uint8)

        # Merge channels and convert back to RGB format
        img_balanced = cv2.merge([b, g, r])
        img_balanced = cv2.cvtColor(img_balanced, cv2.COLOR_BGR2RGB)

        # Convert numpy array back to PIL Image
        return np.array(Image.fromarray(img_balanced)), boxes

class correctImage(ImgAug):
    def __init__(self, ):
        self.augmentations = iaa.Sequential(
            [
            iaa.AddToBrightness((-100, 0)),
            ],
        )

class correctImageAspectRatio(ImgAug):
    def __init__(self, ):
        self.augmentations = iaa.Sequential(
            [
            iaa.Resize({"height": 416, "width": "keep-aspect-ratio"}),
            iaa.CropToFixedSize(height=416, width=416)
            ],
        )

class crop(ImgAug):
    def __init__(self, ):
        height, width = 416, 555
        target_width = 416

        crop_left_right = max(0, (width - target_width) // 2)
        crop_top_bottom = max(0, (height - target_width) // 2)  # Assuming you also want height to be 416

        self.augmentations = iaa.Sequential([
            iaa.Crop(px=(crop_top_bottom, crop_left_right, crop_top_bottom, crop_left_right))
        ])


DEFAULT_TRANSFORMS = transforms.Compose([
    AbsoluteLabels(),
    crop(),
    PadSquare(),
    RelativeLabels(),
    ToTensor(),
])