Merge pull request #10 from nuletizia/main

Skin tuner endpoint

README.md

@@ -5,7 +5,7 @@
 </p>
 
 
-# EraseID - v2.2.0
+# EraseID - v2.3.0
 [![Official Website](https://img.shields.io/badge/Official%20Website-piktid.com-blue?style=flat&logo=world&logoColor=white)](https://piktid.com)
 [![Discord Follow](https://dcbadge.vercel.app/api/server/FJU39e9Z4P?style=flat)](https://discord.com/invite/FJU39e9Z4P)
 
@@ -107,7 +107,7 @@ It is possible to change the default generation parameters, to do that use the c
 $ python3 main.py --filepath 'mydir/myfile.jpg' --identity_filepath 'mydir/myfile.jpg' --identity_name 'myidentityname' --guidance_scale '1.5' --controlnet_scale '0.1' --prompt_strength '0.5'
 ```
 
-## Change facial expression (keeping the identity)
+## Change facial expression
 <a target="_blank" href="https://colab.research.google.com/drive/1d6YT3pt7M4bacAgy0zdr-qYjS57KymLw?usp=sharing">
   <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/>
 </a>
@@ -119,5 +119,12 @@ Choose the EXPRESSION value from the ones available in `cfe_keywords.py`
 $ python3 main.py --all_faces --change_expression_flag --new_expression EXPRESSION
 ```
 
+## Full-body skin tuner (BETA)
+It is now possible to automatically adapt the full-body skin tone to the newly generated face! It is particularly useful when generating different ethnicities. 
+
+```bash
+$ python3 main.py --skin
+```
+
 ## Contact
 office@piktid.com

eraseid_api.py

@@ -246,6 +246,27 @@ def update_data_generation_call(data, PARAM_DICTIONARY):
     return data
 
 
+def update_data_skin_call(data, PARAM_DICTIONARY, TOKEN_DICTIONARY):
+
+    CUSTOM_PROMPT_FLAG = PARAM_DICTIONARY.get('CUSTOM_PROMPT_FLAG')
+    SEED = PARAM_DICTIONARY.get('SEED')
+
+    if CUSTOM_PROMPT_FLAG is True:
+        # overwrite the keyword mechanism
+        extra_data = {'description_type': 'txt'}
+        data.update(extra_data)
+
+    OPTIONS_DICT = {}
+    if SEED is not None:
+        OPTIONS_DICT = {**OPTIONS_DICT, 'seed': SEED}
+
+    OPTIONS = json.dumps(OPTIONS_DICT)
+    extra_options = {'options': OPTIONS}
+    data.update(extra_options)
+
+    return data
+
+
 def generation_call(image_address, idx_face, prompt, PARAM_DICTIONARY, TOKEN_DICTIONARY):
 
     SEED = PARAM_DICTIONARY.get('SEED') 
@@ -277,10 +298,9 @@ def generation_call(image_address, idx_face, prompt, PARAM_DICTIONARY, TOKEN_DIC
 
 def change_expression_call(image_address, idx_face, prompt, PARAM_DICTIONARY, TOKEN_DICTIONARY):
 
-    FLAG_SYNC = PARAM_DICTIONARY.get('FLAG_SYNC')
     SEED = PARAM_DICTIONARY.get('SEED')
 
-    data = {'flag_sync': FLAG_SYNC, 'id_image': image_address, 'id_face': idx_face, 'prompt': prompt, 'seed': SEED}
+    data = {'flag_sync': False, 'id_image': image_address, 'id_face': idx_face, 'prompt': prompt, 'seed': SEED}
     # data = update_data_generation_call(data, PARAM_DICTIONARY, TOKEN_DICTIONARY)
     print(f'data to send to cfe: {data}')
 
@@ -301,7 +321,34 @@ def change_expression_call(image_address, idx_face, prompt, PARAM_DICTIONARY, TO
                                  headers={'Authorization': 'Bearer '+TOKEN},
                                  json=data,
                                  )
-    print(response.text)
+    # print(response.text)
+    response_json = json.loads(response.text)
+    return response_json
+
+
+def change_skin_call(image_address, idx_face, idx_generation, prompt, PARAM_DICTIONARY, TOKEN_DICTIONARY):
+
+    data = {'id_image': image_address, 'id_face': idx_face, 'id_generation': idx_generation, 'prompt': prompt}
+    data = update_data_skin_call(data, PARAM_DICTIONARY, TOKEN_DICTIONARY)
+    print(f'data to send to skin editing: {data}')
+
+    # start the generation process given the image parameters
+    TOKEN = TOKEN_DICTIONARY.get('access_token', '')
+    URL_API = TOKEN_DICTIONARY.get('url_api')
+
+    response = requests.post(URL_API+'/ask_generate_skin_full_body',
+                             headers={'Authorization': 'Bearer '+TOKEN},
+                             json=data,
+                             )
+    # if the access token is expired
+    if response.status_code == 401:
+        TOKEN_DICTIONARY = refresh_call(TOKEN_DICTIONARY)
+        TOKEN = TOKEN_DICTIONARY.get('access_token', '')
+        # try with new TOKEN
+        response = requests.post(URL_API+'/ask_generate_skin_full_body',
+                                 headers={'Authorization': 'Bearer '+TOKEN},
+                                 json=data,
+                                 )
     response_json = json.loads(response.text)
     return response_json
 
@@ -472,3 +519,32 @@ def handle_notifications_new_generation(image_id, idx_face, TOKEN_DICTIONARY):
         sleep(60)
 
     return False, {}
+
+
+def handle_notifications_new_skin(image_id, idx_face, TOKEN_DICTIONARY):
+    # check notifications to verify the generation status
+    i = 0
+    while i < 20:  # max 20 iterations -> then timeout
+        i = i+1
+        notifications_list = get_notification_by_name('new_skin', TOKEN_DICTIONARY)
+        notifications_to_remove = [n for n in notifications_list if (n.get('name') == 'new_skin' and n.get('data').get('address') == image_id and n.get('data').get('f') == idx_face and n.get('data').get('msg') == 'done')]
+
+        print(f'notifications_to_remove: {notifications_to_remove}')
+        # remove notifications
+        result_delete = [delete_notification(n.get('id'), TOKEN_DICTIONARY) for n in notifications_to_remove ]
+        # print(result_delete)
+
+        if len(notifications_to_remove) > 0:
+            print(f'replace for face {idx_face} with full skin completed')
+            return True, {**notifications_to_remove[0].get('data', {})}
+
+        # check iteration
+        if i >= 10:
+            print('Timeout. Error in editing skin')
+            return False, {}
+
+        # wait
+        print('waiting for notification...')
+        sleep(30)
+
+    return False, {}

eraseid_utils.py

@@ -1,6 +1,6 @@
 import json
 
-from eraseid_api import open_image_from_url, upload_and_detect_call, upload_reference_face_call, selection_call, get_identities_call, generation_call, handle_notifications_new_generation, get_generated_faces, get_last_generated_face, set_identity_call, replace_call
+from eraseid_api import upload_and_detect_call, upload_reference_face_call, selection_call, get_identities_call, generation_call, change_expression_call, change_skin_call, handle_notifications_new_generation, handle_notifications_new_skin, get_generated_faces, get_last_generated_face, set_identity_call, replace_call
 from cfe_keywords import cfe_dict
 
 
@@ -95,6 +95,8 @@ def process_single_face(idx_face, count, PARAM_DICTIONARY, TOKEN_DICTIONARY):
 
     CHANGE_EXPRESSION_FLAG = PARAM_DICTIONARY.get('CHANGE_EXPRESSION_FLAG')
 
+    CHANGE_SKIN = PARAM_DICTIONARY.get('CHANGE_SKIN')
+
     KEYWORDS_LIST = PARAM_DICTIONARY.get('KEYWORDS_LIST')
 
     image_id = PARAM_DICTIONARY.get('IMAGE_ID')
@@ -114,10 +116,15 @@ def process_single_face(idx_face, count, PARAM_DICTIONARY, TOKEN_DICTIONARY):
     keywords_to_send = KEYWORDS_LIST[count]
     keywords_to_send = json.dumps(keywords_to_send.get('a'))
 
-    print('Generating new faces')
-    response = generation_call(image_address=image_id, idx_face=idx_face, prompt=keywords_to_send, PARAM_DICTIONARY=PARAM_DICTIONARY, TOKEN_DICTIONARY=TOKEN_DICTIONARY)
+    if CHANGE_EXPRESSION_FLAG:
+        print('Changing facial expressions')
+        response = change_expression_call(image_address=image_id, idx_face=idx_face, prompt=keywords_to_send, PARAM_DICTIONARY=PARAM_DICTIONARY, TOKEN_DICTIONARY=TOKEN_DICTIONARY)
+        print(f'Cfe response:{response}')
 
-    print(f'Generation response:{response}')
+    else:
+        print('Generating new faces')
+        response = generation_call(image_address=image_id, idx_face=idx_face, prompt=keywords_to_send, PARAM_DICTIONARY=PARAM_DICTIONARY, TOKEN_DICTIONARY=TOKEN_DICTIONARY)
+        print(f'Generation response:{response}')
 
     # Asynchronous API call
     response_notifications = handle_notifications_new_generation(image_id, idx_face, TOKEN_DICTIONARY)
@@ -131,6 +138,19 @@ def process_single_face(idx_face, count, PARAM_DICTIONARY, TOKEN_DICTIONARY):
     idx_generation_to_replace = [get_last_generated_face(list_generated_faces.get('links'), idx_face)]
     print(f'Replace generation {idx_generation_to_replace}')
 
+    if CHANGE_SKIN:
+        for idx_generation in idx_generation_to_replace:
+
+            print('Editing the skin')
+            response = change_skin_call(image_address=image_id, idx_face=idx_face, idx_generation=idx_generation, prompt=keywords_to_send, PARAM_DICTIONARY=PARAM_DICTIONARY, TOKEN_DICTIONARY=TOKEN_DICTIONARY)
+            print(f'Skin editing response:{response}')
+
+            # Asynchronous API call
+            response_notifications = handle_notifications_new_skin(image_id, idx_face, TOKEN_DICTIONARY)
+            if response_notifications is False:
+                # Error
+                return False
+
     # Store the last generated face as 'pippo'
     if IDENTITY_NAME is None:
         if STORE_IDENTITY_FLAG:
@@ -142,7 +162,6 @@ def process_single_face(idx_face, count, PARAM_DICTIONARY, TOKEN_DICTIONARY):
     links = replace_call(image_id, idx_face, idx_generation_to_replace, TOKEN_DICTIONARY)
 
     # download the output from EraseID
-    # output_img = open_image_from_url(links[-1])
     print(f'Download the generated image here: {links[-1]}')
 
     return True

main.py

@@ -16,18 +16,24 @@
     parser.add_argument('--url', help='Image file url', type=str, default='https://images.pexels.com/photos/733872/pexels-photo-733872.jpeg?auto=compress&cs=tinysrgb&w=1260&h=750&dpr=1')
     parser.add_argument('--filepath', help='Input image file absolute path', type=str, default=None)
 
+    # Consistent identity parameters
     parser.add_argument('--identity_filepath', help='Input identity file absolute path', type=str, default=None)
     parser.add_argument('--identity_name', help='Use the face from the stored identities', default=None)
     parser.add_argument('--store_identity', help='Save the generated identity under the name pippo', action='store_true')
 
+    # Change expression parameters
     parser.add_argument('--change_expression_flag', help='Change only the facial expression', action='store_true')
     parser.add_argument('--new_expression', help='Desired facial expression', type=str, default='happy')
 
+    # Random generation parameters
     parser.add_argument('--guidance_scale', help='Guidance scale', type=str, default=None)
     parser.add_argument('--prompt_strength', help='Description strength', type=str, default=None)
     parser.add_argument('--controlnet_scale', help='Conditioning scale', type=str, default=None)
     parser.add_argument('--seed', help='Generation seed', type=int, default=randint(0, 100000))
 
+    # Skin parameters
+    parser.add_argument('--skin', help='Change also the skin', action='store_true')
+
     args = parser.parse_args()
 
     # be sure to export your email and psw as environmental variables
@@ -53,6 +59,9 @@
     CONTROLNET_SCALE = args.controlnet_scale
     SEED = args.seed
 
+    # Skin parameters
+    CHANGE_SKIN = args.skin  # True if also the skin is anonymized
+
     # Image parameters
     URL = args.url 
     INPUT_PATH = args.filepath
@@ -100,6 +109,7 @@
             'CONTROLNET_SCALE': CONTROLNET_SCALE,
             'CHANGE_EXPRESSION_FLAG': CHANGE_EXPRESSION_FLAG,
             'NEW_EXPRESSION': NEW_EXPRESSION,
+            'CHANGE_SKIN': CHANGE_SKIN,
         }
 
     response = process_single_image(input_image, PARAM_DICTIONARY, TOKEN_DICTIONARY)