6

app.py

@@ -1,44 +1,66 @@
+import torch.nn as nn
+import torch
+
+from model.DCAMA import DCAMA
+from common.logger import Logger, AverageMeter
+from common.vis import Visualizer
+from common.evaluation import Evaluator
+from common.config import parse_opts
+from common import utils
+from data.dataset import FSSDataset
+import cv2
+import numpy as np
+import os
+
 import gradio as gr
 import time
 
 
-def sleep(im):
-    time.sleep(5)
-    return [im["background"], im["layers"][0], im["layers"][1], im["composite"]]
+def inference_mask1(
+            query_img,
+            *prompt,
+              ):
+    support_masks = []
+    support_imgs = []
+    for i in range(len(prompt)):
+        mask = np.stack(prompt[i]['layers'], axis=0).any(0).any(-1)
+        support_masks.append(mask)
+        support_imgs.append(prompt[i]['background'])
+    model = DCAMA('resnet50', 'resnet50_a1h-35c100f8.pth', True)
+    model.eval()
+    model.cuda()
+    state_dict = torch.load('model_45.pt')
+    if 'state_dict' in state_dict.keys():
+        state_dict = state_dict['state_dict']
+    state_dict2 = {}
+    for k, v in state_dict.items():
+        if 'scorer' not in k:
+            state_dict2[k] = v
+    state_dict = state_dict2
+
+    for k1, k2 in zip(list(state_dict.keys()), params.keys()):
+        state_dict[k2] = state_dict.pop(k1)
+
+    try:
+        model.load_state_dict(state_dict, strict=True)
+    except:
+        for k in params.keys():
+            if k not in state_dict.keys():
+                state_dict[k] = params[k]
+        model.load_state_dict(state_dict, strict=True)
 
+    pred_mask, simi, simi_map = model.module.predict_mask_nshot(batch, nshot=nshot)
 
-support_im_masks = [None for _ in range(1000)]
-def predict(im):
-  for i in range(len(support_im_masks)):
-    if support_im_masks[i] is None:
-      break
-    support_im_mask = support_im_masks[i]
     import ipdb;ipdb.set_trace()
-    pass
-
-with gr.Blocks() as demo:
-    b = gr.Button("Add Textbox")
-    b2 = gr.Button("Generate Masks")
-    b2.click(predict)
-  
-    num = gr.State(0)
-    
-    b.click(lambda x:x+1, num, num)
-
-    with gr.Row():
-      query_im = gr.Image(label='query image')
-      
-      @gr.render(inputs=num)
-      def show_support_imgs(n):
-          with gr.Column():
-              for i in range(n):
-                  support_im = gr.ImageEditor(
-                      label="support image {}".format(i),
-                      type="numpy",
-                      crop_size="1:1",
-                  )
-                  support_im_masks[i] = support_im
-
-
-if __name__ == "__main__":
-    demo.launch(debug=True)
+
+
+inputs = [gr.Image(label='query')]
+for i in range(5):
+    inputs.append(gr.ImageMask(label='prompt{}'.format(i)))
+demo_mask = gr.Interface(fn=inference_mask1,
+                   inputs=inputs,
+    outputs=[gr.Image(label="output")],
+)
+
+demo = gr.TabbedInterface([demo_mask], ['demo'])
+demo.launch()