import torch import torchvision.transforms.functional as F import numpy as np import cv2 import matplotlib.pyplot as plt import streamlit as st # Define dictionaries to map class indices to their corresponding names object_dict = { 0: 'background', 1: 'task', 2: 'exclusiveGateway', 3: 'event', 4: 'parallelGateway', 5: 'messageEvent', 6: 'pool', 7: 'lane', 8: 'dataObject', 9: 'dataStore', 10: 'subProcess', 11: 'eventBasedGateway', 12: 'timerEvent', } arrow_dict = { 0: 'background', 1: 'sequenceFlow', 2: 'dataAssociation', 3: 'messageFlow', } class_dict = { 0: 'background', 1: 'task', 2: 'exclusiveGateway', 3: 'event', 4: 'parallelGateway', 5: 'messageEvent', 6: 'pool', 7: 'lane', 8: 'dataObject', 9: 'dataStore', 10: 'subProcess', 11: 'eventBasedGateway', 12: 'timerEvent', 13: 'sequenceFlow', 14: 'dataAssociation', 15: 'messageFlow', } def is_inside(box1, box2): """Check if the center of box1 is inside box2.""" x_center = (box1[0] + box1[2]) / 2 y_center = (box1[1] + box1[3]) / 2 return box2[0] <= x_center <= box2[2] and box2[1] <= y_center <= box2[3] def is_vertical(box): """Determine if the text in the bounding box is vertically aligned.""" width = box[2] - box[0] height = box[3] - box[1] return (height > 2 * width) def rescale_boxes(scale, boxes): """Rescale the bounding boxes by a given scale factor.""" for i in range(len(boxes)): boxes[i] = [boxes[i][0] * scale, boxes[i][1] * scale, boxes[i][2] * scale, boxes[i][3] * scale] return boxes def iou(box1, box2): """Calculate the Intersection over Union (IoU) of two bounding boxes.""" inter_box = [max(box1[0], box2[0]), max(box1[1], box2[1]), min(box1[2], box2[2]), min(box1[3], box2[3])] inter_area = max(0, inter_box[2] - inter_box[0]) * max(0, inter_box[3] - inter_box[1]) box1_area = (box1[2] - box1[0]) * (box1[3] - box1[1]) box2_area = (box2[2] - box2[0]) * (box2[3] - box2[1]) union_area = box1_area + box2_area - inter_area return inter_area / union_area def proportion_inside(box1, box2): """Calculate the proportion of the smaller box inside the larger box.""" box1_area = (box1[2] - box1[0]) * (box1[3] - box1[1]) box2_area = (box2[2] - box2[0]) * (box2[3] - box2[1]) big_box, small_box = (box1, box2) if box1_area > box2_area else (box2, box1) inter_box = [max(small_box[0], big_box[0]), max(small_box[1], big_box[1]), min(small_box[2], big_box[2]), min(small_box[3], big_box[3])] inter_area = max(0, inter_box[2] - inter_box[0]) * max(0, inter_box[3] - inter_box[1]) proportion = inter_area / ((small_box[2] - small_box[0]) * (small_box[3] - small_box[1])) return min(proportion, 1.0) def resize_boxes(boxes, original_size, target_size): """ Resizes bounding boxes according to a new image size. Parameters: - boxes (np.array): The original bounding boxes as a numpy array of shape [N, 4]. - original_size (tuple): The original size of the image as (width, height). - target_size (tuple): The desired size to resize the image to as (width, height). Returns: - np.array: The resized bounding boxes as a numpy array of shape [N, 4]. """ orig_width, orig_height = original_size target_width, target_height = target_size width_ratio = target_width / orig_width height_ratio = target_height / orig_height boxes[:, 0] *= width_ratio boxes[:, 1] *= height_ratio boxes[:, 2] *= width_ratio boxes[:, 3] *= height_ratio return boxes def resize_keypoints(keypoints, original_size, target_size): """ Resize keypoints based on the original and target dimensions of an image. Parameters: - keypoints (np.ndarray): The array of keypoints, where each keypoint is represented by its (x, y) coordinates. - original_size (tuple): The width and height of the original image (width, height). - target_size (tuple): The width and height of the target image (width, height). Returns: - np.ndarray: The resized keypoints. """ orig_width, orig_height = original_size target_width, target_height = target_size width_ratio = target_width / orig_width height_ratio = target_height / orig_height keypoints[:, 0] *= width_ratio keypoints[:, 1] *= height_ratio return keypoints def write_results(name_model, metrics_list, start_epoch): """Write training results to a text file.""" with open('./results/' + name_model + '.txt', 'w') as f: for i in range(len(metrics_list[0])): f.write(f"{i + 1 + start_epoch},{metrics_list[0][i]},{metrics_list[1][i]},{metrics_list[2][i]},{metrics_list[3][i]},{metrics_list[4][i]},{metrics_list[5][i]},{metrics_list[6][i]},{metrics_list[7][i]},{metrics_list[8][i]},{metrics_list[9][i]} \n") def find_other_keypoint(idx, keypoints, boxes): """ Find the opposite keypoint to the center of the box. Parameters: - idx (int): The index of the box and keypoints. - keypoints (np.ndarray): The array of keypoints. - boxes (np.ndarray): The array of bounding boxes. Returns: - tuple: The coordinates of the new keypoint and the average keypoint. """ box = boxes[idx] key1, key2 = keypoints[idx] x1, y1, x2, y2 = box center = ((x1 + x2) // 2, (y1 + y2) // 2) average_keypoint = (key1 + key2) // 2 if average_keypoint[0] < center[0]: x = center[0] + abs(center[0] - average_keypoint[0]) else: x = center[0] - abs(center[0] - average_keypoint[0]) if average_keypoint[1] < center[1]: y = center[1] + abs(center[1] - average_keypoint[1]) else: y = center[1] - abs(center[1] - average_keypoint[1]) return x, y, average_keypoint[0], average_keypoint[1] def filter_overlap_boxes(boxes, scores, labels, keypoints, iou_threshold=0.5): """ Filters overlapping boxes based on the Intersection over Union (IoU) metric, keeping only the boxes with the highest scores. Parameters: - boxes (np.ndarray): Array of bounding boxes with shape (N, 4), where each row contains [x_min, y_min, x_max, y_max]. - scores (np.ndarray): Array of scores for each box, reflecting the confidence of detection. - labels (np.ndarray): Array of labels corresponding to each box. - keypoints (np.ndarray): Array of keypoints associated with each box. - iou_threshold (float): Threshold for IoU above which a box is considered overlapping. Returns: - tuple: Filtered boxes, scores, labels, and keypoints. """ areas = (boxes[:, 2] - boxes[:, 0]) * (boxes[:, 3] - boxes[:, 1]) order = scores.argsort()[::-1] keep = [] while order.size > 0: i = order[0] keep.append(i) xx1 = np.maximum(boxes[i, 0], boxes[order[1:], 0]) yy1 = np.maximum(boxes[i, 1], boxes[order[1:], 1]) xx2 = np.minimum(boxes[i, 2], boxes[order[1:], 2]) yy2 = np.minimum(boxes[i, 3], boxes[order[1:], 3]) w = np.maximum(0.0, xx2 - xx1) h = np.maximum(0.0, yy2 - yy1) inter = w * h iou = inter / (areas[i] + areas[order[1:]] - inter) inds = np.where(iou <= iou_threshold)[0] order = order[inds + 1] boxes = boxes[keep] scores = scores[keep] labels = labels[keep] keypoints = keypoints[keep] return boxes, scores, labels, keypoints def draw_annotations(image, target=None, prediction=None, full_prediction=None, text_predictions=None, model_dict=class_dict, draw_keypoints=False, draw_boxes=False, draw_text=False, draw_links=False, draw_twins=False, write_class=False, write_score=False, write_text=False, write_idx=False, score_threshold=0.4, keypoints_correction=False, only_print=None, axis=False, return_image=False, new_size=(1333, 800), resize=False): """ Draws annotations on images including bounding boxes, keypoints, links, and text. Parameters: - image (np.array): The image on which annotations will be drawn. - target (dict): Ground truth data containing boxes, labels, etc. - prediction (dict): Prediction data from a model. - full_prediction (dict): Additional detailed prediction data, potentially including relationships. - text_predictions (tuple): OCR text predictions containing bounding boxes and texts. - model_dict (dict): Mapping from class IDs to class names. - draw_keypoints (bool): Flag to draw keypoints. - draw_boxes (bool): Flag to draw bounding boxes. - draw_text (bool): Flag to draw text annotations. - draw_links (bool): Flag to draw links between annotations. - draw_twins (bool): Flag to draw twin keypoints. - write_class (bool): Flag to write class names near the annotations. - write_score (bool): Flag to write scores near the annotations. - write_text (bool): Flag to write OCR recognized text. - score_threshold (float): Threshold for scores above which annotations will be drawn. - only_print (str): Specific class name to filter annotations by. - resize (bool): Whether to resize annotations to fit the image size. """ # Convert image to RGB (if not already in that format) if prediction is None: image = image.squeeze(0).permute(1, 2, 0).cpu().numpy() image = cv2.cvtColor(image, cv2.COLOR_BGR2RGB) image_copy = image.copy() scale = max(image.shape[0], image.shape[1]) / 1000 # Helper function to draw annotations based on provided data def draw(data, is_prediction=False): for i in range(len(data['boxes'])): box = data['boxes'][i].tolist() x1, y1, x2, y2 = box if resize: x1, y1, x2, y2 = resize_boxes(np.array([box]), new_size, (image_copy.shape[1], image_copy.shape[0]))[0] if is_prediction: score = data['scores'][i].item() if score < score_threshold: continue if draw_boxes: if only_print is not None: if data['labels'][i] != list(model_dict.values()).index(only_print): continue cv2.rectangle(image_copy, (int(x1), int(y1)), (int(x2), int(y2)), (0, 0, 0) if is_prediction else (0, 0, 0), int(2 * scale)) if is_prediction and write_score: cv2.putText(image_copy, str(round(score, 2)), (int(x1), int(y1) + int(15 * scale)), cv2.FONT_HERSHEY_SIMPLEX, scale / 2, (100, 100, 255), 2) if write_class and 'labels' in data: class_id = data['labels'][i].item() cv2.putText(image_copy, model_dict[class_id], (int(x1), int(y1) - int(2 * scale)), cv2.FONT_HERSHEY_SIMPLEX, scale / 2, (255, 100, 100), 2) if write_idx: cv2.putText(image_copy, str(i), (int(x1) + int(15 * scale), int(y1) + int(15 * scale)), cv2.FONT_HERSHEY_SIMPLEX, 2 * scale, (0, 0, 0), 2) # Draw keypoints if available if draw_keypoints and 'keypoints' in data: if is_prediction and keypoints_correction: for idx, (key1, key2) in enumerate(data['keypoints']): if data['labels'][idx] not in [list(model_dict.values()).index('sequenceFlow'), list(model_dict.values()).index('messageFlow'), list(model_dict.values()).index('dataAssociation')]: continue distance = np.linalg.norm(key1[:2] - key2[:2]) if distance < 5: x_new, y_new, x, y = find_other_keypoint(idx, data['keypoints'], data['boxes']) data['keypoints'][idx][0] = torch.tensor([x_new, y_new, 1]) data['keypoints'][idx][1] = torch.tensor([x, y, 1]) print("keypoint has been changed") for i in range(len(data['keypoints'])): kp = data['keypoints'][i] for j in range(kp.shape[0]): if is_prediction and data['labels'][i] not in [list(model_dict.values()).index('sequenceFlow'), list(model_dict.values()).index('messageFlow'), list(model_dict.values()).index('dataAssociation')]: continue if is_prediction: score = data['scores'][i] if score < score_threshold: continue x, y, v = np.array(kp[j]) if resize: x, y, v = resize_keypoints(np.array([kp[j]]), new_size, (image_copy.shape[1], image_copy.shape[0]))[0] if j == 0: cv2.circle(image_copy, (int(x), int(y)), int(5 * scale), (0, 0, 255), -1) else: cv2.circle(image_copy, (int(x), int(y)), int(5 * scale), (255, 0, 0), -1) # Draw text predictions if available if (draw_text or write_text) and text_predictions is not None: for i in range(len(text_predictions[0])): x1, y1, x2, y2 = text_predictions[0][i] text = text_predictions[1][i] if resize: x1, y1, x2, y2 = resize_boxes(np.array([[float(x1), float(y1), float(x2), float(y2)]]), new_size, (image_copy.shape[1], image_copy.shape[0]))[0] if draw_text: cv2.rectangle(image_copy, (int(x1), int(y1)), (int(x2), int(y2)), (0, 255, 0), int(2 * scale)) if write_text: cv2.putText(image_copy, text, (int(x1 + int(2 * scale)), int((y1 + y2) / 2)), cv2.FONT_HERSHEY_SIMPLEX, scale / 2, (0, 0, 0), 2) def draw_with_links(full_prediction): """Draws links between objects based on the full prediction data.""" if draw_twins and full_prediction is not None: circle_color = (0, 255, 0) circle_radius = int(10 * scale) for idx, (key1, key2) in enumerate(full_prediction['keypoints']): if full_prediction['labels'][idx] not in [list(model_dict.values()).index('sequenceFlow'), list(model_dict.values()).index('messageFlow'), list(model_dict.values()).index('dataAssociation')]: continue distance = np.linalg.norm(key1[:2] - key2[:2]) if distance < 10: x_new, y_new, x, y = find_other_keypoint(idx, full_prediction['keypoints'], full_prediction['boxes']) cv2.circle(image_copy, (int(x), int(y)), circle_radius, circle_color, -1) cv2.circle(image_copy, (int(x_new), int(y_new)), circle_radius, (0, 0, 0), -1) if draw_links and full_prediction is not None: for i, (start_idx, end_idx) in enumerate(full_prediction['links']): if start_idx is None or end_idx is None: continue start_box = full_prediction['boxes'][start_idx] end_box = full_prediction['boxes'][end_idx] current_box = full_prediction['boxes'][i] start_center = ((start_box[0] + start_box[2]) // 2, (start_box[1] + start_box[3]) // 2) end_center = ((end_box[0] + end_box[2]) // 2, (end_box[1] + end_box[3]) // 2) current_center = ((current_box[0] + current_box[2]) // 2, (current_box[1] + current_box[3]) // 2) cv2.line(image_copy, (int(start_center[0]), int(start_center[1])), (int(current_center[0]), int(current_center[1])), (0, 0, 255), int(2 * scale)) cv2.line(image_copy, (int(current_center[0]), int(current_center[1])), (int(end_center[0]), int(end_center[1])), (255, 0, 0), int(2 * scale)) i += 1 if target is not None: draw(target, is_prediction=False) if prediction is not None: draw(prediction, is_prediction=True) if full_prediction is not None: draw_with_links(full_prediction) image_copy = cv2.cvtColor(image_copy, cv2.COLOR_BGR2RGB) plt.figure(figsize=(12, 12)) plt.imshow(image_copy) if not axis: plt.axis('off') plt.show() if return_image: return image_copy def find_closest_object(keypoint, boxes, labels): """ Find the closest object to a keypoint based on their proximity. Parameters: - keypoint (numpy.ndarray): The coordinates of the keypoint. - boxes (numpy.ndarray): The bounding boxes of the objects. Returns: - int or None: The index of the closest object to the keypoint, or None if no object is found. """ closest_object_idx = None best_point = None min_distance = float('inf') for i, box in enumerate(boxes): if labels[i] in [list(class_dict.values()).index('sequenceFlow'), list(class_dict.values()).index('messageFlow'), list(class_dict.values()).index('dataAssociation'), list(class_dict.values()).index('lane')]: continue x1, y1, x2, y2 = box top = ((x1 + x2) / 2, y1) bottom = ((x1 + x2) / 2, y2) left = (x1, (y1 + y2) / 2) right = (x2, (y1 + y2) / 2) points = [left, top, right, bottom] pos_dict = {0: 'left', 1: 'top', 2: 'right', 3: 'bottom'} for pos, point in enumerate(points): distance = np.linalg.norm(keypoint[:2] - point) if distance < min_distance: min_distance = distance closest_object_idx = i best_point = pos_dict[pos] return closest_object_idx, best_point def error(text='There is an error in the detection'): """Display an error message using Streamlit.""" st.error(text, icon="🚨") def warning(text='Some element are maybe not detected, verify the results, try to modify the parameters or try to add it in the method and style step.'): """Display a warning message using Streamlit.""" st.warning(text, icon="⚠️")