change of the model object

modules/streamlit_utils.py

@@ -53,7 +53,7 @@ def load_models():
         model_arrow = get_arrow_model(len(arrow_dict),2)
 
         url_arrow = 'https://drive.google.com/uc?id=1vv1X_r_lZ8gnzMAIKxcVEb_T_Qb-NkyA'
-        url_object = 'https://drive.google.com/uc?id=1b1bqogxqdPS-SnvaOfWJGV1I1qOrTKh5'
+        url_object = 'https://drive.google.com/uc?id=1ciSS7H5baqXf8sRilcAjWBM9jZ0DQhsP'
 
         # Define paths to save models
         output_arrow = 'model_arrow.pth'

modules/utils.py

@@ -411,7 +411,7 @@ def rotate_vertical(image, target):
     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,800),keep_ratio=False,resize=True, model_type='object', rotate_vertical=False):
+    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,800),keep_ratio=0.1,resize=True, model_type='object'):
         self.annotations = annotations
         print(f"Loaded {len(self.annotations)} annotations.")
         self.transform = transform
@@ -420,7 +420,6 @@ class BPMN_Dataset(Dataset):
         self.flip_transform = flip_transform
         self.rotate_transform = rotate_transform
         self.resize = resize
-        self.rotate_vertical = rotate_vertical
         self.new_size = new_size
         self.keep_ratio = keep_ratio
         self.model_type = model_type
@@ -494,11 +493,11 @@ class BPMN_Dataset(Dataset):
             image, target = self.crop_transform(image, target)
             
         # Rotate vertical image
-        if self.rotate_vertical and random.random() < self.rotate_90_proba:
+        if random.random() < self.rotate_90_proba:
             image, target = rotate_vertical(image, target)
 
         if self.resize:
-            if self.keep_ratio:
+            if random.random() < self.keep_ratio:
                 original_size = image.size
                 # Calculate scale to fit the new size while maintaining aspect ratio
                 scale = min(self.new_size[0] / original_size[0], self.new_size[1] / original_size[1])
@@ -557,7 +556,7 @@ def collate_fn(batch):
 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, rotate_vertical=False, keep_ratio=False, model_type = 'object'):
+                  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.
 
@@ -592,7 +591,6 @@ def create_loader(new_size,transformation, annotations1, annotations2=None,
         rotate_90_proba=rotate_90_proba,
         flip_transform=custom_flip_transform,
         rotate_transform=custom_rotate_transform,
-        rotate_vertical=rotate_vertical,
         new_size=new_size,
         keep_ratio=keep_ratio,
         model_type=model_type,
@@ -608,7 +606,6 @@ def create_loader(new_size,transformation, annotations1, annotations2=None,
             crop_prob=crop_prob,
             rotate_90_proba=rotate_90_proba,
             flip_transform=custom_flip_transform,
-            rotate_vertical=rotate_vertical,
             new_size=new_size,
             keep_ratio=keep_ratio,
             model_type=model_type,