add masking api

CHANGELOG.md

@@ -42,7 +42,10 @@
   - support models with their own YAML model config files
   - support models with their own JSON per-component config files, for example: `playground-v2.5_vae.config`
 - **API**
-  - add preprocessor api endpoints: `/sdapi/v1/preprocessors`, `/sdapi/v1/preprocess`
+  - add preprocessor api endpoints  
+    GET:`/sdapi/v1/preprocessors`, POST:`/sdapi/v1/preprocess`, sample script:`cli/simple-preprocess.py`
+  - add masking api endpoints  
+    GET:`/sdapi/v1/masking`, POST:`/sdapi/v1/mask`, sample script:`cli/simple-mask.py`
 - **Internal**
   - remove obsolete textual inversion training code
   - remove obsolete hypernetworks training code

TODO.md

@@ -18,4 +18,3 @@ Main ToDo list can be found at [GitHub projects](https://github.com/users/vladma
 
 - second pass: <https://github.com/vladmandic/automatic/issues/2783>  
 - control api  
-- masking api  

cli/simple-mask.py

@@ -0,0 +1,78 @@
+#!/usr/bin/env python
+import io
+import os
+import time
+import base64
+import logging
+import argparse
+import requests
+import urllib3
+from PIL import Image
+
+
+sd_url = os.environ.get('SDAPI_URL', "http://127.0.0.1:7860")
+sd_username = os.environ.get('SDAPI_USR', None)
+sd_password = os.environ.get('SDAPI_PWD', None)
+
+
+logging.basicConfig(level = logging.INFO, format = '%(asctime)s %(levelname)s: %(message)s')
+log = logging.getLogger(__name__)
+urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+
+
+def auth():
+    if sd_username is not None and sd_password is not None:
+        return requests.auth.HTTPBasicAuth(sd_username, sd_password)
+    return None
+
+
+def get(endpoint: str, dct: dict = None):
+    req = requests.get(f'{sd_url}{endpoint}', json=dct, timeout=300, verify=False, auth=auth())
+    if req.status_code != 200:
+        return { 'error': req.status_code, 'reason': req.reason, 'url': req.url }
+    else:
+        return req.json()
+
+
+def post(endpoint: str, dct: dict = None):
+    req = requests.post(f'{sd_url}{endpoint}', json = dct, timeout=300, verify=False, auth=auth())
+    if req.status_code != 200:
+        return { 'error': req.status_code, 'reason': req.reason, 'url': req.url }
+    else:
+        return req.json()
+
+
+def info(args): # pylint: disable=redefined-outer-name
+    t0 = time.time()
+    with open(args.input, 'rb') as f:
+        image = base64.b64encode(f.read()).decode()
+    if args.mask:
+        with open(args.mask, 'rb') as f:
+            mask = base64.b64encode(f.read()).decode()
+    else:
+        mask = None
+    options = get('/sdapi/v1/masking')
+    log.info(f'options: {options}')
+    req = {
+        'image': image,
+        'mask': mask,
+        'type': 'Composite',
+        'params': { 'auto_mask': 'Grayscale' if mask is None else None },
+    }
+    data = post('/sdapi/v1/mask', req)
+    t1 = time.time()
+    if 'mask' in data:
+        b64 = data['mask'].split(',',1)[0]
+        image = Image.open(io.BytesIO(base64.b64decode(b64)))
+        log.info(f'received image: size={image.size} time={t1-t0:.2f}')
+    else:
+        log.info(f'received: {data} time={t1-t0:.2f}')
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description = 'simple-info')
+    parser.add_argument('--input', required=True, help='input image')
+    parser.add_argument('--mask', required=False, help='input mask')
+    args = parser.parse_args()
+    log.info(f'info: {args}')
+    info(args)

modules/api/api.py

@@ -52,13 +52,15 @@ class Api:
         self.add_api_route("/sdapi/v1/extra-single-image", self.extras_single_image_api, methods=["POST"], response_model=models.ResProcessImage)
         self.add_api_route("/sdapi/v1/extra-batch-images", self.extras_batch_images_api, methods=["POST"], response_model=models.ResProcessBatch)
         self.add_api_route("/sdapi/v1/preprocess", control_api.post_preprocess, methods=["POST"])
-        self.add_api_route("/sdapi/v1/preprocessors", control_api.get_preprocess, methods=["GET"])
+        self.add_api_route("/sdapi/v1/mask", control_api.post_mask, methods=["POST"])
 
         # api dealing with optional scripts
         self.add_api_route("/sdapi/v1/scripts", script.get_scripts_list, methods=["GET"], response_model=models.ResScripts)
         self.add_api_route("/sdapi/v1/script-info", script.get_script_info, methods=["GET"], response_model=List[models.ItemScript])
 
         # enumerator api
+        self.add_api_route("/sdapi/v1/preprocessors", control_api.get_preprocess, methods=["GET"])
+        self.add_api_route("/sdapi/v1/masking", control_api.get_mask, methods=["GET"])
         self.add_api_route("/sdapi/v1/interrogate", endpoints.get_interrogate, methods=["GET"], response_model=List[str])
         self.add_api_route("/sdapi/v1/samplers", endpoints.get_samplers, methods=["GET"], response_model=List[models.ItemSampler])
         self.add_api_route("/sdapi/v1/upscalers", endpoints.get_upscalers, methods=["GET"], response_model=List[models.ItemUpscaler])

modules/control/api.py

@@ -1,6 +1,5 @@
 from typing import Optional
 from pydantic import BaseModel, Field # pylint: disable=no-name-in-module
-from fastapi.exceptions import HTTPException
 from fastapi.responses import JSONResponse
 from modules.api.helpers import decode_base64_to_image, encode_pil_to_base64
 
@@ -9,14 +8,14 @@ processor = None # cached instance of processor
 
 
 class ReqPreprocess(BaseModel):
-    image: str = Field(default=None, title="Image", description="The base64 encoded image")
-    model: str = Field(default=None, title="Model", description="The model to use for preprocessing")
+    image: str = Field(title="Image", description="The base64 encoded image")
+    model: str = Field(title="Model", description="The model to use for preprocessing")
     params: Optional[dict] = Field(default={}, title="Settings", description="Preprocessor settings")
 
 
 class ResPreprocess(BaseModel):
-    model: str = Field(default=None, title="Model", description="The processor model used")
-    image: str = Field(default=None, title="Image", description="The processed image in base64 format")
+    model: str = Field(default='', title="Model", description="The processor model used")
+    image: str = Field(default='', title="Image", description="The processed image in base64 format")
 
 
 def get_preprocess():
@@ -32,13 +31,59 @@ def post_preprocess(req: ReqPreprocess):
     from modules.control import processors
     models = list(processors.config)
     if req.model not in models:
-        raise HTTPException(status_code=404, detail=f"Processor model not found: id={req.model}")
+        return JSONResponse(status_code=400, content={"error": f"Processor model not found: id={req.model}"})
     image = decode_base64_to_image(req.image)
     if processor is None or processor.processor_id != req.model:
         processor = processors.Processor(req.model)
     for k, v in req.params.items():
         if k not in processors.config[processor.processor_id]['params']:
-            raise HTTPException(status_code=400, detail=f"Processor invalid parameter: id={req.model} {k}={v}")
+            return JSONResponse(status_code=400, content={"error": f"Processor invalid parameter: id={req.model} {k}={v}"})
     processed = processor(image, local_config=req.params)
     image = encode_pil_to_base64(processed)
     return ResPreprocess(model=processor.processor_id, image=image)
+
+
+class ReqMask(BaseModel):
+    image: str = Field(title="Image", description="The base64 encoded image")
+    type: str = Field(title="Mask type", description="Type of masking image to return")
+    mask: Optional[str] = Field(title="Mask", description="If optional maks image is not provided auto-masking will be performed")
+    model: Optional[str] = Field(title="Model", description="The model to use for preprocessing")
+    params: Optional[dict] = Field(default={}, title="Settings", description="Preprocessor settings")
+
+
+class ResMask(BaseModel):
+    mask: str = Field(default='', title="Image", description="The processed image in base64 format")
+
+
+def get_mask():
+    from modules import masking
+    res = {
+        'models': list(masking.MODELS),
+        'colormaps': masking.COLORMAP,
+        'params': vars(masking.opts),
+        'types': masking.TYPES,
+    }
+    return JSONResponse(res)
+
+
+def post_mask(req: ReqMask):
+    from modules import masking
+    if req.model:
+        if req.model not in masking.MODELS:
+            return JSONResponse(status_code=400, content={"error": f"Mask model not found: id={req.model}"})
+        else:
+            masking.init_model(req.model)
+    if req.type not in masking.TYPES:
+        return JSONResponse(status_code=400, content={"error": f"Mask type not found: id={req.type}"})
+    image = decode_base64_to_image(req.image)
+    mask = decode_base64_to_image(req.mask) if req.mask else None
+    for k, v in req.params.items():
+        if not hasattr(masking.opts, k):
+            return JSONResponse(status_code=400, content={"error": f"Mask invalid parameter: {k}={v}"})
+        else:
+            setattr(masking.opts, k, v)
+    processed = masking.run_mask(input_image=image, input_mask=mask, return_type=req.type)
+    if processed is None:
+        return JSONResponse(status_code=400, content={"error": "Mask is none"})
+    image = encode_pil_to_base64(processed)
+    return ResMask(mask=image)

modules/masking.py

@@ -141,6 +141,7 @@ MODELS = {
     # "isnet-anime",
 }
 COLORMAP = ['autumn', 'bone', 'jet', 'winter', 'rainbow', 'ocean', 'summer', 'spring', 'cool', 'hsv', 'pink', 'hot', 'parula', 'magma', 'inferno', 'plasma', 'viridis', 'cividis', 'twilight', 'shifted', 'turbo', 'deepgreen']
+TYPES = ['None', 'Opaque', 'Binary', 'Masked', 'Grayscale', 'Color', 'Composite']
 cache_dir = 'models/control/segment'
 generator: MaskGenerationPipeline = None
 busy = False