Push localization with load management to master

localization.py

@@ -13,14 +13,10 @@ import jax
 import gc
 
 
-preprocess = transforms.Compose(
-    [
-        transforms.ToTensor(),
-        transforms.Normalize(
-            (0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)
-        ),
-    ]
-)
+preprocess = transforms.Compose([
+    transforms.ToTensor(),
+    transforms.Normalize((0.48145466, 0.4578275, 0.40821073), (0.26862954, 0.26130258, 0.27577711)),
+])
 
 
 def pad_to_square(image, size=224):
@@ -54,19 +50,19 @@ def gen_image_batch(image_url, image_size=224, pixel_size=10):
     masks.append(mask)
 
     for i in range(0, n_pixels):
-        for j in range(i + 1, n_pixels):
+        for j in range(i+1, n_pixels):
             m = mask.copy()
-            m[: min(i * pixel_size, image_size) + 1, :] = 0
-            m[min(j * pixel_size, image_size) + 1 :, :] = 0
+            m[:min(i*pixel_size, image_size) + 1, :] = 0
+            m[min(j*pixel_size, image_size) + 1:, :] = 0
             neg_m = 1 - m
             image_batch.append(image * m + gray * neg_m)
             masks.append(m)
 
-    for i in range(0, n_pixels + 1):
-        for j in range(i + 1, n_pixels + 1):
+    for i in range(0, n_pixels+1):
+        for j in range(i+1, n_pixels+1):
             m = mask.copy()
-            m[:, : min(i * pixel_size + 1, image_size)] = 0
-            m[:, min(j * pixel_size + 1, image_size) :] = 0
+            m[:, :min(i*pixel_size + 1, image_size)] = 0
+            m[:, min(j*pixel_size + 1, image_size):] = 0
             neg_m = 1 - m
             image_batch.append(image * m + gray * neg_m)
             masks.append(m)
@@ -79,9 +75,7 @@ def get_heatmap(image_url, text, pixel_size=10, iterations=3):
     model = get_model()
     image_size = model.config.vision_config.image_size
     text_embedding = text_encoder(text, model, tokenizer)
-    images, masks = gen_image_batch(
-        image_url, image_size=image_size, pixel_size=pixel_size
-    )
+    images, masks = gen_image_batch(image_url, image_size=image_size, pixel_size=pixel_size)
 
     input_image = images[0].copy()
     images = np.stack([preprocess(image) for image in images], axis=0)
@@ -118,8 +112,6 @@ def app():
 
         For example, try typing "gatto" (cat) or "cane" (dog) in the space for label and click "locate"!
 
-        *Depending on the server load, the computation time may vary. With normal load and pixel size 10, it can take up to two minutes.
-        *
         """
     )
 
@@ -133,9 +125,13 @@ def app():
     col1, col2 = st.beta_columns([3, 1])
 
     with col2:
-        pixel_size = st.selectbox("Pixel Size", options=range(10, 21, 5), index=1)
+        pixel_size = st.selectbox(
+            "Pixel Size", options=range(10, 21, 5), index=0
+        )
 
-        iterations = st.selectbox("Refinement Steps", options=range(3, 30, 3), index=0)
+        iterations = st.selectbox(
+            "Refinement Steps", options=range(3, 30, 3), index=0
+        )
 
         compute = st.button("LOCATE")
 
@@ -152,7 +148,7 @@ def app():
 
 
         if not caption or not image_url:
-            st.error("Please specify an image URL and a label")
+            st.error("Please choose one image and at least one label")
         else:
             with st.spinner("Computing..."):
                 heatmap, image = get_heatmap(image_url, caption, pixel_size, iterations)
@@ -164,7 +160,7 @@ def app():
         gc.collect()
 
     elif image_url:
-        image_raw = requests.get(image_url, stream=True,).raw
+        image_raw = requests.get(image_url, stream=True, ).raw
         image = Image.open(image_raw).convert("RGB")
         with col1:
             st.image(image)