Update app.py

app.py

@@ -1,3 +1,4 @@
+import functools
 import re
 import PIL.Image
 import gradio as gr
@@ -36,12 +37,62 @@ def parse_segmentation(input_image, input_text, max_new_tokens=100):
     has_annotations = bool(annotated_img[1])
     return annotated_img
 
-INTRO_TEXT="🔬🧠 CellVision AI -- Intelligent Cell Imaging Analysis 🤖🧫"
-IMAGE_PROMPT="""
+INTRO_TEXT = "🔬🧠 CellVision AI -- Intelligent Cell Imaging Analysis 🤖🧫"
+IMAGE_PROMPT = """
 Describe the morphological characteristics and visible interactions between different cell types.
 Assess the biological context to identify signs of cancer and the presence of antigens.
 """
 
+def extract_objs(text, width, height, unique_labels=False):
+    """Returns objs for a string with "<loc>" and "<seg>" tokens."""
+    objs = []
+    seen = set()
+    while text:
+        m = _SEGMENT_DETECT_RE.match(text)
+        if not m:
+            break
+        print("m", m)
+        gs = list(m.groups())
+        before = gs.pop(0)
+        name = gs.pop()
+        y1, x1, y2, x2 = [int(x) / 1024 for x in gs[:4]]
+
+        y1, x1, y2, x2 = map(round, (y1 * height, x1 * width, y2 * height, x2 * width))
+        seg_indices = gs[4:20]
+        if seg_indices[0] is None:
+            mask = None
+        else:
+            seg_indices = np.array([int(x) for x in seg_indices], dtype=np.int32)
+            m64, = vae_model.reconstruct_masks(seg_indices[None])[..., 0]
+            m64 = np.clip(np.array(m64) * 0.5 + 0.5, 0, 1)
+            m64 = PIL.Image.fromarray((m64 * 255).astype('uint8'))
+            mask = np.zeros([height, width])
+            if y2 > y1 and x2 > x1:
+                mask[y1:y2, x1:x2] = np.array(m64.resize([x2 - x1, y2 - y1])) / 255.0
+
+        content = m.group()
+        if before:
+            objs.append(dict(content=before))
+            content = content[len(before):]
+        while unique_labels and name in seen:
+            name = (name or '') + "'"
+        seen.add(name)
+        objs.append(dict(
+            content=content, xyxy=(x1, y1, x2, y2), mask=mask, name=name))
+        text = text[len(before) + len(content):]
+
+    if text:
+        objs.append(dict(content=text))
+
+    return objs
+
+_SEGMENT_DETECT_RE = re.compile(
+    r'(.*?)' +
+    r'<loc(\d{4})>' * 4 + r'\s*' +
+    '(?:%s)?' % (r'<seg(\d{3})>' * 16) +
+    r'\s*([^;<>]+)? ?(?:; )?',
+)
+
 with gr.Blocks(css="style.css") as demo:
     gr.Markdown(INTRO_TEXT)
     with gr.Tab("Segment/Detect"):
@@ -117,58 +168,5 @@ with gr.Blocks(css="style.css") as demo:
             inputs=chat_inputs,
         )
 
-### Postprocessing Utils for Segmentation Tokens
-### Segmentation tokens are passed to another VAE which decodes them to a mask
-
-def extract_objs(text, width, height, unique_labels=False):
-    """Returns objs for a string with "<loc>" and "<seg>" tokens."""
-    objs = []
-    seen = set()
-    while text:
-        m = _SEGMENT_DETECT_RE.match(text)
-        if not m:
-            break
-        print("m", m)
-        gs = list(m.groups())
-        before = gs.pop(0)
-        name = gs.pop()
-        y1, x1, y2, x2 = [int(x) / 1024 for x in gs[:4]]
-
-        y1, x1, y2, x2 = map(round, (y1*height, x1*width, y2*height, x2*width))
-        seg_indices = gs[4:20]
-        if seg_indices[0] is None:
-            mask = None
-        else:
-            seg_indices = np.array([int(x) for x in seg_indices], dtype=np.int32)
-            m64, = vae_model.reconstruct_masks(seg_indices[None])[..., 0]
-            m64 = np.clip(np.array(m64) * 0.5 + 0.5, 0, 1)
-            m64 = PIL.Image.fromarray((m64 * 255).astype('uint8'))
-            mask = np.zeros([height, width])
-            if y2 > y1 and x2 > x1:
-                mask[y1:y2, x1:x2] = np.array(m64.resize([x2 - x1, y2 - y1])) / 255.0
-
-        content = m.group()
-        if before:
-            objs.append(dict(content=before))
-            content = content[len(before):]
-        while unique_labels and name in seen:
-            name = (name or '') + "'"
-        seen.add(name)
-        objs.append(dict(
-            content=content, xyxy=(x1, y1, x2, y2), mask=mask, name=name))
-        text = text[len(before) + len(content):]
-
-    if text:
-        objs.append(dict(content=text))
-
-    return objs
-
-_SEGMENT_DETECT_RE = re.compile(
-    r'(.*?)' +
-    r'<loc(\d{4})>' * 4 + r'\s*' +
-    '(?:%s)?' % (r'<seg(\d{3})>' * 16) +
-    r'\s*([^;<>]+)? ?(?:; )?',
-)
-
 if __name__ == "__main__":
     demo.queue(max_size=10).launch(debug=True)