mirror of
https://github.com/vladmandic/automatic
synced 2026-08-31 09:31:00 +02:00
Vladimir Mandic
90 lines
3.7 KiB
Python
90 lines
3.7 KiB
Python
from typing import Optional
|
|
from pydantic import BaseModel, Field # pylint: disable=no-name-in-module
|
|
from fastapi.responses import JSONResponse
|
|
from modules.api.helpers import decode_base64_to_image, encode_pil_to_base64
|
|
|
|
|
|
processor = None # cached instance of processor
|
|
|
|
|
|
class ReqPreprocess(BaseModel):
|
|
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='', title="Model", description="The processor model used")
|
|
image: str = Field(default='', title="Image", description="The processed image in base64 format")
|
|
|
|
|
|
def get_preprocess():
|
|
from modules.control import processors
|
|
p = {}
|
|
for k, v in processors.config.items():
|
|
p[k] = v.get('params')
|
|
return JSONResponse(p)
|
|
|
|
|
|
def post_preprocess(req: ReqPreprocess):
|
|
global processor # pylint: disable=global-statement
|
|
from modules.control import processors
|
|
models = list(processors.config)
|
|
if req.model not in models:
|
|
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']:
|
|
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)
|