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