Add clean image to cli

README.md

@@ -2,13 +2,13 @@
 [![Python Package](https://img.shields.io/pypi/v/abraia.svg)](https://pypi.org/project/abraia/)
 ![Package Downloads](https://img.shields.io/pypi/dm/abraia)
 
-# Abraia Python SDK image analysis toolbox
+# Abraia SDK and CLI
 
-The Abraia Python SDK provides and easy and practical way to develop and deploy Machine Learning image applications on the edge. You can easily annotate and train your custom deep learning model with [DeepLab](https://abraia.me/deeplab/), and deploy the model with this Python SDK.
+The Abraia SDK and CLI is a Python package which provides a set of tools to develop and deploy advanced Machine Learning image applications on the edge. Moreover, with [Abraia DeepLab](https://abraia.me/deeplab/) you can easily annotate and train, your own versions of some of the best state of the art deep learning models, and get them ready to deploy with this Python SDK.
 
 ![people walking](https://github.com/abraia/abraia-multiple/raw/master/images/people-walking.gif)
 
-Just install the Abraia Python SDK and CLI on Windows, Mac, or Linux:
+Just install the Abraia SDK and CLI on Windows, Mac, or Linux:
 
 ```sh
 python -m pip install -U abraia
@@ -82,32 +82,6 @@ save_image(out, 'images/rolling-stones-identified.jpg')
 
 ![rolling stones identified](https://github.com/abraia/abraia-multiple/raw/master/images/rolling-stones-identified.jpg)
 
-### License plates blurring
-
-Automatically blur car license plates in videos with just a few lines of code.
-
-```python
-import numpy as np
-
-from abraia import detect
-from abraia import draw
-
-model_uri = 'multiple/models/alpd-seg.onnx'
-model = detect.load_model(model_uri)
-
-src = 'images/cars.mp4'
-video = detect.Video(src, output='images/blur.mp4')
-for k, frame in enumerate(video):
-    results = model.run(frame, approx=0.02)
-    mask = np.zeros(frame.shape[:2], np.uint8)
-    [draw.draw_filled_polygon(mask, result['polygon'], 255) for result in results]
-    frame = draw.draw_blurred_mask(frame, mask)
-    video.write(frame)
-    video.show(frame)
-```
-
-![car license plate blurred](https://github.com/abraia/abraia-multiple/raw/master/images/blur.jpg)
-
 ### License plates recognition
 
 Automatically recognize car license plates in images and video streams.
@@ -131,34 +105,6 @@ show_image(img)
 
 ![car license plate recognition](https://github.com/abraia/abraia-multiple/raw/master/images/car-plate.jpg)
 
-## Remove unwanted objects
-
-Directly remove unwanted objects in images and photos locally. Just click on the object and press the "spacebar" to automatically select and delete the object from the image. Finally, press "s" to save the final image.
-
-```python
-from abraia.utils import load_image, Sketcher
-from abraia.editing.inpaint import LAMA
-from abraia.editing.sam import SAM
-
-
-img = load_image('images/dog.jpg')
-
-sam = SAM()
-lama = LAMA()
-sam.encode(img)
-
-sketcher = Sketcher(img)
-
-def on_click(point):
-    mask = sam.predict(img, f'[{{"type":"point","data":[{point[0]},{point[1]}],"label":1}}]')
-    sketcher.mask = sketcher.dilate(mask)
-
-sketcher.on_click(on_click)
-sketcher.run(lama.predict)
-```
-
-![inpaint output](https://github.com/abraia/abraia-multiple/raw/master/images/inpaint-output.jpg)
-
 ## Gender Age model
 
 Model to predict gender and age. It can be useful to anonymize minors faces.
@@ -185,6 +131,16 @@ show_image(img)
 
 The Abraia CLI provides access to the Abraia Cloud Platform through the command line. It makes simple to manage your files and enables bulk image editing capabilities. It provides and easy way to resize, convert, and compress your images - JPEG, WebP, or PNG -, and get them ready to publish on the web. Moreover, you can automatically remove the background, upscale, or anonymize your images in bulk.
 
+### Remove unwanted objects
+
+Remove unwanted objects in images and photos locally. Just click on the object to automatically select and delete it from the image. Finally, press "s" to save the output image.
+
+```sh
+abraia editing clean dog.jpg
+```
+
+![inpaint output](https://github.com/abraia/abraia-multiple/raw/master/images/inpaint-output.jpg)
+
 ### Remove background
 
 Automatically remove images background and make them transparent in bulk.
@@ -207,7 +163,7 @@ abraia editing upscale "*.jpg"
 
 ### Anonymize images
 
-Anonymize images in bulk, automatically blurring faces, car license plates, and removing metadata.
+Automatically blur car license plates and faces to anonymize images in bulk.
 
 ```sh
 abraia editing anonymize "*.jpg"

abraia/__init__.py

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

abraia/editing/__init__.py

@@ -5,10 +5,11 @@
 from .upscale import ESRGAN, SwinIR
 from .smartcrop import Smartcrop
 from .inpaint import LAMA
+from .sam import SAM
 
 from ..detect import load_model
 from ..faces import Recognition
-from ..utils import draw
+from ..utils import draw, Sketcher
 
 
 def detect_faces(img):
@@ -65,4 +66,20 @@ def smartcrop_image(img, size):
 
 def inpaint_image(img, mask):
     lama = LAMA()
-    return lama.predict(img, mask) 
+    return lama.inpaint(img, mask)
+
+
+def clean_image(img):
+    sam = SAM()
+    lama = LAMA()
+    sam.encode(img)
+
+    def handle_click(point):
+        mask = sam.predict(img, f'[{{"type":"point","data":[{point[0]},{point[1]}],"label":1}}]')
+        sketcher.mask = cv2.bitwise_or(sketcher.dilate(mask), sketcher.mask)
+        sketcher.show(sketcher.img, sketcher.mask)
+        return lama.inpaint(img, sketcher.mask)
+
+    sketcher = Sketcher(img)
+    sketcher.on_click(handle_click)
+    sketcher.run()

abraia/editing/inpaint.py

@@ -5,6 +5,37 @@
 from ..utils import download_file
 
 
+def tiling(img, tile_size, overlap = 8):
+    """Return an image in a tiled manner."""
+    tile_width, tile_height = tile_size
+    height, width = img.shape[:2]
+    for y in range(0, height, tile_height - overlap):
+        y_end = min(y + tile_height, height)
+        y = max(0, y_end - tile_height)
+        for x in range(0, width, tile_width - overlap):
+            x_end = min(x + tile_width, width)
+            x = max(0, x_end - tile_width)
+            tile = img[y:y_end, x:x_end]
+            yield tile
+
+
+def stitching(tiles, img_size, overlap = 8):
+    """Return an image from a set of tiles."""
+    tile_height, tile_width = tiles[0].shape[:2]
+    width, height = img_size
+    out = np.empty((height, width, 3), dtype=np.uint8)
+    k = 0
+    for y in range(0, height, tile_height - overlap):
+        for x in range(0, width, tile_width - overlap):
+            x_end = min(x + tile_width, width)
+            y_end = min(y + tile_height, height)
+            x = max(0, x_end - tile_width)
+            y = max(0, y_end - tile_height)
+            out[y:y_end, x:x_end, :] = tiles[k]
+            k += 1
+    return out
+
+
 def ceil_modulo(x, mod):
     return x if x % mod == 0 else (x // mod + 1) * mod
 
@@ -45,10 +76,19 @@ def postprocess(self, output, size):
         output = cv2.resize(output, size, interpolation=cv2.INTER_CUBIC)
         return output
 
-    def predict(self, image, mask):
-        h, w = image.shape[:2]
-        image, mask = self.preprocess(image, mask)
+    def process(self, img, mask):
+        h, w = img.shape[:2]
+        image, mask = self.preprocess(img, mask)
         outputs = self.session.run(None, {'image': image, 'mask': mask})
         output = self.postprocess(outputs[0], (w, h))
         return output
+
+    def inpaint(self, img, mask):
+        outs = []
+        tile_size = self.image_size
+        img_size = (img.shape[1], img.shape[0])
+        for im, msk in zip(tiling(img, tile_size), tiling(mask, tile_size)):
+            outs.append(self.process(im, msk))
+        out = stitching(outs, img_size)
+        return out
 

abraia/editing/upscale.py

@@ -43,14 +43,12 @@ def create_gradient_mask(shape, feather):
 
 def tiled_upscale(samples, function, scale, tile_size, overlap = 8):
     """Apply a scaling function to image samples in a tiled manner."""
+    height, width = samples.shape[2:]
     tile_width, tile_height = tile_size
-    _batch, _channels, height, width = samples.shape
     out_height, out_width = round(height * scale), round(width * scale)
-    # Initialize output tensors
     output = np.empty((1, 3, out_height, out_width))
     out = np.zeros((1, 3, out_height, out_width))
     out_div = np.zeros_like(output)
-    # Process the image in tiles
     for y in range(0, height, tile_height - overlap):
         for x in range(0, width, tile_width - overlap):
             # Ensure we don't go out of bounds

abraia/faces/transform.py

@@ -23,16 +23,18 @@ def similarity_transform(src_pts, ref_pts):
     return M
 
 
-def affine_transform(src_pts, ref_pts):
-    # M, _ = cv2.estimateAffine2D(src_pts, ref_pts)
-    M, _ = cv2.estimateAffinePartial2D(src_pts, ref_pts)
+def affine_transform(src_pts, dst_pts):
+    src_tri = np.array([src_pts[0], src_pts[1], (src_pts[3] + src_pts[4]) / 2]).astype(np.float32)
+    dst_tri = np.array([dst_pts[0], dst_pts[1], (dst_pts[3] + dst_pts[4]) / 2]).astype(np.float32)
+    M = cv2.getAffineTransform(src_tri, dst_tri)
+    # M, _ = cv2.estimateAffinePartial2D(src_pts, dst_pts)
     return M
 
 
 def align_face(img, src_pts, size):
     dst_pts = ref_pts * size / 112 if size != 112 else ref_pts
-    M = similarity_transform(src_pts, dst_pts)
-    # M = affine_transform(src_pts, ref_pts)
+    # M = similarity_transform(src_pts, dst_pts)
+    M = affine_transform(src_pts, dst_pts)
     return cv2.warpAffine(img, M, (size, size), borderValue=0.0)
 
 

abraia/utils/sketcher.py

@@ -1,10 +1,9 @@
 '''
-Sketcher.
+Magic Eraser.
 
 Keys:
-  SPACE - callback
-  r     - reset the mask
-  s     - save output
+  s     - save & exit
+  r     - reset
   ESC   - exit
 '''
 
@@ -15,7 +14,7 @@
 
 
 class Sketcher:
-    def __init__(self, img, radius=11):
+    def __init__(self, img, radius=7):
         print(__doc__)
         self.prev_pt = None
         self.win_name = 'Image'
@@ -38,34 +37,24 @@ def show(self, img, mask=None):
             img = draw_overlay_mask(img, mask, (255, 0, 0), 0.5)
         self.output = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
         cv2.imshow(self.win_name, self.output)
+        cv2.waitKey(1)
 
     def on_click(self, callback):
         self.handle_click = callback
 
     def on_mouse(self, event, x, y, flags, param):
-        pt = (x, y)
         if event == cv2.EVENT_LBUTTONDOWN:
-            self.prev_pt = pt
-        if self.prev_pt and flags & cv2.EVENT_FLAG_LBUTTON:
-            cv2.line(self.mask, self.prev_pt, pt, 255, self.radius)
-            self.prev_pt = pt
-        else:
-            self.prev_pt = None
-        if event == cv2.EVENT_LBUTTONUP:
             if self.handle_click:
-                self.handle_click(pt)
-        if self.prev_pt:
-            self.show(self.img, self.mask)
+                self.show(self.handle_click([x, y]))
 
-    def run(self, callback):
+    def run(self):
         while True:
             ch = 0xFF & cv2.waitKey()
             if ch == 27 or ch == ord('q'):
                 break
-            if ch == ord(' '):
-                self.show(callback(self.img, self.mask))
             if ch == ord('r'):
                 self.load(self.img)
             if ch == ord('s'):
                 cv2.imwrite('output.png', self.output)
+                break
         cv2.destroyWindow(self.win_name)

scripts/abraia

@@ -10,6 +10,9 @@ from tqdm import tqdm
 from urllib.parse import urlparse
 from concurrent.futures import ProcessPoolExecutor
 
+from abraia.utils import load_image
+from abraia.editing import clean_image
+
 from abraia import config
 from abraia import Abraia
 from abraia import APIError
@@ -146,11 +149,19 @@ def upscale(src):
 
 @cli_editing.command()
 @click.argument('src')
-def anonymize(src, ):
+def anonymize(src):
     """Anonymize image blurring faces and car license plates."""
     editing_files(src, 'anonymize', desc="Anonymizing")
 
 
+@cli_editing.command()
+@click.argument('src')
+def clean(src):
+    """Clean images removing unwanted objects with inpating."""
+    img = load_image(src)
+    clean_image(img)
+
+
 @cli.group('files')
 def cli_files():
     """Commands related to cloud storage."""