Develop

cucaracha/__init__.py

@@ -7,9 +7,9 @@
 from rich import print
 from rich.progress import track
 
-from cucaracha.aligment import inplane_deskew
-from cucaracha.noise_removal import sparse_dots
-from cucaracha.threshold import otsu
+from cucaracha.tasks.aligment import inplane_deskew
+from cucaracha.tasks.noise_removal import sparse_dots
+from cucaracha.tasks.threshold import otsu
 
 
 class Document:

cucaracha/configuration.py

@@ -0,0 +1 @@
+ALLOWED_IMAGE_EXTENSIONS = ('.png', '.jpg', '.jpeg')

cucaracha/ml_models/__init__.py

@@ -1,22 +1,27 @@
+import os
+
 VALID_MODALITIES = [
     'image_classification',
     'image_keypoint_detection',
     'image_object_detection',
     'image_segmentation',
 ]
+DEFAULT_MODEL_LAKE = os.path.join(
+    os.path.expanduser('~'), '.cache', 'kagglehub', 'models'
+)
 
 # Pre-treined cucahacha models
 CUCARACHA_PRESETS = {
     'image_classification': {
         'doc_is_signed': {
             'variation': 'cucaracha-project/cucaracha-imgclass-document-is-signed/tensorFlow2/cucaracha-imgclass-document_is_signed-v0.1.0',
             'dataset': 'cucaracha-project/cucaracha-mod-imgclass-constains-signature',
-        }
-    },
-    'image_segmentation': {
-        'TBD': {
-            'variation': 'cucaracha-project/TBD',
-            'dataset': 'cucaracha-project/TBD',
-        }
-    },
+            'labels': {0: 'unsigned', 1: 'signed'},
+        },
+        'cnh_cpf_rg': {
+            'variation': 'cucaracha-project/cucaracha-imgclass-brazilian-personal-document/keras/cucaracha-imgclass-brazilian-personal-document',
+            'dataset': 'cucaracha-project/cucaracha-mod-imgclass-brazilian-personal-doc',
+            'labels': {0: 'rg', 1: 'others', 2: 'cpf', 3: 'cnh'},
+        },
+    }
 }

cucaracha/ml_models/image_classification/alex_net.py

@@ -77,5 +77,6 @@ def get_model(self):
         return keras.Model(inputs, outputs)
 
     def __str__(self):
-        super().__str__()
+        output = super().__str__()
         self.get_model().summary(show_trainable=True)
+        return output

cucaracha/ml_models/image_classification/dense_net_121.py

@@ -40,5 +40,6 @@ def get_model(self):
         )
 
     def __str__(self):
-        super().__str__()
+        output = super().__str__()
         self.get_model().summary(show_trainable=True)
+        return output

cucaracha/ml_models/image_classification/model_soup.py

@@ -47,5 +47,6 @@ def get_model(self):
         return keras.Model(model.input, prediction)
 
     def __str__(self):
-        super().__str__()
+        output = super().__str__()
         self.get_model().summary(show_trainable=True)
+        return output

cucaracha/ml_models/image_classification/res_net_50.py

@@ -40,5 +40,6 @@ def get_model(self):
         )
 
     def __str__(self):
-        super().__str__()
+        output = super().__str__()
         self.get_model().summary(show_trainable=True)
+        return output

cucaracha/ml_models/image_classification/small_xception.py

@@ -70,5 +70,6 @@ def get_model(self):
         return keras.Model(inputs, outputs)
 
     def __str__(self):
-        super().__str__()
+        output = super().__str__()
         self.get_model().summary(show_trainable=True)
+        return output

cucaracha/ml_models/image_segmentation/unet_xception.py

@@ -91,5 +91,6 @@ def get_model(self):
         return keras.Model(inputs, outputs)
 
     def __str__(self):
-        super().__str__()
+        output = super().__str__()
         self.get_model().summary(show_trainable=True)
+        return output

cucaracha/ml_models/model_download.py → cucaracha/ml_models/kaggle_helpers.py

@@ -1,6 +1,8 @@
+import os
+
 import kagglehub
 
-from cucaracha.ml_models import CUCARACHA_PRESETS
+from cucaracha.ml_models import CUCARACHA_PRESETS, DEFAULT_MODEL_LAKE
 
 
 def download_cucaracha_model(model_url: str):
@@ -74,3 +76,52 @@ def download_cucaracha_dataset(dataset_url: str):
         raise ValueError(f'Error downloading the dataset: {e}')
 
     return path
+
+
+def collect_cucaracha_model(cucaracha_preset: str):
+    """
+    Collects the Cucaracha model from the given preset.
+
+    This function checks if the provided model preset is present in the CUCARACHA_PRESETS dictionary.
+    If the preset is valid, it attempts to download the model using the kagglehub library.
+
+    The downloaded files are located in the home/.cache folder.
+
+    Note:
+        We used the kagglehub library to make all the operations here. The
+        CUCARACHA_PRESETS dictionary is expected to have a nested structure where
+        the model variations are stored under a 'variation' key. If the URL is
+        valid, it attempts to download the model using the kagglehub library.
+
+    Args:
+        cucaracha_preset (str): The name of the model preset to be downloaded.
+    Returns:
+        str: The path where the model is downloaded.
+    Raises:
+        ValueError: If the model preset is not present in CUCARACHA_PRESETS or if there is an error during download.
+    """
+    found = False
+    modality = None
+    for mod in CUCARACHA_PRESETS.values():
+        if cucaracha_preset in mod:
+            found = True
+            modality = next(
+                key
+                for key, value in CUCARACHA_PRESETS.items()
+                if cucaracha_preset in value
+            )
+            break
+
+    if not found:
+        raise ValueError(
+            f'Model preset {cucaracha_preset} is not present in CUCARACHA_PRESETS'
+        )
+
+    model_url = CUCARACHA_PRESETS[modality][cucaracha_preset]['variation']
+
+    output = {
+        'model_path': download_cucaracha_model(model_url),
+        'modality': modality,
+        'labels': CUCARACHA_PRESETS[modality][cucaracha_preset]['labels'],
+    }
+    return output

cucaracha/ml_models/model_architect.py

@@ -4,17 +4,27 @@
 
 
 class ModelArchitect(ABC):
+    """
+    Abstract base class for defining model architectures.
+    Attributes:
+        modality (str): The modality of the model architecture. Expected values are defined in VALID_MODALITIES.
+    Methods:
+        get_model():
+            Abstract method to be implemented by subclasses to return the model architecture.
+        __str__():
+            Returns a string representation of the model architecture, including its modality.
+    """
+
     def __init__(self, **kwargs):
         self.modality = kwargs.get('modality', None)
-        # valid_modalities = ['image_classification', 'image_keypoint_detection', 'image_object_detection']
         if self.modality is None or self.modality not in VALID_MODALITIES:
             raise ValueError(
                 f'Invalid modality. Expected one of {VALID_MODALITIES}, got {self.modality}'
             )
 
     @abstractmethod
-    def get_model(self):
+    def get_model(self):   # pragma: no cover
         pass
 
-    def __str__(self):
+    def __str__(self):   # pragma: no cover
         return f'Model Architecture modality: {self.modality}'