Add train_model to training

Author: jrrodri

README.md

@@ -236,17 +236,6 @@ pc_img = hsi.principal_components(img)
 hsi.plot_image(pc_img, 'Principal components')
 ```
 
-### Classification model
-
-Two classification models are directly available for automatic identification on hysperspectral images. One is based on support vector machines ('svm') while the other is based on deep image classification ('hsn'). Both models are available under a simple interface like bellow:
-
-```python
-n_bands, n_classes = 30, 17
-model = hsi.create_model('hsn', (25, 25, n_bands), n_classes)
-model.train(X, y, train_ratio=0.3, epochs=5)
-y_pred = model.predict(X)
-```
-
 ## License
 
 This software is licensed under the MIT License. [View the license](LICENSE).

abraia/__init__.py

@@ -2,7 +2,7 @@
 from dotenv import load_dotenv
 load_dotenv()
 
-__version__ = '0.25.3'
+__version__ = '0.25.4'
 
 from . import config
 from .client import Abraia, APIError

abraia/training/__init__.py

@@ -154,12 +154,35 @@ def split_dataset(annotations):
     return train, val, test
 
 
-def create_dataset(dataset, task, classes):
-    if not os.path.exists(dataset):
-        annotations = load_annotations(dataset)
+def create_dataset(project, task, classes):
+    if not os.path.exists(project):
+        annotations = load_annotations(project)
         train, val, test = split_dataset(annotations)
         data_annotations = {'train': train, 'val': val, 'test': test}
         #TODO: Download files in one single step
         for x in ['train', 'val', 'test']:
-            save_data(data_annotations[x], f"{dataset}/{x}", classes, task)
-        save_config(dataset, classes)
+            save_data(data_annotations[x], f"{project}/{x}", classes, task)
+        save_config(project, classes)
+
+
+def train_model(project, task, classes, epochs, batch=32, imgsz=640):
+    if task == 'classify':
+        training_session = classify.Model()
+        dataloaders, classes = training_session.create_dataset(project)
+        model = training_session.train(project, epochs=epochs)
+        # training.classify.visualize_data(dataloaders['train'])
+        #training.classify.visualize_model(model, dataloaders['val'])
+    else:
+        training_session = detect.Model(task, imgsz=imgsz)
+        def print_train_end(trainer):
+            print('# End training')
+            print('Metrics:', trainer.metrics)
+        #training_session.model.add_callback('on_train_start', print_train_start)
+        #training_session.model.add_callback('on_train_epoch_start', print_train_epoch)
+        training_session.model.add_callback('on_train_end', print_train_end)
+        metrics = training_session.train(project, epochs=epochs, batch=batch)
+        print("Train metrics")
+        print(training_session.test('val'))
+        #TODO: Save metrics with model
+        training_session.metrics = training_session.test('test')
+    return training_session

abraia/training/detect.py

@@ -27,16 +27,17 @@ def build_model_name(model_name, task):
 
 
 class Model:
-    def __init__(self, task, model_type='yolov8n'):
+    def __init__(self, task, model_type='yolov8n', imgsz=640):
         model_name = build_model_name(model_type, task)
         self.model = YOLO(f"{model_name}.pt", verbose=False)
         self.model_name = model_name
         self.metrics = {}
         self.task = task
+        self.imgsz = imgsz
 
-    def train(self, dataset, epochs=100, batch=32, imgsz=640):
+    def train(self, dataset, epochs=100, batch=32):
         data = f"{dataset}" if self.task == 'classify' else f"{dataset}/data.yaml"
-        results = self.model.train(data=data, batch=batch, epochs=epochs, imgsz=imgsz)
+        results = self.model.train(data=data, batch=batch, epochs=epochs, imgsz=self.imgsz)
         # TODO: Merge with test and add parse metrics
         metrics = self.model.val(data=data)
         # self.metrics = self.test(split='val')
@@ -46,15 +47,17 @@ def test(self, split='val'):
         out = io.StringIO()
         with contextlib.redirect_stderr(out):
             metrics = self.model.val(split=split)
-        return {'mAP': metrics.box.map50, 'P': metrics.box.p, 'R': metrics.box.r, 
-                'confusionMatrix': metrics.confusion_matrix.matrix}
+        return {'mAP': float(metrics.box.map50), 'P': metrics.box.p.tolist(), 'R': metrics.box.r.tolist(), 
+                'confusionMatrix': metrics.confusion_matrix.matrix.tolist()}
 
-    def save(self, dataset, classes, imgsz=640, device="cpu"):
+    def save(self, dataset, classes, device="cpu"):
         # TODO: Add versioning
-        model_src = self.model.export(format="onnx", device=device, opset=21, half=True)
+        out = io.StringIO()
+        with contextlib.redirect_stdout(out):
+            model_src = self.model.export(format="onnx", device=device, opset=19, half=True)
         abraia.upload_file(model_src, f"{dataset}/{self.model_name}.onnx")
         abraia.save_json(f"{dataset}/{self.model_name}.json", 
-                         {'task': self.task, 'inputShape': [1, 3, imgsz, imgsz], 
+                         {'task': self.task, 'inputShape': [1, 3, self.imgsz, self.imgsz], 
                           'classes': classes, 'metrics': self.metrics})
 
     def run(self, img):

scripts/abraia

@@ -286,6 +286,27 @@ def annotate(project, label):
     dataset.save_annotations(project, annotations)
 
 
+@cli_dataset.command()
+@click.argument('project')
+def train(project):
+    """Train a model on the specified dataset."""
+    from abraia import training
+    annotations = training.load_annotations(project)
+    classes = training.load_labels(annotations)
+    tasks = training.load_tasks(annotations)
+    task = tasks[-1] if len(tasks) else 'detect'
+    epochs = 25 if task == 'classify' else 250
+    imgsz = 224 if task == 'classify' else 640
+    click.echo('Loading dataset...')
+    if not os.path.exists(project):
+        training.create_dataset(project, task, classes)
+    click.echo('Dataset loaded.')
+    training_session = training.train_model(project, task, classes, epochs, imgsz=imgsz)
+    click.echo('Trained model.')
+    training_session.save(project, classes)
+    click.echo('Model saved.')
+
+
 if __name__ == '__main__':
     if not abraia.userid:
         configure()

setup.py

@@ -14,7 +14,7 @@
 
 extras_require = {
     'multiple': ['spectral>=0.23.1', 'scipy>=1.14.1', 'tifffile>=2024.8.30'],
-    'dev': ['opencv-python>=4.10.0.84', 'ImageHash>=4.3.2', 'tifffile>=2024.8.30', 'ultralytics>=8.3.59', 'onnx>=1.18.0', 'transformers>=4.57.1'],#, 'onnxsim>=0.4.36'],
+    'dev': ['opencv-python>=4.10.0.84', 'ImageHash>=4.3.2', 'tifffile>=2024.8.30', 'ultralytics>=8.3.161', 'onnx>=1.18.0', 'transformers>=4.57.1'],#, 'onnxsim>=0.4.36'],
 }
 
 setup(