Add predict command to cli dataset

Author: jrrodri

abraia/editing/__init__.py

@@ -5,7 +5,7 @@
 from .upscale import Upscaler, SwinIR
 from .smartcrop import Smartcrop
 from .inpaint import LAMA
-from .sam import SAM
+from ..inference.sam import SAM
 
 from ..inference import PlateDetector
 from ..inference.faces import FaceRecognizer

abraia/inference/detect.py

@@ -1,11 +1,13 @@
 import os
 import cv2
 import math
+import json
 import numpy as np
 import onnxruntime as ort
 
 from .ops import non_maximum_suppression, normalize, mask_to_polygon, softmax
 from ..utils import download_file, load_json, get_color, get_providers
+from .sam import SAM
 
 
 def resize(img, size):
@@ -116,3 +118,14 @@ def run(self, img, conf_threshold=0.35, iou_threshold=0.7, approx=0.001):
             return process_output(outputs, img_size, self.input_shape, self.config['classes'], conf_threshold, iou_threshold, approx)
         outputs = self.session.run(None, {self.input_name: preprocess(img)})
         return postprocess(outputs, self.config['classes'])
+
+
+def segment_objects(frame, results):
+    sam = SAM()
+    sam.encode(frame)
+    for result in results:
+        x, y, w, h = result['box']
+        mask = sam.predict(frame, prompt=json.dumps([{"type": "rectangle", "data": [x, y, x+w, y+h]}]))
+        result['polygon'] = mask_to_polygon(mask[y:y+h, x:x+w], (x, y))
+        # result['mask'] = mask[y:y+h, x:x+w]
+    return results

abraia/editing/sam.py abraia/inference/sam.pyabraia/editing/sam.py renamed to abraia/inference/sam.py

@@ -9,6 +9,7 @@
 from PIL import Image
 from tqdm.contrib.concurrent import process_map
 from sklearn.model_selection import train_test_split
+from typing import Dict, Any
 
 
 abraia = Abraia()
@@ -31,6 +32,7 @@ def load_annotations(project):
     annotations = abraia.load_json(f"{project}/annotations.json")
     for annotation in annotations:
         annotation['path'] = f"{project}/{annotation['filename']}"
+        annotation['url'] = url_path(f"{abraia.userid}/{annotation['path']}")
     return annotations
 
 
@@ -154,35 +156,41 @@ def split_dataset(annotations):
     return train, val, test
 
 
-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"{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
+class ModelTrainer:
+    """High-level trainer orchestrator using models and dataset utilities."""
+    def __init__(self, task: str, imgsz: int = 640):
+        self.task = task
+        if task == 'classify':
+            self.model = classify.Model()
+        else:
+            self.model = detect.Model(task, imgsz=imgsz)
+    
+    def prepare_dataset(self, project: str, classes, force: bool = False):
+        if force or 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"{project}/{x}", classes, self.task)
+            save_config(project, classes)
+
+    def train(self, dataset: str, epochs: int, batch: int = 32, **kwargs) -> None:
+        if self.task == 'classify':
+            self.model.train(dataset, epochs=epochs)
+        else:
+            self.model.train(dataset, epochs=epochs, batch=batch)
+
+    def test(self, split: str = 'val') -> Dict[str, Any]:
+        if self.task == 'classify':
+            # classify.Model currently has no .test
+            return {}
+        else:
+            return self.model.test(split=split)
+
+    def save(self, dataset: str, classes, device='cpu') -> None:
+        self.model.save(dataset, classes, device=device)
+
+    def run(self, img):
+        return self.model.run(img)
+    

abraia/training/__init__.py

@@ -72,7 +72,7 @@ def __getitem__(self, idx):
 
 
 def create_model(class_names, pretrained=True):
-    model = models.resnet18(pretrained=pretrained)
+    model = models.resnet18(weights=models.ResNet18_Weights.IMAGENET1K_V1 if pretrained else None)
     for param in model.parameters():
         param.requires_grad = False
     num_ftrs = model.fc.in_features
@@ -230,7 +230,10 @@ def train(self, dataset, epochs=25):
         model_conv = create_model(classes)
         self.classes = classes
         self.model = train_model(model_conv, dataloaders, num_epochs=epochs)
-        return self.model
+        #dataloaders, classes = training_session.create_dataset(project)
+        #training_session.train(project, epochs=epochs)
+        # training.classify.visualize_data(dataloaders['train'])
+        #training.classify.visualize_model(model, dataloaders['val'])
 
     def save(self, dataset, classes, device='cpu'):
         self.model.to(device)

abraia/training/classify.py

@@ -35,13 +35,16 @@ def __init__(self, task, model_type='yolov8n', imgsz=640):
         self.task = task
         self.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)
+
     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=self.imgsz)
-        # TODO: Merge with test and add parse metrics
-        metrics = self.model.val(data=data)
-        # self.metrics = self.test(split='val')
-        return metrics
 
     def test(self, split='val'):
         out = io.StringIO()

abraia/training/detect.py

@@ -4,7 +4,7 @@
 from abraia.inference import Model, Tracker
 from abraia.utils.draw import render_results, draw_overlay_mask, get_color, hex_to_rgb
 from abraia.inference.ops import count_objects
-from abraia.editing.sam import SAM
+from abraia.editing import SAM
 
 
 model_uri = 'multiple/models/yolov8n.onnx'