modules/dataset_loader.py

from torchvision.models.detection import keypointrcnn_resnet50_fpn
from torchvision.models.detection.faster_rcnn import FastRCNNPredictor
from torchvision.models.detection.keypoint_rcnn import KeypointRCNNPredictor
from torchvision.models.detection import KeypointRCNN_ResNet50_FPN_Weights
import random
import torch
from torch.utils.data import Dataset
import torchvision.transforms.functional as F
import numpy as np
from torch.utils.data.dataloader import default_collate
import cv2
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader, Subset, ConcatDataset
import streamlit as st
from modules.utils import object_dict, arrow_dict, resize_boxes, resize_keypoints

class RandomCrop:
    def __init__(self, new_size=(1333,800),crop_fraction=0.5, min_objects=4):
        self.crop_fraction = crop_fraction
        self.min_objects = min_objects
        self.new_size = new_size

    def __call__(self, image, target):
        new_w1, new_h1 = self.new_size
        w, h = image.size
        new_w = int(w * self.crop_fraction)
        new_h = int(new_w*new_h1/new_w1)

        i=0
        for i in range(4): 
          if new_h >= h:
            i += 0.05
            new_w = int(w * (self.crop_fraction - i))
            new_h = int(new_w*new_h1/new_w1)
          if new_h < h:
            continue

        if new_h >= h:
          return image, target

        boxes = target["boxes"]
        if 'keypoints' in target:
            keypoints = target["keypoints"]
        else:
            keypoints = []
            for i in range(len(boxes)):
                keypoints.append(torch.zeros((2,3)))
        

        # Attempt to find a suitable crop region
        success = False
        for _ in range(100):  # Max 100 attempts to find a valid crop
            top = random.randint(0, h - new_h)
            left = random.randint(0, w - new_w)
            crop_region = [left, top, left + new_w, top + new_h]

            # Check how many objects are fully contained in this region
            contained_boxes = []
            contained_keypoints = []
            for box, kp in zip(boxes, keypoints):
                if box[0] >= crop_region[0] and box[1] >= crop_region[1] and box[2] <= crop_region[2] and box[3] <= crop_region[3]:
                    # Adjust box and keypoints coordinates
                    new_box = box - torch.tensor([crop_region[0], crop_region[1], crop_region[0], crop_region[1]])
                    new_kp = kp - torch.tensor([crop_region[0], crop_region[1], 0])
                    contained_boxes.append(new_box)
                    contained_keypoints.append(new_kp)

            if len(contained_boxes) >= self.min_objects:
                success = True
                break

        if success:
            # Perform the actual crop
            image = F.crop(image, top, left, new_h, new_w)
            target["boxes"] = torch.stack(contained_boxes) if contained_boxes else torch.zeros((0, 4))
            if 'keypoints' in target:
                target["keypoints"] = torch.stack(contained_keypoints) if contained_keypoints else torch.zeros((0, 2, 4))

        return image, target


class RandomFlip:
    def __init__(self, h_flip_prob=0.5, v_flip_prob=0.5):
        """
        Initializes the RandomFlip with probabilities for flipping.

        Parameters:
        - h_flip_prob (float): Probability of applying a horizontal flip to the image.
        - v_flip_prob (float): Probability of applying a vertical flip to the image.
        """
        self.h_flip_prob = h_flip_prob
        self.v_flip_prob = v_flip_prob

    def __call__(self, image, target):
        """
        Applies random horizontal and/or vertical flip to the image and updates target data accordingly.

        Parameters:
        - image (PIL Image): The image to be flipped.
        - target (dict): The target dictionary containing 'boxes' and 'keypoints'.

        Returns:
        - PIL Image, dict: The flipped image and its updated target dictionary.
        """
        if random.random() < self.h_flip_prob:
            image = F.hflip(image)
            w, _ = image.size  # Get the new width of the image after flip for bounding box adjustment
            # Adjust bounding boxes for horizontal flip
            for i, box in enumerate(target['boxes']):
                xmin, ymin, xmax, ymax = box
                target['boxes'][i] = torch.tensor([w - xmax, ymin, w - xmin, ymax], dtype=torch.float32)

            # Adjust keypoints for horizontal flip
            if 'keypoints' in target:
                new_keypoints = []
                for keypoints_for_object in target['keypoints']:
                    flipped_keypoints_for_object = []
                    for kp in keypoints_for_object:
                        x, y = kp[:2]
                        new_x = w - x
                        flipped_keypoints_for_object.append(torch.tensor([new_x, y] + list(kp[2:])))
                    new_keypoints.append(torch.stack(flipped_keypoints_for_object))
                target['keypoints'] = torch.stack(new_keypoints)

        if random.random() < self.v_flip_prob:
            image = F.vflip(image)
            _, h = image.size  # Get the new height of the image after flip for bounding box adjustment
            # Adjust bounding boxes for vertical flip
            for i, box in enumerate(target['boxes']):
                xmin, ymin, xmax, ymax = box
                target['boxes'][i] = torch.tensor([xmin, h - ymax, xmax, h - ymin], dtype=torch.float32)

            # Adjust keypoints for vertical flip
            if 'keypoints' in target:
                new_keypoints = []
                for keypoints_for_object in target['keypoints']:
                    flipped_keypoints_for_object = []
                    for kp in keypoints_for_object:
                        x, y = kp[:2]
                        new_y = h - y
                        flipped_keypoints_for_object.append(torch.tensor([x, new_y] + list(kp[2:])))
                    new_keypoints.append(torch.stack(flipped_keypoints_for_object))
                target['keypoints'] = torch.stack(new_keypoints)

        return image, target
    

class RandomRotate:
    def __init__(self, max_rotate_deg=20, rotate_proba=0.3):
        """
        Initializes the RandomRotate with a maximum rotation angle and probability of rotating.

        Parameters:
        - max_rotate_deg (int): Maximum degree to rotate the image.
        - rotate_proba (float): Probability of applying rotation to the image.
        """
        self.max_rotate_deg = max_rotate_deg
        self.rotate_proba = rotate_proba

    def __call__(self, image, target):
        """
        Randomly rotates the image and updates the target data accordingly.

        Parameters:
        - image (PIL Image): The image to be rotated.
        - target (dict): The target dictionary containing 'boxes', 'labels', and 'keypoints'.

        Returns:
        - PIL Image, dict: The rotated image and its updated target dictionary.
        """
        if random.random() < self.rotate_proba:
            angle = random.uniform(-self.max_rotate_deg, self.max_rotate_deg)
            image = F.rotate(image, angle, expand=False, fill=200)

            # Rotate bounding boxes
            w, h = image.size
            cx, cy = w / 2, h / 2
            boxes = target["boxes"]
            new_boxes = []
            for box in boxes:
                new_box = self.rotate_box(box, angle, cx, cy)
                new_boxes.append(new_box)
            target["boxes"] = torch.stack(new_boxes)

            # Rotate keypoints
            if 'keypoints' in target:
                new_keypoints = []
                for keypoints in target["keypoints"]:
                    new_kp = self.rotate_keypoints(keypoints, angle, cx, cy)
                    new_keypoints.append(new_kp)
                target["keypoints"] = torch.stack(new_keypoints)

        return image, target

    def rotate_box(self, box, angle, cx, cy):
        """
        Rotates a bounding box by a given angle around the center of the image.
        """
        x1, y1, x2, y2 = box
        corners = torch.tensor([
            [x1, y1],
            [x2, y1],
            [x2, y2],
            [x1, y2]
        ])
        corners = torch.cat((corners, torch.ones(corners.shape[0], 1)), dim=1)
        M = cv2.getRotationMatrix2D((cx, cy), angle, 1)
        corners = torch.matmul(torch.tensor(M, dtype=torch.float32), corners.T).T
        x_ = corners[:, 0]
        y_ = corners[:, 1]
        x_min, x_max = torch.min(x_), torch.max(x_)
        y_min, y_max = torch.min(y_), torch.max(y_)
        return torch.tensor([x_min, y_min, x_max, y_max], dtype=torch.float32)

    def rotate_keypoints(self, keypoints, angle, cx, cy):
        """
        Rotates keypoints by a given angle around the center of the image.
        """
        new_keypoints = []
        for kp in keypoints:
            x, y, v = kp
            point = torch.tensor([x, y, 1])
            M = cv2.getRotationMatrix2D((cx, cy), angle, 1)
            new_point = torch.matmul(torch.tensor(M, dtype=torch.float32), point)
            new_keypoints.append(torch.tensor([new_point[0], new_point[1], v], dtype=torch.float32))
        return torch.stack(new_keypoints)

def rotate_90_box(box, angle, w, h):
    x1, y1, x2, y2 = box
    if angle == 90:
        return torch.tensor([y1,h-x2,y2,h-x1])
    elif angle == 270 or angle == -90:
        return torch.tensor([w-y2,x1,w-y1,x2])
    else:
        print("angle not supported")

def rotate_90_keypoints(kp, angle, w, h):
    # Extract coordinates and visibility from each keypoint tensor
    x1, y1, v1 = kp[0][0], kp[0][1], kp[0][2]
    x2, y2, v2 = kp[1][0], kp[1][1], kp[1][2]
    # Swap x and y coordinates for each keypoint
    if angle == 90:
        new = [[y1, h-x1, v1], [y2, h-x2, v2]]
    elif angle == 270 or angle == -90:
        new = [[w-y1, x1, v1], [w-y2, x2, v2]]

    return torch.tensor(new, dtype=torch.float32)
    

def rotate_vertical(image, target):
    # Rotate the image and target if the image is vertical
    new_boxes = []
    angle = random.choice([-90,90])
    image = F.rotate(image, angle, expand=True, fill=200)
    for box in target["boxes"]:
        new_box = rotate_90_box(box, angle, image.size[0], image.size[1])
        new_boxes.append(new_box)
    target["boxes"] = torch.stack(new_boxes)
    
    if 'keypoints' in target:
        new_kp = []  
        for kp in target['keypoints']:                   
            new_key = rotate_90_keypoints(kp, angle, image.size[0], image.size[1])
            new_kp.append(new_key)
        target['keypoints'] = torch.stack(new_kp)
    return image, target


import torchvision.transforms.functional as F
import torch

def resize_and_pad(image, target, new_size=(1333, 800)):
    original_size = image.size
    # Calculate scale to fit the new size while maintaining aspect ratio
    scale = min(new_size[0] / original_size[0], new_size[1] / original_size[1])
    new_scaled_size = (int(original_size[0] * scale), int(original_size[1] * scale))

    # Resize image to new scaled size
    image = F.resize(image, (new_scaled_size[1], new_scaled_size[0]))

    # Calculate padding to center the image
    pad_left = (new_size[0] - new_scaled_size[0]) // 2
    pad_top = (new_size[1] - new_scaled_size[1]) // 2
    pad_right = new_size[0] - new_scaled_size[0] - pad_left
    pad_bottom = new_size[1] - new_scaled_size[1] - pad_top

    # Pad the resized image to make it exactly the desired size
    image = F.pad(image, (pad_left, pad_top, pad_right, pad_bottom), fill=250, padding_mode='constant')

    # Adjust bounding boxes
    target['boxes'] = resize_boxes(target['boxes'], original_size, new_scaled_size)
    target['boxes'][:, 0::2] += pad_left
    target['boxes'][:, 1::2] += pad_top

    # Adjust keypoints if they exist in the target
    if 'keypoints' in target:
        for i in range(len(target['keypoints'])):
            target['keypoints'][i] = resize_keypoints(target['keypoints'][i], original_size, new_scaled_size)
            target['keypoints'][i][:, 0] += pad_left
            target['keypoints'][i][:, 1] += pad_top

    return image, target

class BPMN_Dataset(Dataset):
    def __init__(self, annotations, transform=None, crop_transform=None, crop_prob=0.3, rotate_90_proba=0.2, 
                 flip_transform=None, rotate_transform=None, new_size=(1333,1333), keep_ratio=0.1, resize=True, model_type='object'):
        self.annotations = annotations
        print(f"Loaded {len(self.annotations)} annotations.")
        self.transform = transform
        self.crop_transform = crop_transform
        self.crop_prob = crop_prob
        self.flip_transform = flip_transform
        self.rotate_transform = rotate_transform
        self.resize = resize
        self.new_size = new_size
        self.keep_ratio = keep_ratio
        self.model_type = model_type
        if model_type == 'object':
            self.dict = object_dict
        elif model_type == 'arrow':
            self.dict = arrow_dict
        self.rotate_90_proba = rotate_90_proba

    def __len__(self):
        return len(self.annotations)

    def __getitem__(self, idx):
        annotation = self.annotations[idx]
        image = annotation.img.convert("RGB")
        boxes = torch.tensor(np.array(annotation.boxes_ltrb), dtype=torch.float32)
        labels_names = [ann for ann in annotation.categories]

        # Only keep the labels, boxes and keypoints that are in the class_dict
        kept_indices = [i for i, ann in enumerate(annotation.categories) if ann in self.dict.values()]
        boxes = boxes[kept_indices]
        labels_names = [ann for i, ann in enumerate(labels_names) if i in kept_indices]
        # Replace any subprocess by task
        labels_names = ['task' if ann == 'subProcess' else ann for ann in labels_names]

        labels_id = torch.tensor([(list(self.dict.values()).index(ann)) for ann in labels_names], dtype=torch.int64)

        # Initialize keypoints tensor
        max_keypoints = 2
        keypoints = torch.zeros((len(labels_id), max_keypoints, 3), dtype=torch.float32)

        ii = 0
        for i, ann in enumerate(annotation.annotations):
            # Only keep the keypoints that are in the kept indices
            if i not in kept_indices:
                continue
            if ann.category in ["sequenceFlow", "messageFlow", "dataAssociation"]:
                # Fill the keypoints tensor for this annotation, mark as visible (1)
                kp = np.array(ann.keypoints, dtype=np.float32).reshape(-1, 3)
                kp = kp[:,:2]
                visible = np.ones((kp.shape[0], 1), dtype=np.float32)
                kp = np.hstack([kp, visible])
                keypoints[ii, :kp.shape[0], :] = torch.tensor(kp, dtype=torch.float32)
                ii += 1

        area = (boxes[:, 3] - boxes[:, 1]) * (boxes[:, 2] - boxes[:, 0])

        if self.model_type == 'object':        
            target = {
                "boxes": boxes,
                "labels": labels_id,
                #"area": area,
            }
        elif self.model_type == 'arrow':
            target = {
                "boxes": boxes,
                "labels": labels_id,
                #"area": area,
                "keypoints": keypoints,
            }

        # Randomly apply flip transform
        if self.flip_transform:
            image, target = self.flip_transform(image, target)

        # Randomly apply rotate transform
        if self.rotate_transform:
            image, target = self.rotate_transform(image, target)

        # Randomly apply the custom cropping transform
        if self.crop_transform and random.random() < self.crop_prob:
            image, target = self.crop_transform(image, target)
            
        # Rotate vertical image
        if random.random() < self.rotate_90_proba:
            image, target = rotate_vertical(image, target)

        if self.resize:
            if random.random() < self.keep_ratio:
                # Center and pad the image while keeping the aspect ratio
                image, target = resize_and_pad(image, target, self.new_size)
            else:
                target['boxes'] = resize_boxes(target['boxes'], (image.size[0],image.size[1]), self.new_size)
                if 'area' in target:
                    target['area'] = (target['boxes'][:, 3] - target['boxes'][:, 1]) * (target['boxes'][:, 2] - target['boxes'][:, 0])
                if 'keypoints' in target:
                    for i in range(len(target['keypoints'])):
                        target['keypoints'][i] = resize_keypoints(target['keypoints'][i], (image.size[0],image.size[1]), self.new_size)
                image = F.resize(image, (self.new_size[1], self.new_size[0]))

        return self.transform(image), target


def collate_fn(batch):
    """
    Custom collation function for DataLoader that handles batches of images and targets.

    This function ensures that images are properly batched together using PyTorch's default collation,
    while keeping the targets (such as bounding boxes and labels) in a list of dictionaries, 
    as each image might have a different number of objects detected.

    Parameters:
    - batch (list): A list of tuples, where each tuple contains an image and its corresponding target dictionary.

    Returns:
    - Tuple containing:
      - Tensor: Batched images.
      - List of dicts: Targets corresponding to each image in the batch.
    """
    images, targets = zip(*batch)  # Unzip the batch into separate lists for images and targets.

    # Batch images using the default collate function which handles tensors, numpy arrays, numbers, etc.
    images = default_collate(images)

    return images, targets



def create_loader(new_size,transformation, annotations1, annotations2=None, 
                  batch_size=4, crop_prob=0.2, crop_fraction=0.7, min_objects=3, 
                  h_flip_prob=0.3, v_flip_prob=0.3, max_rotate_deg=20, rotate_90_proba=0.2, rotate_proba=0.3, 
                  seed=42, resize=True, keep_ratio=0.1, model_type = 'object'):
    """
    Creates a DataLoader for BPMN datasets with optional transformations and concatenation of two datasets.

    Parameters:
    - transformation (callable): Transformation function to apply to each image (e.g., normalization).
    - annotations1 (list): Primary list of annotations.
    - annotations2 (list, optional): Secondary list of annotations to concatenate with the first.
    - batch_size (int): Number of images per batch.
    - crop_prob (float): Probability of applying the crop transformation.
    - crop_fraction (float): Fraction of the original width to use when cropping.
    - min_objects (int): Minimum number of objects required to be within the crop.
    - h_flip_prob (float): Probability of applying horizontal flip.
    - v_flip_prob (float): Probability of applying vertical flip.
    - seed (int): Seed for random number generators for reproducibility.
    - resize (bool): Flag indicating whether to resize images after transformations.

    Returns:
    - DataLoader: Configured data loader for the dataset.
    """

    # Initialize custom transformations for cropping and flipping
    custom_crop_transform = RandomCrop(new_size,crop_fraction, min_objects)
    custom_flip_transform = RandomFlip(h_flip_prob, v_flip_prob)
    custom_rotate_transform = RandomRotate(max_rotate_deg, rotate_proba)

    # Create the primary dataset
    dataset = BPMN_Dataset(
        annotations=annotations1,
        transform=transformation,
        crop_transform=custom_crop_transform,
        crop_prob=crop_prob,
        rotate_90_proba=rotate_90_proba,
        flip_transform=custom_flip_transform,
        rotate_transform=custom_rotate_transform,
        new_size=new_size,
        keep_ratio=keep_ratio,
        model_type=model_type,
        resize=resize
    )

    # Optionally concatenate a second dataset
    if annotations2:
        dataset2 = BPMN_Dataset(
            annotations=annotations2,
            transform=transformation,
            crop_transform=custom_crop_transform,
            crop_prob=crop_prob,
            rotate_90_proba=rotate_90_proba,
            flip_transform=custom_flip_transform,
            new_size=new_size,
            keep_ratio=keep_ratio,
            model_type=model_type,
            resize=resize
        )
        dataset = ConcatDataset([dataset, dataset2])  # Concatenate the two datasets

    # Set the seed for reproducibility in random operations within transformations and data loading
    random.seed(seed)
    torch.manual_seed(seed)

    # Create the DataLoader with the dataset
    data_loader = DataLoader(dataset, batch_size=batch_size, shuffle=True, collate_fn=collate_fn)

    return data_loader