add footnote

app.py

@@ -22,7 +22,7 @@ cfg = setup(args)
 
 multi_classes = False
 
-device = "cpu"  # "cuda" if torch.cuda.is_available() else "cpu"
+device = "cuda" if torch.cuda.is_available() else "cpu"
 Ours, preprocess = models.load("CS-ViT-B/16", device=device, cfg=cfg, train_bool=False)
 state_dict = torch.load("sketch_seg_best_miou.pth", map_location=device)
 
@@ -37,6 +37,9 @@ print("Model loaded successfully")
 
 
 def run(sketch, caption, threshold, seed):
+    # select a random seed between 1 and 10 for the color
+    color_seed = np.random.randint(0, 9)
+    
     # set the condidate classes here
     caption = caption.replace('\n',' ')
     translator = str.maketrans('', '', string.punctuation)
@@ -49,7 +52,7 @@ def run(sketch, caption, threshold, seed):
     # print(classes)
     
     colors = plt.get_cmap("Set1").colors
-    classes_colors = colors[:len(classes)]
+    classes_colors = colors[color_seed:len(classes)+color_seed]
 
     sketch2 = sketch['composite']   
 
@@ -127,9 +130,11 @@ def run(sketch, caption, threshold, seed):
     plt.imshow(rgb_image)
     plt.axis('off')
     plt.tight_layout()
+    # plt.savefig(f'poster_vis/{classes[0]}.png', bbox_inches='tight', pad_inches=0)
     plt.savefig('output.png', bbox_inches='tight', pad_inches=0)
     plt.close()
     
+    # rgb_image = Image.open(f'poster_vis/{classes[0]}.png')    
     rgb_image = Image.open('output.png')
 
     return rgb_image
@@ -212,9 +217,7 @@ async () => {
         element.textContent = firstWord;
     });
     
-    document.querySelector('input[type=number]').disabled = true;
-    
-    
+    document.querySelector('input[type=number]').disabled = true;   
 }
 """
 
@@ -401,13 +404,16 @@ with gr.Blocks(js=scripts, css=css, theme='gstaff/xkcd') as demo:
     with gr.Row():
         with gr.Column():
             # in_image = gr.Image( label="Sketch", type="pil", sources="upload" , height=512 )
-            in_canvas_image = gr.Sketchpad(  brush=gr.Brush(colors=["#000000"], color_mode="fixed" , default_size=2), 
+            in_canvas_image = gr.Sketchpad(
+                # value=Image.new('RGB', (512, 512), color=(255, 255, 255)),
+                brush=gr.Brush(colors=["#000000"], color_mode="fixed" , default_size=2), 
                 image_mode="RGBA",elem_classes=["white-bg", "gray-border" , "border-radius" ,"own-shadow" ] ,  
-                label="Sketch" , canvas_size=(512,512) , sources=['upload'], 
-                interactive=True , layers= False, transforms=[] )
+                label="Sketch" , canvas_size=(512,512) ,sources=['upload'], 
+                interactive=True , layers= False, transforms=[] 
+                )
             query_selector = 'button[aria-label="Upload button"]'
+            
             # with gr.Row():
-
                 # segment_btn.click(fn=run, inputs=[in_image, in_textbox, in_slider], outputs=[out_image])
             upload_draw_btn = gr.HTML(f"""
                 <div id="upload_draw_group" class="svelte-15lo0d8 stretch">
@@ -441,7 +447,7 @@ with gr.Blocks(js=scripts, css=css, theme='gstaff/xkcd') as demo:
                                     value=0.5 , interactive=True,  step=0.05, minimum=0, maximum=1)
 
     with gr.Row():
-        segment_btn = gr.Button( 'Segment it !' , elem_classes=["white-bg", "gray-border" , "border-radius" ,"own-shadow" , 'bold' , 'mt-mb-1' ] , size="sm")
+        segment_btn = gr.Button( 'Segment it* !' , elem_classes=["white-bg", "gray-border" , "border-radius" ,"own-shadow" , 'bold' , 'mt-mb-1' ] , size="sm")
         segment_btn.click(fn=run, inputs=[in_canvas_image , in_textbox , in_slider  ], outputs=[out_image])
     gallery_label = gr.HTML("<h3 class='black-text'> <span class='black-text underline'>Gallery:</span> <span style='color: grey;'>you can click on any of the example sketches below to start segmenting them (or even drawing over them)</span> </div>")
 
@@ -498,5 +504,7 @@ with gr.Blocks(js=scripts, css=css, theme='gstaff/xkcd') as demo:
         fn=run,
         # cache_examples=True,
     )
-
+    
+    gr.HTML("<h5 class='black-text' style='text-align: left;'>*For optimal performance, use a commercial Nvidia RTX 3090 (this demo runs on a basic 2 vCPU).</h5>")
+    gr.HTML("<h5 class='black-text' style='text-align: left;'>*We compare the entire caption to the scene sketch and threshold most similar pixels, without extracting individual classes.</h5>")
 demo.launch(share=False)