add more options for GIS

app.py

@@ -4,7 +4,7 @@ import time
 import random
 import torch
 import torchvision.transforms as transforms
-import gradio as gr
+#import gradio as gr
 import matplotlib.pyplot as plt
 
 from models import get_model
@@ -83,7 +83,9 @@ _search_params = {
 
 
 # Gradio UI
-def inference(query, labels, n_supp=10):
+def inference(query, labels, n_supp=10,
+              file_type='png', rights='cc_publicdomain',
+              image_type='photo', color_type='color'):
     '''
     query: PIL image
     labels: list of class names
@@ -91,6 +93,12 @@ def inference(query, labels, n_supp=10):
     labels = labels.split(',')
     n_supp = int(n_supp)
 
+    _search_params['num'] = n_supp
+    _search_params['fileType'] = file_type
+    _search_params['rights'] = rights
+    _search_params['imgType'] = image_type
+    _search_params['imgColorType'] = color_type
+
     fig, axs = plt.subplots(len(labels), n_supp, figsize=(n_supp*4, len(labels)*4))
 
     with torch.no_grad():
@@ -104,9 +112,8 @@ def inference(query, labels, n_supp=10):
         for idx, y in enumerate(labels):
             gis = GoogleImagesSearch(args.api_key, args.cx)
             _search_params['q'] = y
-            _search_params['num'] = n_supp
             gis.search(search_params=_search_params, custom_image_name='my_image')
-            gis._custom_image_name = 'my_image'
+            gis._custom_image_name = 'my_image' # fix: image name sometimes too long
 
             for j, x in enumerate(gis.results()):
                 x.download('./')
@@ -135,9 +142,10 @@ def inference(query, labels, n_supp=10):
 
 
 # DEBUG
-#query = Image.open('../labrador-puppy.jpg')
+##query = Image.open('../labrador-puppy.jpg')
+#query = Image.open('/Users/hushell/Documents/Dan_tr.png')
 ##labels = 'dog, cat'
-#labels = 'girl, boy'
+#labels = 'girl, sussie'
 #output = inference(query, labels, n_supp=2)
 #print(output)
 
@@ -146,7 +154,11 @@ gr.Interface(fn=inference,
              inputs=[
                  gr.inputs.Image(label="Image to classify", type="pil"),
                  gr.inputs.Textbox(lines=1, label="Class hypotheses:", placeholder="Enter class names separated by ','",),
-                 gr.inputs.Slider(minimum=2, maximum=10, step=1, label="Number of support examples from Google")
+                 gr.inputs.Slider(minimum=2, maximum=10, step=1, label="GIS: Number of support examples per class"),
+                 gr.inputs.Dropdown(['png', 'jpg'], default='png', label='GIS: Image file type'),
+                 gr.inputs.Dropdown(['cc_publicdomain', 'cc_attribute', 'cc_sharealike', 'cc_noncommercial', 'cc_nonderived'], default='cc_publicdomain', label='GIS: Copy rights'),
+                 gr.inputs.Dropdown(['clipart', 'face', 'lineart', 'stock', 'photo', 'animated', 'imgTypeUndefined'], default='photo', label='GIS: Image type'),
+                 gr.inputs.Dropdown(['color', 'gray', 'mono', 'trans', 'imgColorTypeUndefined'], default='color', label='GIS: Image color type'),
              ],
              theme="grass",
              outputs=[