correct bug with Vizifile

modules/toWizard.py

@@ -1,6 +1,7 @@
 import xml.etree.ElementTree as ET
 from modules.utils import class_dict
 from xml.dom import minidom
+from modules.utils import error
 
 def rescale(scale, boxes):
     for i in range(len(boxes)):
@@ -52,6 +53,9 @@ def check_end(val):
 
 def connect(data, text_mapping, i):
     target_idx = data['links'][i][1]  
+    if target_idx >= len(data['links']):
+        error('There is an error with the Vizi file, care when you download it.')
+        return None, None
     current_id = data['BPMN_id'][i]
     next_idx = data['links'][target_idx][1]
     next_id = data['BPMN_id'][next_idx]
@@ -116,11 +120,11 @@ def create_wizard_file(data, text_mapping):
     
     for idx, activity_name in enumerate(data['BPMN_id']):
         if activity_name.startswith('task'):
-            target_text = ' FINISH'
             activity = ET.SubElement(activities, 'activity', attrib={'name': text_mapping.get(activity_name, activity_name), 'performer': ''})
             endStates = ET.SubElement(activity, 'endStates')
             current_text, next_text = connect(data, text_mapping, idx)
-            ET.SubElement(endStates, 'endState', attrib={'name': next_text, 'isRegular': 'True'})
+            if next_text is not None:
+                ET.SubElement(endStates, 'endState', attrib={'name': next_text, 'isRegular': 'True'})
             ET.SubElement(activity, 'subActivities')
             ET.SubElement(activity, 'subActivityFlows')
             ET.SubElement(activity, 'messageFlows')
@@ -130,10 +134,14 @@ def create_wizard_file(data, text_mapping):
     for i, link in enumerate(data['links']):
         if link[0] is None and link[1] is not None and (data['BPMN_id'][i].split('_')[0] == 'event' or data['BPMN_id'][i].split('_')[0] == 'message'):
             current_text, next_text = connect(data, text_mapping, i)
+            if current_text is None or next_text is None:
+                continue
             ET.SubElement(activityFlows, 'activityFlow', attrib={'startEvent': current_text, 'endState': '---', 'target': next_text, 'isMerging': 'False', 'isPredefined': 'True'})
             i+=1
         if link[0] is not None and link[1] is not None and data['BPMN_id'][i].split('_')[0] == 'task':
             current_text, next_text = connect(data, text_mapping, i)
+            if current_text is None or next_text is None:
+                continue
             ET.SubElement(activityFlows, 'activityFlow', attrib={'activity': current_text, 'endState': '---', 'target': next_text, 'isMerging': 'False', 'isPredefined': 'True'})
             i+=1
     

modules/train.py

@@ -7,9 +7,10 @@ import matplotlib.pyplot as plt
 
 from modules.eval import main_evaluation
 from torch.optim import SGD, AdamW
-from torchvision.models.detection import keypointrcnn_resnet50_fpn, KeypointRCNN_ResNet50_FPN_Weights
+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 fasterrcnn_resnet50_fpn
 from tqdm import tqdm
 from modules.utils import write_results
 
@@ -36,7 +37,6 @@ def get_arrow_model(num_classes, num_keypoints=2):
        This is necessary to tailor the model to specific tasks that may have different keypoint structures.
     """
     # Load a model pre-trained on COCO, initialized without pre-trained weights
-    device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu')
     model = keypointrcnn_resnet50_fpn(weights=None)
 
     # Get the number of input features for the classifier in the box predictor.
@@ -50,8 +50,7 @@ def get_arrow_model(num_classes, num_keypoints=2):
 
     return model
 
-from torchvision.models.detection import fasterrcnn_resnet50_fpn
-from torchvision.models.detection import FasterRCNN_ResNet50_FPN_Weights
+
 def get_faster_rcnn_model(num_classes):
     """
     Configures and returns a modified Faster R-CNN model based on ResNet-50 with FPN, adapted for a custom number of classes.