Feature/emojich rgba

.gitlab-ci.yml

@@ -11,6 +11,7 @@ all_branch_test:
   - apt-get install ffmpeg libsm6 libxext6 -y
   - pip install cython
   - pip install -r requirements-test.txt --no-cache-dir
+  - pip install timm==0.4.12
   - pip install codecov
   - pytest --cov=rudalle tests/
   - bash <(curl -s https://codecov.io/bash) -t $CODECOV_TOKEN

Emojich.md

@@ -51,6 +51,19 @@ show(pil_images, 4)
 ```
 ![](./pics/emojich/emoji-Donald.png)
 
+### Converting to Telegram Stickers format (512x512 RGBA)
+Also, don't forget you can use super resolution before :)
+```python
+from rudalle.pipelines import convert_emoji_to_rgba, show_rgba
+from rudalle import get_emojich_unet
+
+device = 'cuda'
+emojich_unet = get_emojich_unet('unet_effnetb5').to(device)
+rgba_images = convert_emoji_to_rgba(pil_images, emojich_unet,  device=device)
+for rgba_image in rgba_images:
+    show_rgba(rgba_image);
+```
+![](./pics/emojich/emojich-stickers.png)
 
 ### Examples of generated emojis
 

README.md

@@ -7,7 +7,7 @@
 [![pre-commit.ci status](https://results.pre-commit.ci/badge/github/sberbank-ai/ru-dalle/master.svg)](https://results.pre-commit.ci/latest/github/sberbank-ai/ru-dalle/master)
 
 ```
-pip install rudalle==0.1.0
+pip install rudalle==0.2.1
 ```
 ### 🤗 HF Models:
 [ruDALL-E Malevich (XL)](https://huggingface.co/sberbank-ai/rudalle-Malevich) \

requirements.txt

@@ -5,6 +5,8 @@ youtokentome~=1.0.6
 omegaconf>=2.0.0
 einops~=0.3.2
 PyWavelets==1.1.1
+segmentation-models-pytorch==0.1.3
+opencv-python==4.5.4.60
 torch
 torchvision
 matplotlib

rudalle/__init__.py

@@ -4,6 +4,7 @@
 from .tokenizer import get_tokenizer
 from .realesrgan import get_realesrgan
 from .ruclip import get_ruclip
+from .emojich_unet import get_emojich_unet
 from . import vae, dalle, tokenizer, realesrgan, pipelines, ruclip, image_prompts
 
 
@@ -13,6 +14,7 @@
     'get_tokenizer',
     'get_realesrgan',
     'get_ruclip',
+    'get_emojich_unet',
     'vae',
     'dalle',
     'ruclip',
@@ -22,4 +24,4 @@
     'image_prompts',
 ]
 
-__version__ = '0.1.0'
+__version__ = '0.2.1'

rudalle/emojich_unet/__init__.py

@@ -0,0 +1,34 @@
+# -*- coding: utf-8 -*-
+import os
+
+import torch
+import segmentation_models_pytorch as smp
+from huggingface_hub import hf_hub_url, cached_download
+
+
+MODELS = {
+    'unet_effnetb5': dict(
+        encoder_name='efficientnet-b5',
+        repo_id='sberbank-ai/rudalle-Emojich',
+        filename='pytorch_model.bin',
+    ),
+}
+
+
+def get_emojich_unet(name, cache_dir='/tmp/rudalle'):
+    assert name in MODELS
+    config = MODELS[name]
+    model = smp.Unet(
+        encoder_name=config['encoder_name'],
+        encoder_weights=None,
+        in_channels=3,
+        classes=1,
+    )
+    cache_dir = os.path.join(cache_dir, name)
+    filename = config['filename']
+    config_file_url = hf_hub_url(repo_id=config['repo_id'], filename=f'{name}/{filename}')
+    cached_download(config_file_url, cache_dir=cache_dir, force_filename=filename)
+    checkpoint = torch.load(os.path.join(cache_dir, config['filename']), map_location='cpu')
+    model.load_state_dict(checkpoint)
+    print(f'{name} --> ready')
+    return model

rudalle/pipelines.py

@@ -3,13 +3,15 @@
 from glob import glob
 from os.path import join
 
+import cv2
 import torch
 import torchvision
 import transformers
 import more_itertools
 import numpy as np
 import matplotlib.pyplot as plt
 from tqdm.auto import tqdm
+from PIL import Image
 
 from . import utils
 
@@ -114,3 +116,43 @@ def show(pil_images, nrow=4, size=14, save_dir=None, show=True):
     if show:
         fix.show()
         plt.show()
+
+
+def convert_emoji_to_rgba(pil_images, emojich_unet,  device='cpu', bs=4):
+    final_images = []
+    with torch.no_grad():
+        for chunk in more_itertools.chunked(pil_images, bs):
+            images = []
+            for pil_image in chunk:
+                image = np.array(pil_image.resize((512, 512)))[:, :, :3]
+                image = image.astype(np.float32) / 255.0
+                image = torch.from_numpy(image).permute(2, 0, 1)
+                images.append(image)
+            images = torch.nn.utils.rnn.pad_sequence(images, batch_first=True)
+            pred_masks = emojich_unet(images.to(device))[:, 0, :, :]
+            pred_masks = torch.sigmoid(pred_masks)
+            pred_masks = (pred_masks > 0.5).int().cpu().numpy()
+            pred_masks = (pred_masks * 255).astype(np.uint8)
+            for pil_image, pred_mask in zip(chunk, pred_masks):
+                ret, thresh = cv2.threshold(pred_mask, 0, 255, 0)
+                contours, hierarchy = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_NONE)
+                cv2.drawContours(pred_mask, contours, -1, (0, 0, 0), 1)
+                final_image = np.zeros((512, 512, 4), np.uint8)
+                final_image[:, :, :3] = np.array(pil_image.resize((512, 512)))[:, :, :3]
+                final_image[:, :, -1] = pred_mask
+                final_image = Image.fromarray(final_image)
+                final_images.append(final_image)
+    return final_images
+
+
+def show_rgba(rgba_pil_image):
+    img = np.array(rgba_pil_image)
+    fig, ax = plt.subplots(1, 3, figsize=(10, 10), dpi=100)
+    ax[0].imshow(img[:, :, :3])
+    ax[1].imshow(img[:, :, -1])
+    mask = np.repeat(np.expand_dims(img[:, :, -1] < 128, -1), 3, axis=-1)
+    img = img[:, :, :3]
+    img[mask[:, :, 0], 0] = 64
+    img[mask[:, :, 0], 1] = 255
+    img[mask[:, :, 0], 2] = 64
+    ax[2].imshow(img)

setup.py

@@ -47,7 +47,8 @@ def get_version():
     author='SberAI, SberDevices',
     author_email='[email protected]',
     description='ruDALL-E generate images from texts in Russian language',
-    packages=['rudalle', 'rudalle/dalle', 'rudalle/realesrgan', 'rudalle/ruclip', 'rudalle/vae'],
+    packages=['rudalle', 'rudalle/dalle', 'rudalle/realesrgan', 'rudalle/ruclip', 'rudalle/vae',
+              'rudalle/emojich_unet'],
     package_data={'rudalle/vae': ['*.yml']},
     install_requires=get_requirements(),
     dependency_links=get_links(),

tests/conftest.py

@@ -6,7 +6,7 @@
 import pytest
 import requests
 
-from rudalle import get_tokenizer, get_rudalle_model, get_vae, get_realesrgan
+from rudalle import get_tokenizer, get_rudalle_model, get_vae, get_realesrgan, get_emojich_unet
 
 
 TEST_ROOT = dirname(abspath(__file__))
@@ -48,4 +48,10 @@ def sample_image():
 @pytest.fixture(scope='module')
 def small_dalle():
     model = get_rudalle_model('small', pretrained=False, fp16=False, device='cpu')
-    return model
+    yield model
+
+
+@pytest.fixture(scope='module')
+def emojich_unet():
+    model = get_emojich_unet('unet_effnetb5')
+    yield model

tests/test_emojich_unet.py

@@ -0,0 +1,13 @@
+# -*- coding: utf-8 -*-
+import numpy as np
+
+from rudalle.pipelines import convert_emoji_to_rgba
+
+
+def test_convert_emoji_to_rgba(sample_image, emojich_unet):
+    img = sample_image.copy()
+    img = img.resize((512, 512))
+    rgba_img = convert_emoji_to_rgba([img], emojich_unet)[0]
+    assert rgba_img.size[0] == 512
+    assert rgba_img.size[1] == 512
+    assert np.array(rgba_img).shape[-1] == 4