Refactor to group inference models

README.md

@@ -25,7 +25,7 @@ Annotate your images and train a state-of-the-art model for classification, dete
 Detect objects with a pre-trained YOLOv8 model on images, videos, or even camera streams.
 
 ```python
-from abraia import detect
+from abraia.inference import detect
 
 model_uri = f"multiple/models/yolov8n.onnx"
 model = detect.load_model(model_uri)
@@ -41,7 +41,7 @@ detect.show_image(img)
 To run a multi-object detector on video or directly on a camera stream, you just need to use the Video class to process every frame as is done for images.
 
 ```python
-from abraia import detect
+from abraia.inference import detect
 
 model_uri = f"multiple/models/yolov8n.onnx"
 model = detect.load_model(model_uri)
@@ -60,13 +60,13 @@ Identify people on images with face recognition as shown bellow.
 ```python
 import os
 
-from abraia.faces import Recognition
+from abraia.inference import FaceRecognizer
 from abraia.utils import load_image, save_image, render_results
 
 img = load_image('images/rolling-stones.jpg')
 out = img.copy()
 
-recognition = Recognition()
+recognition = FaceRecognizer()
 results = recognition.represent_faces(img)
 
 index = []
@@ -87,10 +87,10 @@ save_image(out, 'images/rolling-stones-identified.jpg')
 Automatically recognize car license plates in images and video streams.
 
 ```python
-from abraia.alpr import ALPR
+from abraia.inference import PlateRecognizer
 from abraia.utils import load_image, show_image, render_results
 
-alpr = ALPR()
+alpr = PlateRecognizer()
 
 img = load_image('images/car.jpg')
 results = alpr.detect(img)
@@ -110,11 +110,11 @@ show_image(img)
 Model to predict gender and age. It can be useful to anonymize minors faces.
 
 ```python
-from abraia.faces import Recognition, Attribute
+from abraia.inference import FaceRecognizer, FaceAttribute
 from abraia.utils import load_image, show_image, render_results
 
-recognition = Recognition()
-attribute = Attribute()
+recognition = FaceRecognizer()
+attribute = FaceAttribute()
 
 img = load_image('images/image.jpg')
 results = recognition.detect_faces(img)

abraia/__init__.py

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

abraia/editing/__init__.py

@@ -1,25 +1,25 @@
 import cv2
 import numpy as np
 
-from .removebg import RemoveBG
-from .upscale import ESRGAN, SwinIR
+from .removebg import BackgroundRemover
+from .upscale import Upscaler, SwinIR
 from .smartcrop import Smartcrop
 from .inpaint import LAMA
 from .sam import SAM
 
-from ..detect import load_model
-from ..faces import Recognition
+from ..inference import PlateDetector
+from ..inference.faces import FaceRecognizer
 from ..utils import draw, Sketcher
 
 
 def detect_faces(img):
-    recognition = Recognition()
+    recognition = FaceRecognizer()
     return recognition.detect_faces(img)
 
 
 def detect_plates(img):
-    detection = load_model('multiple/models/alpd-seg.onnx')
-    return detection.run(img, approx=0.02)
+    plate = PlateDetector()
+    return plate.detect(img)
 
 
 def detect_smartcrop(img, size):
@@ -43,7 +43,7 @@ def anonymize_image(img):
 
 
 def remove_background(img):
-    removebg = RemoveBG()
+    removebg = BackgroundRemover()
     out = removebg.remove(img)
     return out
 
@@ -62,7 +62,7 @@ def upscale_image(img):
         scale = 1920 / max(img.shape)
         size = (round(scale * w), round(scale * h))
         img = cv2.resize(img, size, cv2.INTER_LINEAR)
-    upscaler = ESRGAN()
+    upscaler = Upscaler()
     out = upscaler.upscale(img)
     return out
 
@@ -91,3 +91,6 @@ def handle_click(point):
     sketcher = Sketcher(img)
     sketcher.on_click(handle_click)
     sketcher.run()
+
+
+# __all__ = [clean_image]

abraia/editing/removebg.py

@@ -29,7 +29,7 @@ def naive_cutout(img, mask):
     return cutout
 
 
-class RemoveBG:
+class BackgroundRemover:
 
     def __init__(self):
         self.image_size = (1024, 1024)

abraia/editing/smartcrop.py

@@ -3,7 +3,7 @@
 import numpy as np
 from PIL import Image
 
-from ..faces import Retinaface
+from ..inference.faces import Retinaface
 
 
 def spectral_residual_mono(gray):

abraia/editing/upscale.py

@@ -8,7 +8,6 @@
 
 
 class SwinIR:
-
     def __init__(self):
         model_src = download_file('multiple/models/editing/003_realSR_BSRGAN_DFO_s64w8_SwinIR-M_x4_GAN.onnx')
         self.session = ort.InferenceSession(model_src, sess_options)
@@ -72,8 +71,7 @@ def tiled_upscale(samples, function, scale, tile_size, overlap = 8):
     return output
 
 
-class ESRGAN:
-
+class Upscaler:
     def __init__(self, overlap = 8):
         self.scale = 2
         self.overlap = overlap