commit

app.py

@@ -4,6 +4,7 @@ from gradio_client import Client, handle_file
 from pathlib import Path
 from gradio.utils import get_cache_folder
 
+from PIL import Image
 
 class Examples(gr.helpers.Examples):
     def __init__(self, *args, cached_folder=None, **kwargs):
@@ -21,37 +22,104 @@ client = Client("Canyu/Diception",
                 hf_token=HF_TOKEN)
 
 
-def process_image_check(path_input):
+map_prompt = {
+    'depth': '[[image2depth]]',
+    'normal': '[[image2normal]]',
+    'pose': '[[image2pose]]',
+    'entity segmentation': '[[image2panoptic coarse]]',
+    'point segmentation': '[[image2segmentation]]',
+    'semantic segmentation': '[[image2semantic]]',
+}
+
+def download_additional_params(model_name, filename="add_params.bin"):
+    # 下载文件并返回文件路径
+    file_path = hf_hub_download(repo_id=model_name, filename=filename, use_auth_token=HF_TOKEN)
+    return file_path
+
+# 加载 additional_params.bin 文件
+def load_additional_params(model_name):
+    # 下载 additional_params.bin
+    params_path = download_additional_params(model_name)
+    
+    # 使用 torch.load() 加载文件内容
+    additional_params = torch.load(params_path, map_location='cpu')
+    
+    # 返回加载的参数内容
+    return additional_params
+
+def process_image_check(path_input, prompt):
     if path_input is None:
         raise gr.Error(
             "Missing image in the left pane: please upload an image first."
         )
+    if len(prompt) == 0:
+        raise gr.Error(
+            "At least 1 prediction type is needed."
+        )
+
+
+
+def process_image_4(image_path, prompt):
+
+    inputs = []
+    for p in prompt:
+        image = Image.open(image_path)
+
+        w, h = image.size
+
+        coor_point = torch.zeros((1,5,2)).to(torch.float32)
+        point_labels = torch.zeros((1,5,1)).to(torch.float32)
 
-def infer_image_matting(matting_image_input):
+        image = image.resize((768, 768), Image.LANCZOS).convert('RGB')
+        to_tensor = transforms.ToTensor()
+        image = (to_tensor(image) - 0.5) * 2
+
+        cur_input = {
+                'input_images': image.unsqueeze(0),
+                'original_size': torch.tensor([[w,h]]),
+                'target_size': torch.tensor([[768, 768]]),
+                'prompt': [p],
+                'coor_point': coor_point,
+                'point_labels': point_labels,
+                'generator': generator
+            }
+        inputs.append(cur_input)
+
+    return inputs
+
+
+def infer_image_matting(image_path, prompt):
+    inputs = process_image_4(image_path, prompt)
+    return None
     return client.predict(
-      prompt=handle_file(matting_image_input),
-      api_name="/infer_image_matting"
+      batch=inputs,
+      api_name="/inf"
     )
 
 def clear_cache():
     return None, None
 
 def run_demo_server():
+    options = ['depth', 'normal', 'entity', 'pose']
     gradio_theme = gr.themes.Default()
     with gr.Blocks(
         theme=gradio_theme,
         title="Matting",
     ) as demo:
         with gr.Row():
-            gr.Markdown("# Matting Demo")
+            gr.Markdown("# Diception Demo")
         with gr.Row():
-            gr.Markdown("### Due to the GPU quota limit, if an error occurs, please wait for 5 minutes before retrying.")
+            gr.Markdown("### All results are generated using the same single model. To facilitate input processing, we separate point-prompted segmentation and semantic segmentation, as they require input points and segmentation targets.")
+        with gr.Row():
+            checkbox_group = gr.CheckboxGroup(choices=options, label="Select options:")
+
         with gr.Row():
             with gr.Column():
                 matting_image_input = gr.Image(
                     label="Input Image",
                     type="filepath",
                 )
+
                 with gr.Row():
                     matting_image_submit_btn = gr.Button(
                         value="Estimate Matting", variant="primary"
@@ -80,7 +148,7 @@ def run_demo_server():
 
         matting_image_submit_btn.click(
             fn=process_image_check,
-            inputs=matting_image_input,
+            inputs=[matting_image_input, checkbox_group],
             outputs=None,
             preprocess=False,
             queue=False,