api changes + notes

typo

updating depth api

	modified:   scripts/depthmap_api.py

Author: graemeniedermayer

.gitignore

@@ -1 +1,3 @@
-__pycache__/
+__pycache__/
+models/
+ouputs/

scripts/depthmap_api.py

@@ -6,17 +6,20 @@
 from fastapi import FastAPI, Body
 from fastapi.exceptions import HTTPException
 from PIL import Image
+from itertools import tee
 
 import gradio as gr
 
 from modules.api.models import List, Dict
 from modules.api import api
+from typing_extensions import Annotated
+from pydantic import BaseModel
 
-from src.core import core_generation_funnel
+from src.core import core_generation_funnel, CoreGenerationFunnelInp
 from src.misc import SCRIPT_VERSION
 from src import backbone
 from src.common_constants import GenerationOptions as go
-
+from src.api_constants import Api_Defaults, Api_Forced, Api_options
 
 def encode_to_base64(image):
     if type(image) is str:
@@ -28,48 +31,77 @@ def encode_to_base64(image):
     else:
         return ""
 
-
 def encode_np_to_base64(image):
     pil = Image.fromarray(image)
     return api.encode_pil_to_base64(pil)
 
-
 def to_base64_PIL(encoding: str):
     return Image.fromarray(np.array(api.decode_base64_to_image(encoding)).astype('uint8'))
 
 
+def api_gen(depth_input_images, options):
+    
+    default_options = CoreGenerationFunnelInp({Api_Defaults}).values
+
+    #TODO try-catch type errors here
+    for key, value in options.items():
+        default_options[key] = value
+
+    for key, value in Api_Forced.items():
+        default_options[key.lower()] = value
+
+    if len(depth_input_images) == 0:
+        raise HTTPException(status_code=422, detail="No images supplied")
+        
+    print(f"Processing {str(len(depth_input_images))} images through the API")
+
+    pil_images = []
+    for input_image in depth_input_images:
+        pil_images.append(to_base64_PIL(input_image))
+    outpath = backbone.get_outpath()
+    gen_obj = core_generation_funnel(outpath, pil_images, None, None, options)
+    return gen_obj
+
 def depth_api(_: gr.Blocks, app: FastAPI):
     @app.get("/depth/version")
     async def version():
         return {"version": SCRIPT_VERSION}
 
     @app.get("/depth/get_options")
     async def get_options():
-        return {"options": sorted([x.name.lower() for x in go])}
+        return {
+            "api_options": Api_options,
+            "gen_options": [x.name.lower() for x in go]
+        }
 
-    # TODO: some potential inputs not supported (like custom depthmaps)
     @app.post("/depth/generate")
     async def process(
         depth_input_images: List[str] = Body([], title='Input Images'),
-        options: Dict[str, object] = Body("options", title='Generation options'),
+        api_options: Dict[str, object] = Body({'outputs': ["depth"]}, title='Api options', options= Api_options),
+        gen_options: Dict[str, object] = Body({}, title='Generation options', options= [x.name.lower() for x in go])
     ):
-        # TODO: restrict mesh options
-
-        if len(depth_input_images) == 0:
-            raise HTTPException(status_code=422, detail="No images supplied")
-        print(f"Processing {str(len(depth_input_images))} images trough the API")
-
-        pil_images = []
-        for input_image in depth_input_images:
-            pil_images.append(to_base64_PIL(input_image))
-        outpath = backbone.get_outpath()
-        gen_obj = core_generation_funnel(outpath, pil_images, None, None, options)
-
-        results_based = []
-        for count, type, result in gen_obj:
-            if type not in ['simple_mesh', 'inpainted_mesh']:
-                results_based += [encode_to_base64(result)]
-        return {"images": results_based, "info": "Success"}
+        gen_obj = api_gen(depth_input_images, gen_options)
+        #NOTE Work around yield. (Might not be necessary, not sure if yield caches)
+        _, gen_obj = tee (gen_obj)
+
+        if len(api_options["outputs"])>1:
+            results_based = {}
+
+            for type in api_options["outputs"]:
+                result_per_type = []
+
+                for count, img_type, result in gen_obj:
+                    if img_type == type:
+                        result_per_type += result
+                
+                if len(result_per_type)==0:
+                    results_based[type] = "Check options. no img-type of " + str(type) + " where generated"
+                else:
+                    results_based[type] = map(encode_to_base64, result_per_type)
+
+            return {"images": results_based, "info": "Success"}
+        else:
+            return {"images": {}, "info": "api_options.output is empty"}
 
 
 try:

src/api_constants.py

@@ -0,0 +1,17 @@
+# TODO Maybe these have a better home
+Api_options = {
+    'outputs': ["depth"], # list of outputs to send in response. examples ["depth", "normalmap", 'heatmap', "normal", 'background_removed'] etc
+    #'conversions': "", #TODO implement. it's a good idea to give some options serverside for because often that's challenging in js/clientside 
+    'save':"" #TODO implement. To save on local machine. Can be very helpful for debugging.
+}
+
+# TODO: These two are intended to be temporary
+Api_Defaults={
+    "BOOST": False,
+    "NET_SIZE_MATCH": True
+}
+#These are enforced after user inputs
+Api_Forced={
+    "GEN_SIMPLE_MESH": False,
+    "GEN_INPAINTED_MESH": False
+}