Added n class support

app.py

@@ -25,7 +25,7 @@ args.cx = '06d75168141bc47f1'
 
 
 # model
-device = 'cpu' #torch.device("cuda" if torch.cuda.is_available() else "cpu")
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = get_model(args)
 model.to(device)
 checkpoint = torch.hub.load_state_dict_from_url(args.resume, map_location='cpu')
@@ -57,25 +57,29 @@ def denormalize(x, mean, std):
 
 
 # Gradio UI
-def inference(query, class1_name="class1", support_imgs=None, class2_name="class2", support_imgs2=None):
+def inference(query, *support_text_box_and_files):
     '''
     query: PIL image
-    labels: list of class names
+    class_names: list of class names
     '''
 
 
-    #first, open the images
-    support_imgs = [Image.open(img) for img in support_imgs]
-    support_imgs2 = [Image.open(img) for img in support_imgs2]
 
+    labels = support_text_box_and_files[0::2]
+    support_images = support_text_box_and_files[1::2]
+
+    print(f"Support images: {support_images}")
+
+    #first, open the images
+    support_images = [[Image.open(img) for img in imgs] for imgs in support_images if imgs != None]
 
-    labels = [class1_name, class2_name]
-    
     supp_x = []
     supp_y = []
-
-    for i, (class_name, support_img) in enumerate(zip([class1_name, class2_name], [support_imgs, support_imgs2])):
-        for img in support_img:
+    for i, support_imgs in enumerate(support_images):
+    #for i, (class_name, support_imgs) in enumerate(zip(class_names, support_images)):
+        if len(support_imgs) == 0:
+            continue 
+        for img in support_imgs:
             x_im = preprocess(img)
             supp_x.append(x_im)
             supp_y.append(i)
@@ -93,11 +97,14 @@ def inference(query, class1_name="class1", support_imgs=None, class2_name="class
 
 
     with torch.cuda.amp.autocast(True):
+        start_time = time.time()
         output = model(supp_x, supp_y, query) # (1, 1, n_labels)
+        exec_time = time.time() - start_time
 
     probs = output.softmax(dim=-1).detach().cpu().numpy()
     
-    return {k: float(v) for k, v in zip(labels, probs[0, 0])}
+
+    return {k: float(v) for k, v in zip(labels, probs[0, 0])}, exec_time
 
 
 # DEBUG
@@ -109,25 +116,84 @@ def inference(query, class1_name="class1", support_imgs=None, class2_name="class
 #print(output)
 
 
-title = "P>M>F few-shot learning pipeline"
+title = "# P>M>F few-shot learning pipeline"
 description = "Short description: We take a ViT-small backbone, which is pre-trained with DINO, and meta-trained on Meta-Dataset; for few-shot classification, we use a ProtoNet classifier. The demo can be viewed as zero-shot since the support set is built by searching images from Google. Note that you may need to play with GIS parameters to get good support examples. Besides, GIS is not very stable as search requests may fail for many reasons (e.g., number of requests reaches the limit of the day). This code is heavely inspired from the original HF space <a href='https://huggingface.co/spaces/hushell/pmf_with_gis' target='_blank'>here</a>"
 article = "<p style='text-align: center'><a href='http://arxiv.org/abs/2204.07305' target='_blank'>Arxiv</a></p>"
-gr.Interface(fn=inference,
-             inputs=[
-                 gr.Image(label="Image to classify", type="pil"),
-                 gr.Textbox(lines=1, label="First class name :", placeholder="Enter first class name",),
-                 gr.File(label="First class example images", file_types=["image"], file_count="multiple"),
-                 gr.Textbox(lines=1, label="Second class name :", placeholder="Enter second class name",),
-                 gr.File(label="Second class example iamges", file_types=["image"], file_count="multiple"),
-             ],
-             theme="grass",
-             outputs=[
-                 gr.Label(label="Predicted class probabilities"),
-             ],
-             title=title,
-             description=description,
-             article=article,
-             examples=[
-                 ["./example_images/2007_000033.jpg", "plane", ["./example_images/2007_000738.jpg", "./example_images/2007_000256.jpg"], "cat", ["./example_images/2007_000528.jpg", "./example_images/2007_000549.jpg"]]
-             ]
-            ).launch(debug=True)
+
+min_classes = 2
+max_classes = 10
+
+
+def variable_outputs(k):
+    k = int(k)
+    inputs = []
+    for _ in range(k):
+        inputs.append(gr.Textbox(visible=True))
+        inputs.append(gr.File(visible=True))
+    
+    for _ in range(max_classes-k):
+        inputs.append(gr.Textbox(visible=False))
+        inputs.append(gr.File(visible=False))
+
+    return inputs
+
+with gr.Blocks() as demo:
+
+    
+    with gr.Row():
+        gr.Markdown(title)
+    with gr.Row():
+        gr.Markdown(description)
+    with gr.Row():
+        gr.Markdown(article)
+    with gr.Row():
+        with gr.Column():
+
+
+            query = gr.Image(label="Image to classify", type="pil")
+            num_classes_slider = gr.Slider(minimum=min_classes, maximum=10, value=2, label="Number of classes", step=1)
+            
+            #set_number_classes_btn = gr.Button("Set number of classes")
+
+            textboxes_and_files = []
+            for i in range(max_classes):
+                is_visible = (i < 2) 
+                t = gr.Textbox(label=f"Class {i+1} name", placeholder=f"Enter class {i+1} name", visible=is_visible)
+                textboxes_and_files.append(t)
+                f = gr.File(label=f"Support image for class {i+1}", type="filepath", visible=is_visible, file_count="multiple")
+                textboxes_and_files.append(f)
+            
+
+
+
+            num_classes_slider.change(variable_outputs, inputs=[num_classes_slider], outputs=textboxes_and_files)
+
+
+            run_button = gr.Button("Run Inference")
+
+        with gr.Column():
+            output = gr.Label(label="Predicted class probabilities")
+            exec_time = gr.Textbox(label="Execution time (s)")
+
+
+#    def run_inference(query, *example_inputs):
+#        
+#        print("len(example_inputs) : ")
+#        print(len(example_inputs))
+#
+#        class_names = [example_inputs[i].value for i in range(0, len(example_inputs), 2)]
+#        support_images = [example_inputs[i].value for i in range(1, len(example_inputs), 2)]
+#        return inference(query, class_names, support_images)
+
+    run_button.click(
+        fn=inference,
+        inputs=[query] + textboxes_and_files,
+        outputs=[output, exec_time]
+    )
+    
+
+    # this does nothing it seems
+    demo.examples = [
+        ["./example_images/2007_000033.jpg", "plane", ["./example_images/2007_000738.jpg", "./example_images/2007_000256.jpg"], "cat", ["./example_images/2007_000528.jpg", "./example_images/2007_000549.jpg"]]
+    ]
+    demo.launch()