fix curse issue

app.py

@@ -4,7 +4,6 @@ import time
 import random
 import torch
 import torchvision.transforms as transforms
-#import requests
 import gradio as gr
 import matplotlib.pyplot as plt
 
@@ -12,16 +11,13 @@ from models import get_model
 from dotmap import DotMap
 from PIL import Image
 
-
-os.environ['TERM'] = 'linux'
-os.environ['TERMINFO'] = '/etc/terminfo'
-
+#os.environ['TERM'] = 'linux'
+#os.environ['TERMINFO'] = '/etc/terminfo'
 
 # args
 args = DotMap()
 args.deploy = 'vanilla'
 args.arch = 'dino_small_patch16'
-#args.resume = '/fast_scratch/hushell/fluidstack/FS125_few-shot-transformer/outputs/dinosmall_1e-4/best_converted.pth'
 args.resume = 'https://huggingface.co/hushell/pmf_dinosmall_lr1e-4/resolve/main/best_converted.pth'
 args.api_key = 'AIzaSyAFkOGnXhy-2ZB0imDvNNqf2rHb98vR_qY'
 args.cx = '06d75168141bc47f1'
@@ -31,7 +27,6 @@ args.cx = '06d75168141bc47f1'
 device = 'cpu' #torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model = get_model(args)
 model.to(device)
-#checkpoint = torch.load(args.resume, map_location='cpu')
 checkpoint = torch.hub.load_state_dict_from_url(args.resume, map_location='cpu')
 model.load_state_dict(checkpoint['model'], strict=True)
 
@@ -63,6 +58,12 @@ def denormalize(x, mean, std):
 # Google image search
 from google_images_search import GoogleImagesSearch
 
+class MyGIS(GoogleImagesSearch):
+    def __enter__(self):
+        return self
+    def __exit__(self, exc_type, exc_val, exc_tb):
+        return
+
 # define search params
 # option for commonly used search param are shown below for easy reference.
 # For param marked with '##':
@@ -90,7 +91,6 @@ def inference(query, labels, n_supp=10):
     labels = labels.split(',')
     n_supp = int(n_supp)
 
-    #print(f'#rows={len(labels)}, #cols={n_supp}')
     fig, axs = plt.subplots(len(labels), n_supp, figsize=(n_supp*4, len(labels)*4))
 
     with torch.no_grad():
@@ -102,26 +102,24 @@ def inference(query, labels, n_supp=10):
 
         # search support images
         for idx, y in enumerate(labels):
-            with GoogleImagesSearch(args.api_key, args.cx) as gis:
-                _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'
-
-                for j, x in enumerate(gis.results()):
-                    #url = x.url
-                    #x_im = Image.open(requests.get(url, stream=True).raw)
-                    x.download('./')
-                    x_im = Image.open(x.path)
-
-                    # vis
-                    axs[idx, j].imshow(x_im)
-                    axs[idx, j].set_title(f'{y}{j}')
-                    axs[idx, j].axis('off')
-
-                    x_im = preprocess(x_im) # (3, H, W)
-                    supp_x.append(x_im)
-                    supp_y.append(idx)
+            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'
+
+            for j, x in enumerate(gis.results()):
+                x.download('./')
+                x_im = Image.open(x.path)
+
+                # vis
+                axs[idx, j].imshow(x_im)
+                axs[idx, j].set_title(f'{y}{j}')
+                axs[idx, j].axis('off')
+
+                x_im = preprocess(x_im) # (3, H, W)
+                supp_x.append(x_im)
+                supp_y.append(idx)
 
         print('Searching for support images is done.')
 
@@ -148,7 +146,6 @@ 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.Number(default=1, label="Number of support examples from Google")
                  gr.inputs.Slider(minimum=2, maximum=10, step=1, label="Number of support examples from Google")
              ],
              theme="grass",