Add compile to cli

Author: jrrodri

abraia/__init__.py

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

abraia/utils/draw.py

@@ -207,3 +207,15 @@ def render_region(img, region, text='', color=(255, 0, 0), opacity=0.2):
     draw_filled_polygon(img, region, color=color, opacity=opacity)
     draw_text(img, text, point, background_color=color, text_scale=text_scale)
     return img
+
+
+def render_resolution(img):
+    height, width = img.shape[:2]
+    text = f"{width}x{height}"
+    thickness = calculate_optimal_thickness(img.shape[:2])
+    text_scale = calculate_optimal_text_scale(img.shape[:2])
+    text_font, text_thickness = cv2.FONT_HERSHEY_DUPLEX, 1
+    w, h = cv2.getTextSize(text, text_font, text_scale, text_thickness)[0]
+    padding = thickness * 3
+    point = (width - w - 2 * padding, height)
+    return draw_text(img, text, point, background_color=(128, 128, 128), text_scale=text_scale, padding=padding)

abraia/utils/video.py

@@ -2,7 +2,7 @@
 import cv2
 import time
 import psutil
-from .draw import draw_text_multiline
+from .draw import draw_text_multiline, render_resolution
 
 
 def is_raspberry():
@@ -24,7 +24,7 @@ def __init__(self, src=0, resolution=(1920, 1080), fps=30, dest=None):
             from libcamera import Transform
             self.picam2 = Picamera2()
             self.picam2.configure(self.picam2.create_video_configuration(
-                main={"format": 'BGR888', "size": resolution},
+                main={"format": 'RGB888', "size": resolution},
                 transform=Transform(hflip=True, vflip=True),
                 controls={"FrameRate": fps}))
             self.picam2.start()
@@ -94,12 +94,13 @@ def show(self, frame):
         draw_text_multiline(frame, [f"FPS: {round(1 / (t1 - self.t0), 1)}", 
                                     f"CPU: {psutil.cpu_percent()}%", 
                                     f"RAM: {round(psutil.virtual_memory().used / (1024**3), 2)} GB"], (10, 40), background_color=(192, 192, 192))
+        render_resolution(frame)
         self.t0 = t1
         if not self.win_name:
             self.win_name = 'Video'
             cv2.namedWindow(self.win_name, cv2.WINDOW_GUI_NORMAL)
         cv2.imshow(self.win_name, cv2.cvtColor(frame, cv2.COLOR_RGB2BGR))
-        ch = 0xFF & cv2.waitKey(1 if self.picam2 else int(self.fps))
+        ch = 0xFF & cv2.waitKey(int(self.fps))
         if (ch == 27 or ch == ord('q')) or cv2.getWindowProperty(self.win_name, cv2.WND_PROP_VISIBLE) < 1:
             self.quit = True
 

requirements.txt

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

scripts/abraia

@@ -237,6 +237,19 @@ def train(project, epochs):
     training_session.save()
 
 
+@cli.command()
+@click.argument('project')
+@click.option('--device', help='Compilation device (e.g. hailo8)', default='hailo8')
+def compile(project, device):
+    """Compile a model for edge deployment."""
+    from abraia.training import load_dataset, ModelTrainer
+    click.echo('Loading dataset...')
+    dataset = load_dataset(project)
+    click.echo('Compiling model...')
+    training_session = ModelTrainer(project, dataset.task, dataset.classes)
+    training_session.compile(device=device)
+
+
 @cli.command()
 @click.argument('project')
 @click.argument('classes', required=False, default='')