Update app.py

app.py

@@ -141,6 +141,8 @@ class ImageMask(gr.components.Image):
            
             decode_image = processing_utils.decode_base64_to_image(x)
             width, height = decode_image.size
+            img = np.asarray(decode_image)
+            return {'image':img, 'mask':binarize_2(img)}
             
             mask = np.zeros((height, width, 4), dtype=np.uint8)
             
@@ -148,7 +150,9 @@ class ImageMask(gr.components.Image):
             mask = self.postprocess(mask)
             x = {'image': x, 'mask': mask}
             print('vao preprocess-------------------------')
-        hh = super().preprocess(x)
+        if (hh['image'].min()!=255) and (hh['mask'][:,:,:3].max()==0):
+            
+            hh['mask'] = binarize_2(hh['image'])
         
         return hh
 
@@ -446,8 +450,9 @@ def draw(task, input, grounding_texts, new_image_trigger, state, generate_parsed
     
     return [box_image, new_image_trigger, image_scale, state, generate_parsed]
 
-def change_state(bboxes,layout, state, instruction):
-    
+def change_state(bboxes,layout, state, instruction, trigger_stage, boxes):
+    if trigger_stage ==0 :
+        return [boxes, state, 0]
     # mask = 
     state['boxes'] = []
     state['masks'] = []
@@ -470,14 +475,13 @@ def change_state(bboxes,layout, state, instruction):
         grounding_texts += [f'Obj. {bid+1}' for bid in range(len(grounding_texts), len(result))]
 
     box_image = draw_box(result, grounding_texts)
-    mask= np.asarray(box_image)
     
-    mask = binarize_2(mask)
+    mask = binarize_2(layout['image'])
     state['masks'].append(mask.copy())
     # print('done change state', state)
     print('done change state')
     # import pdb; pdb.set_trace()
-    return [box_image,state]
+    return [box_image,state, trigger_stage]
 
 def example_click(name, grounding_instruction, instruction, bboxes,generate_parsed,  trigger_parsed):
     
@@ -497,15 +501,16 @@ def example_click(name, grounding_instruction, instruction, bboxes,generate_pars
     print('done the example click')
     return [box_image, trigger_parsed]
 
-def clear(task, sketch_pad_trigger, batch_size, state, switch_task=False):
+def clear(task, sketch_pad_trigger, batch_size, state,trigger_stage, switch_task=False):
     
     sketch_pad_trigger = sketch_pad_trigger + 1
+    trigger_stage = 0
     blank_samples = batch_size % 2 if batch_size > 1 else 0
     out_images = [gr.Image.update(value=None, visible=True) for i in range(batch_size)] \
                     + [gr.Image.update(value=None, visible=True) for _ in range(blank_samples)] \
                     + [gr.Image.update(value=None, visible=False) for _ in range(4 - batch_size - blank_samples)]
     state = {}
-    return [None, sketch_pad_trigger, None, 1.0] + out_images + [state]
+    return [None, sketch_pad_trigger, None, 1.0] + out_images + [state] + [trigger_stage]
 
 css = """
 #img2img_image, #img2img_image > .fixed-height, #img2img_image > .fixed-height > div, #img2img_image > .fixed-height > div > img
@@ -655,8 +660,8 @@ with Blocks(
         )
         trigger_stage.change(
             change_state,
-            inputs=[text_box,sketch_pad, state, grounding_instruction],
-            outputs=[out_imagebox,state],
+            inputs=[text_box,sketch_pad, state, grounding_instruction, trigger_stage,out_imagebox],
+            outputs=[out_imagebox,state,trigger_stage],
             queue=True
         )
         grounding_instruction.change(
@@ -667,8 +672,8 @@ with Blocks(
         )
         clear_btn.click(
             clear,
-            inputs=[task, sketch_pad_trigger, batch_size, state],
-            outputs=[sketch_pad, sketch_pad_trigger, out_imagebox, image_scale, out_gen_1, out_gen_2, out_gen_3, out_gen_4, state],
+            inputs=[task, sketch_pad_trigger, batch_size,trigger_stage, state],
+            outputs=[sketch_pad, sketch_pad_trigger, out_imagebox, image_scale, out_gen_1, out_gen_2, out_gen_3, out_gen_4, state, trigger_stage],
             queue=False)
     
         sketch_pad_trigger.change(