Refactor dataset preparation

Author: jrrodri

abraia/__init__.py

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

abraia/client.py

@@ -89,20 +89,15 @@ def move_file(self, old_path, new_path):
         resp = resp.json()
         return file_path(resp['file']['source'], self.userid)
 
-    def download_file(self, path, dest='', cache=False):
-        # TODO: Remove BytesIO download option
+    def download_file(self, path, dest, cache=False):
         url = f"{API_URL}/files/{self.userid}/{path}"
-        if cache and dest == '':
-            dest = temporal_src(path)
-            if os.path.exists(dest):
-                return dest
+        if cache and os.path.exists(dest):
+            return dest
         resp = requests.get(url, stream=True, auth=self.auth)
         if resp.status_code != 200:
             raise APIError(resp.text, resp.status_code)
-        if dest:
-            save_data(dest, resp.content)
-            return dest
-        return BytesIO(resp.content)
+        save_data(dest, resp.content)
+        return dest
 
     def remove_file(self, path):
         url = f"{API_URL}/files/{self.userid}/{path}"
@@ -141,11 +136,14 @@ def transform_image(self, path, dest, params={'quality': 'auto'}):
         save_data(dest, resp.content)
 
     def load_file(self, path):
-        stream = self.download_file(path)
+        dest = temporal_src(path)
+        self.download_file(path, dest, cache=True)
         try:
-            return stream.getvalue().decode('utf-8')
+            with open(dest, 'r') as f:
+                return f.read()
         except:
-            return stream
+            with open(dest, 'rb') as f:
+                return BytesIO(f.read())
 
     def save_file(self, path, stream):
         stream =  BytesIO(bytes(stream, 'utf-8')) if isinstance(stream, str) else stream
@@ -158,7 +156,8 @@ def save_json(self, path, values):
         return self.save_file(path, json.dumps(values))
 
     def load_image(self, path):
-        return Image.open(self.download_file(path))
+        dest = temporal_src(path)
+        return Image.open(self.download_file(path, dest, cache=True))
 
     def save_image(self, path, im):
         src = temporal_src(path)

abraia/demo.py

@@ -16,39 +16,39 @@
 DEMOS = {
     'apple': {
         'src': '5479199-hd_1920_1080_25fps.mp4',
-        'label': 'apple'
+        'labels': ['apple']
     },
     'strawberry': {
         'model': 'multiple/strawberry/yolov8n.onnx',
         'src': '9710983-hd_1920_1080_30fps.mp4',
-        'label': 'strawberry'
+        'labels': ['strawberry']
     },
     'grapes': {
         'model': 'multiple/grapes/yolov8n.onnx',
         'src': '5658544-hd_1366_720_24fps.mp4',
-        'label': 'grapes'
+        'labels': ['grapes']
     },
     'tomato': {
         'model': 'multiple/tomato/yolov8n_v6.onnx',
         'src': '10179855-hd_1920_1080_30fps.mp4',
-        'label': 'tomato',
-        'line': [(1440, 0), (1440, 1080)]
+        'labels': ['tomato'],
+        'counter': [(1440, 0), (1440, 1080)]
     },
     'people': {
         'src': '853889-hd_1920_1080_25fps.mp4',
-        'label': 'person',
-        'line': [(0, 650), (1920, 650)],
+        'labels': ['person'],
+        'counter': [(0, 650), (1920, 650)],
         'region': [(0, 600), (1920, 600), (1920, 700), (0, 700)]
     },
     'queue': {
         'src': '4775505-hd_1920_1080_30fps.mp4',
-        'label': 'person',
+        'labels': ['person'],
         'timer': [(10, 600), (1690, 600), (1690, 700), (10, 700)]
     },
     'escalator': {
         'src': '14393755-hd_1920_1080_30fps.mp4',
-        'label': 'person',
-        'line': [(950, 670), (270, 895)],
+        'labels': ['person'],
+        'counter': [(950, 670), (270, 895)],
         'region': [[0, 245], [350, 1080], [1200, 1080], [530, 0], [0, 0]],
         'timer': [[0, 245], [350, 1080], [1200, 1080], [530, 0], [0, 0]]
     }
@@ -67,32 +67,27 @@ def monitor_objects(src=None, demo='detect', resolution=(1280, 720)):
     tracker = Tracker(frame_rate=video.frame_rate)
     model = Model(selected.get('model', 'multiple/models/yolov8n.onnx'))
 
-    line_counter = LineCounter(selected['line']) if selected.get('line') else None
+    line_counter = LineCounter(selected['counter']) if selected.get('counter') else None
     region_filter = RegionFilter(selected['region']) if selected.get('region') else None
     region_timer = RegionTimer(selected['timer']) if selected.get('timer') else None
 
-    labels = [selected.get('label')] if selected.get('label') else None
+    labels = selected.get('labels')
     for k, frame in enumerate(video):
         frame_time = round(k / video.frame_rate, 2)
         t0 = time.time()
         results = model.run(frame, labels=labels)
-        
         if region_filter:
             results, _ = region_filter.update(results)
         results = tracker.update(results)
-        
         out = frame.copy()
         if line_counter:
             in_count, out_count = line_counter.update(results)
-            out = render_counter(out, line_counter.line, f"In: {in_count} | Out: {out_count}" if demo in ['people', 'escalator'] else f"Count: {out_count}")
-        
+            out = render_counter(out, line_counter.line, f"In: {in_count} | Out: {out_count}")
         if region_timer:
             in_objects, out_objects = region_timer.update(results, frame_time)
             out = render_region(out, region_timer.region, f"Count: {len(in_objects)}", color=(255, 255, 0))
-            out = render_results(out, in_objects)
-        else:
-            out = render_results(out, results)
-        
+            results = in_objects
+        out = render_results(out, results)
         print(f"#{k} {round((time.time() - t0) * 1000, 1)}ms {count_objects(results)}")
         video.show(out)
 

abraia/multiple/__init__.py

@@ -139,7 +139,8 @@ def __init__(self, folder=''):
         super(Multiple, self).__init__()
 
     def load_header(self, path):
-        dest = self.download_file(path, cache=True)
+        dest = temporal_src(path)
+        self.download_file(path, dest, cache=True)
         return spectral.io.envi.read_envi_header(dest)
 
     def load_metadata(self, path):
@@ -148,14 +149,18 @@ def load_metadata(self, path):
         return super(Multiple, self).load_metadata(path)
 
     def load_envi(self, path):
-        dest = self.download_file(path, cache=True)
-        raw = self.download_file(f"{path.split('.')[0]}.raw", cache=True)
-        return np.array(spectral.io.envi.open(dest, raw)[:, :, :])
+        dest = temporal_src(path)
+        self.download_file(path, dest, cache=True)
+        raw_path = f"{path.split('.')[0]}.raw"
+        raw_dest = temporal_src(raw_path)
+        self.download_file(raw_path, raw_dest, cache=True)
+        return np.array(spectral.io.envi.open(dest, raw_dest)[:, :, :])
 
     def load_image(self, path):
         if path.lower().endswith('.hdr'):
             return self.load_envi(path)
-        return load_image(self.download_file(path, cache=True))
+        dest = temporal_src(path)
+        return load_image(self.download_file(path, dest, cache=True))
 
     def save_envi(self, path, img, metadata={}):
         src = temporal_src(path)

abraia/training/__init__.py

@@ -6,66 +6,47 @@
 from typing import Dict, Any
 from tqdm.contrib.concurrent import process_map
 
+from ..utils import save_text
 from .dataset import list_datasets, load_dataset, search_images, list_models, download_file
-from ..utils import make_dirs
 
 
 def load_tasks(task):
-    tasks = ['classify', 'detect', 'segment']
     if task:
-        idx = tasks.index(task)
-        return tasks[:idx+1]
+        tasks = ['classify', 'detect', 'segment']
+        return tasks[:tasks.index(task)+1]
     return []
 
 
 def save_annotation(annotation, folder, classes, task):
+    im = Image.open(os.path.join(folder, 'images', annotation['filename']))
+    label_lines = []
+    for object in annotation.get('objects', []):
+        label, box, polygon = object.get('label'), object.get('box'), object.get('polygon')
+        if task == 'segment':
+            if polygon:
+                label_line = f"{classes.index(label)} " + ' '.join([f"{point[0] / im.width} {point[1] / im.height}" for point in polygon])
+                label_lines.append(label_line)
+        elif task == 'detect':
+            if polygon:
+                xx, yy = [point[0] for point in polygon], [point[1] for point in polygon]
+                x1, y1, x2, y2 = min(xx), min(yy), max(xx), max(yy)
+                box = [x1, y1, x2 - x1, y2 - y1]
+            if box:
+                label_line = f"{classes.index(label)} {(box[0] + box[2] / 2) / im.width} {(box[1] + box[3] / 2) / im.height} {box[2] / im.width} {box[3] / im.height}"
+                label_lines.append(label_line)
+    label_path = os.path.join(folder, 'labels',  f"{os.path.splitext(annotation['filename'])[0]}.txt")
+    save_text(label_path, '\n'.join(label_lines))
+
+
+def save_data(annotation, folder, classes, task):
+    path = annotation['path']
+    dest = folder if task == 'classify' else os.path.join(folder, 'images')
     if task == 'classify':
-        for object in annotation.get('objects', []):
-            label = object.get('label')
-            if label:
-                src = os.path.join(folder, annotation['filename'])
-                dest = os.path.join(folder, label, annotation['filename'])
-                shutil.move(src, dest)
-                break
-    else:
-        im = Image.open(os.path.join(folder, 'images', annotation['filename']))
-        label_lines = []
-        for object in annotation.get('objects', []):
-            label, box, polygon = object.get('label'), object.get('box'), object.get('polygon')
-            # Convert polygon or box to yolo format
-            if task == 'segment':
-                if polygon:
-                    label_line = f"{classes.index(label)} " + ' '.join([f"{point[0] / im.width} {point[1] / im.height}" for point in polygon])
-                    label_lines.append(label_line)
-            elif task == 'detect':
-                if polygon:
-                    xx, yy = [point[0] for point in polygon], [point[1] for point in polygon]
-                    x1, y1, x2, y2 = min(xx), min(yy), max(xx), max(yy)
-                    box = [x1, y1, x2 - x1, y2 - y1]
-                if box:
-                    label_line = f"{classes.index(label)} {(box[0] + box[2] / 2) / im.width} {(box[1] + box[3] / 2) / im.height} {box[2] / im.width} {box[3] / im.height}"
-                    label_lines.append(label_line)
-        label_path = os.path.join(folder, 'labels',  f"{os.path.splitext(annotation['filename'])[0]}.txt")
-        with open(label_path, 'w') as f:
-            f.write('\n'.join(label_lines))
-
-
-def save_data(annotations, folder, classes, task):
-    if (task == 'classify'):
-        make_dirs(os.path.join(folder, ''))
-        paths = [annotation['path'] for annotation in annotations]
-        process_map(download_file, paths, itertools.repeat(folder), max_workers=5)
-        for label in classes:
-            make_dirs(os.path.join(folder, label, ''))
-        for annotation in annotations:
-            save_annotation(annotation, folder, classes, task)
-    else:
-        make_dirs(os.path.join(folder, 'images', ''))
-        paths = [annotation['path'] for annotation in annotations]
-        process_map(download_file, paths, itertools.repeat(os.path.join(folder, 'images')), max_workers=5)
-        make_dirs(os.path.join(folder, 'labels', ''))
-        for annotation in annotations:
-            save_annotation(annotation, folder, classes, task)
+        label = next((obj.get('label', '') for obj in annotation.get('objects', [])), '')
+        dest = os.path.join(dest, label)
+    download_file(path, dest)
+    if task != 'classify':
+        save_annotation(annotation, folder, classes, task)
 
 
 def save_config(dataset, classes):
@@ -76,17 +57,20 @@ def save_config(dataset, classes):
     names: {classes}
     '''
     path = os.path.join(dataset, 'data.yaml')
-    with open(path, 'w') as f:
-        f.write(yaml_content)
+    save_text(path, yaml_content)
 
 
 def prepare_dataset(dataset, force = False):
     if force or not os.path.exists(dataset.project):
-        train, val, test = dataset.split()
-        data_annotations = {'train': train, 'val': val, 'test': test}
-        for x in ['train', 'val', 'test']:
-            save_data(data_annotations[x], f"{dataset.project}/{x}", dataset.classes, dataset.task)
-        save_config(dataset.project, dataset.classes)
+        splits = list(zip(['train', 'val', 'test'], dataset.split()))
+        all_annotations, all_folders = [], []
+        for x, annotations in splits:
+            folder = os.path.join(dataset.project, x)
+            all_annotations.extend(annotations)
+            all_folders.extend([folder] * len(annotations))
+        process_map(save_data, all_annotations, all_folders, itertools.repeat(dataset.classes), itertools.repeat(dataset.task), max_workers=5, chunksize=1, desc="Downloading images")
+        if dataset.task != 'classify':
+            save_config(dataset.project, dataset.classes)
 
 
 class ModelTrainer:

abraia/training/detect.py

@@ -2,8 +2,8 @@
 
 import os
 import io
+import shutil
 import contextlib
-import numpy as np
 #os.environ['YOLO_VERBOSE'] = 'False'
 
 from ultralytics import YOLO
@@ -61,7 +61,8 @@ def save(self, project, classes, device='cpu', half=False):
         out = io.StringIO()
         with contextlib.redirect_stdout(out):
             model_src = self.model.export(format="onnx", device=device, opset=11, half=half)
-        abraia.upload_file(model_src, f"{project}/{self.model_name}.onnx")
+            shutil.copy(model_src, f"{self.model_name}.onnx")
+        abraia.upload_file(f"{self.model_name}.onnx", f"{project}/{self.model_name}.onnx")
         abraia.save_json(f"{project}/{self.model_name}.json", 
                          {'task': self.task, 'inputShape': [1, 3, self.imgsz, self.imgsz], 
                           'classes': classes, 'metrics': self.metrics})
@@ -83,6 +84,7 @@ def run(self, img):
         return objects
 
     def compile(self, project, classes, device='hailo8'):
+        abraia.download_file(f"{project}/{self.model_name}.onnx", f"{self.model_name}.onnx")
         print("Compile model for edge deployment to hailo hef format...")
         print(f"hailomz compile yolov8n --ckpt yolov8n.onnx --calib-path {project}/train/images --classes {len(classes)} --hw-arch {device} --performance")
 

abraia/utils/__init__.py

@@ -122,6 +122,18 @@ def save_json(data, dest, gz=False):
     return dest
 
 
+def load_text(src, gz=False):
+    with gzip.open(src, 'rt') if gz else open(src, 'r') as f:
+        return f.read()
+
+
+def save_text(dest, text, gz=False):
+    make_dirs(dest)
+    with gzip.open(dest, 'wt') if gz else open(dest, 'w') as f:
+        f.write(text)
+    return dest
+
+
 def load_data(src, gz=False):
     with gzip.open(src, 'rb') if gz else open(src, 'rb') as f:
         return f.read()

images/screenshot.jpg

@@ -2,7 +2,7 @@ click>=8.0.3
 numpy>=1.21.0
 Pillow>=8.0.0
 filetype>=1.0.7
-onnxruntime==1.18.0
+onnxruntime>=1.18.0
 opencv-python-headless>=4.10.0.84
 python-dotenv>=0.21.1
 requests>=2.27.1

requirements.txt

@@ -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.230', '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.230', 'onnx>=1.16.0', 'transformers>=4.57.1'],
 }
 
 setup(