Upload 5 files

README.md

@@ -1,8 +1,8 @@
 ---
-title: Cellvisionai
-emoji: 🌍
+title: CellVision AI -- Intelligent Cell Imaging Analysis
+emoji: 🔬🧫
 colorFrom: indigo
-colorTo: yellow
+colorTo: gray
 sdk: gradio
 sdk_version: 4.36.1
 app_file: app.py

app.py

@@ -0,0 +1,281 @@
+import functools
+import re
+
+import PIL.Image
+import gradio as gr
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+import flax.linen as nn
+from inference import PaliGemmaModel as pali_gemma_model
+
+COLORS = ['#4285f4', '#db4437', '#f4b400', '#0f9d58', '#e48ef1']
+
+##### Parse segmentation output tokens into masks
+##### Also returns bounding boxes with their labels
+
+def parse_segmentation(input_image, input_text):
+  out = pali_gemma_model.infer(input_image, input_text, max_new_tokens=100)
+  objs = extract_objs(out.lstrip("\n"), input_image.size[0], input_image.size[1], unique_labels=True)
+  labels = set(obj.get('name') for obj in objs if obj.get('name'))
+  color_map = {l: COLORS[i % len(COLORS)] for i, l in enumerate(labels)}
+  highlighted_text = [(obj['content'], obj.get('name')) for obj in objs]
+  annotated_img = (
+    input_image,
+    [
+        (
+            obj['mask'] if obj.get('mask') is not None else obj['xyxy'],
+            obj['name'] or '',
+        )
+        for obj in objs
+        if 'mask' in obj or 'xyxy' in obj
+    ],
+)
+  has_annotations = bool(annotated_img[1])
+  return annotated_img
+
+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.
+"""
+
+with gr.Blocks(css="style.css") as demo:
+  gr.Markdown(INTRO_TEXT)
+  with gr.Tab("Segment/Detect"):
+    with gr.Row():
+        with gr.Column():
+          image = gr.Image(type="pil")
+          seg_input = gr.Text(label="Entities to Segment/Detect")
+    
+        with gr.Column():
+          annotated_image = gr.AnnotatedImage(label="Output")
+
+    seg_btn = gr.Button("Submit")    
+    examples = [["./examples/cart1.jpg", "segment cells"],
+                ["./examples/cart1.jpg", "detect cells"],
+                ["./examples/cart2.jpg", "segment cells"],
+                ["./examples/cart2.jpg", "detect cells"],
+                ["./examples/cart3.jpg", "segment cells"],
+                ["./examples/cart3.jpg", "detect cells"]]
+    gr.Examples(
+        examples=examples,
+        inputs=[image, seg_input],
+    )
+    seg_inputs = [
+        image,
+        seg_input
+        ]
+    seg_outputs = [
+        annotated_image
+    ]
+    seg_btn.click(
+        fn=parse_segmentation,
+        inputs=seg_inputs,
+        outputs=seg_outputs,
+    )
+  with gr.Tab("Text Generation"):
+    with gr.Column():
+        image = gr.Image(type="pil")
+        text_input = gr.Text(label="Input Text")
+
+        text_output = gr.Text(label="Text Output")
+        chat_btn = gr.Button()
+        tokens = gr.Slider(
+            label="Max New Tokens",
+            info="Set to larger for longer generation.",
+            minimum=10,
+            maximum=100,
+            value=50,
+            step=10,
+        )
+
+    chat_inputs = [
+        image,
+        text_input,
+        tokens
+        ]
+    chat_outputs = [
+        text_output
+    ]
+    chat_btn.click(
+        fn=pali_gemma_model.infer,
+        inputs=chat_inputs,
+        outputs=chat_outputs,
+    )
+    
+    examples = [["./examples/cart1.jpg", IMAGE_PROMPT],
+                ["./examples/cart2.jpg", IMAGE_PROMPT],
+                ["./examples/cart3.jpg", IMAGE_PROMPT]]
+    gr.Examples(
+        examples=examples,
+        inputs=chat_inputs,
+    )
+
+
+### Postprocessing Utils for Segmentation Tokens
+### Segmentation tokens are passed to another VAE which decodes them to a mask
+
+_MODEL_PATH = 'vae-oid.npz'
+
+_SEGMENT_DETECT_RE = re.compile(
+    r'(.*?)' +
+    r'<loc(\d{4})>' * 4 + r'\s*' +
+    '(?:%s)?' % (r'<seg(\d{3})>' * 16) +
+    r'\s*([^;<>]+)? ?(?:; )?',
+)
+
+
+def _get_params(checkpoint):
+  """Converts PyTorch checkpoint to Flax params."""
+
+  def transp(kernel):
+    return np.transpose(kernel, (2, 3, 1, 0))
+
+  def conv(name):
+    return {
+        'bias': checkpoint[name + '.bias'],
+        'kernel': transp(checkpoint[name + '.weight']),
+    }
+
+  def resblock(name):
+    return {
+        'Conv_0': conv(name + '.0'),
+        'Conv_1': conv(name + '.2'),
+        'Conv_2': conv(name + '.4'),
+    }
+
+  return {
+      '_embeddings': checkpoint['_vq_vae._embedding'],
+      'Conv_0': conv('decoder.0'),
+      'ResBlock_0': resblock('decoder.2.net'),
+      'ResBlock_1': resblock('decoder.3.net'),
+      'ConvTranspose_0': conv('decoder.4'),
+      'ConvTranspose_1': conv('decoder.6'),
+      'ConvTranspose_2': conv('decoder.8'),
+      'ConvTranspose_3': conv('decoder.10'),
+      'Conv_1': conv('decoder.12'),
+  }
+
+
+def _quantized_values_from_codebook_indices(codebook_indices, embeddings):
+  batch_size, num_tokens = codebook_indices.shape
+  assert num_tokens == 16, codebook_indices.shape
+  unused_num_embeddings, embedding_dim = embeddings.shape
+
+  encodings = jnp.take(embeddings, codebook_indices.reshape((-1)), axis=0)
+  encodings = encodings.reshape((batch_size, 4, 4, embedding_dim))
+  return encodings
+
+
+@functools.cache
+def _get_reconstruct_masks():
+  """Reconstructs masks from codebook indices.
+  Returns:
+    A function that expects indices shaped `[B, 16]` of dtype int32, each
+    ranging from 0 to 127 (inclusive), and that returns a decoded masks sized
+    `[B, 64, 64, 1]`, of dtype float32, in range [-1, 1].
+  """
+
+  class ResBlock(nn.Module):
+    features: int
+
+    @nn.compact
+    def __call__(self, x):
+      original_x = x
+      x = nn.Conv(features=self.features, kernel_size=(3, 3), padding=1)(x)
+      x = nn.relu(x)
+      x = nn.Conv(features=self.features, kernel_size=(3, 3), padding=1)(x)
+      x = nn.relu(x)
+      x = nn.Conv(features=self.features, kernel_size=(1, 1), padding=0)(x)
+      return x + original_x
+
+  class Decoder(nn.Module):
+    """Upscales quantized vectors to mask."""
+
+    @nn.compact
+    def __call__(self, x):
+      num_res_blocks = 2
+      dim = 128
+      num_upsample_layers = 4
+
+      x = nn.Conv(features=dim, kernel_size=(1, 1), padding=0)(x)
+      x = nn.relu(x)
+
+      for _ in range(num_res_blocks):
+        x = ResBlock(features=dim)(x)
+
+      for _ in range(num_upsample_layers):
+        x = nn.ConvTranspose(
+            features=dim,
+            kernel_size=(4, 4),
+            strides=(2, 2),
+            padding=2,
+            transpose_kernel=True,
+        )(x)
+        x = nn.relu(x)
+        dim //= 2
+
+      x = nn.Conv(features=1, kernel_size=(1, 1), padding=0)(x)
+
+      return x
+
+  def reconstruct_masks(codebook_indices):
+    quantized = _quantized_values_from_codebook_indices(
+        codebook_indices, params['_embeddings']
+    )
+    return Decoder().apply({'params': params}, quantized)
+
+  with open(_MODEL_PATH, 'rb') as f:
+    params = _get_params(dict(np.load(f)))
+
+  return jax.jit(reconstruct_masks, backend='cpu')
+
+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, = _get_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
+
+#########
+
+if __name__ == "__main__":
+    demo.queue(max_size=10).launch(debug=True)

inference.py

@@ -0,0 +1,27 @@
+import os
+import torch
+from huggingface_hub import login
+from transformers import PaliGemmaForConditionalGeneration, PaliGemmaProcessor
+import spaces
+
+hf_token = os.getenv("HF_TOKEN")
+login(token=hf_token, add_to_git_credential=True)
+
+class PaliGemmaModel:
+    def __init__(self):
+        self.model_id = "google/paligemma-3b-mix-448"
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.model = PaliGemmaForConditionalGeneration.from_pretrained(self.model_id).eval().to(self.device)
+        self.processor = PaliGemmaProcessor.from_pretrained(self.model_id)
+
+    @spaces.GPU
+    def infer(self, image: PIL.Image.Image, text: str, max_new_tokens: int) -> str:
+        inputs = self.processor(text=text, images=image, return_tensors="pt").to(self.device)
+        with torch.inference_mode():
+            generated_ids = self.model.generate(
+                **inputs,
+                max_new_tokens=max_new_tokens,
+                do_sample=False
+            )
+        result = self.processor.batch_decode(generated_ids, skip_special_tokens=True)
+        return result[0][len(text):].lstrip("\n")

vae-oid.npz

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5586010257b8536dddefab65e7755077f21d5672d5674dacf911f73ae95a4447
+size 8479556