demo buildout

.idea/.gitignore

@@ -0,0 +1,8 @@
+# Default ignored files
+/shelf/
+/workspace.xml
+# Editor-based HTTP Client requests
+/httpRequests/
+# Datasource local storage ignored files
+/dataSources/
+/dataSources.local.xml

.idea/curacel-demo-1.iml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<module type="PYTHON_MODULE" version="4">
+  <component name="NewModuleRootManager">
+    <content url="file://$MODULE_DIR$" />
+    <orderEntry type="inheritedJdk" />
+    <orderEntry type="sourceFolder" forTests="false" />
+  </component>
+</module>

.idea/inspectionProfiles/Project_Default.xml

@@ -0,0 +1,6 @@
+<component name="InspectionProjectProfileManager">
+  <profile version="1.0">
+    <option name="myName" value="Project Default" />
+    <inspection_tool class="Eslint" enabled="true" level="WARNING" enabled_by_default="true" />
+  </profile>
+</component>

.idea/inspectionProfiles/profiles_settings.xml

@@ -0,0 +1,6 @@
+<component name="InspectionProjectProfileManager">
+  <settings>
+    <option name="USE_PROJECT_PROFILE" value="false" />
+    <version value="1.0" />
+  </settings>
+</component>

.idea/modules.xml

@@ -0,0 +1,8 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="ProjectModuleManager">
+    <modules>
+      <module fileurl="file://$PROJECT_DIR$/.idea/curacel-demo-1.iml" filepath="$PROJECT_DIR$/.idea/curacel-demo-1.iml" />
+    </modules>
+  </component>
+</project>

.idea/vcs.xml

@@ -0,0 +1,6 @@
+<?xml version="1.0" encoding="UTF-8"?>
+<project version="4">
+  <component name="VcsDirectoryMappings">
+    <mapping directory="" vcs="Git" />
+  </component>
+</project>

app.py

@@ -1,63 +1,336 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
+import PIL.Image
+import transformers
+from transformers import PaliGemmaForConditionalGeneration, PaliGemmaProcessor
+import torch
+import os
+import string
+import functools
+import re
+import flax.linen as nn
+import jax
+import jax.numpy as jnp
+import numpy as np
+import spaces
 
-"""
-For more information on `huggingface_hub` Inference API support, please check the docs: https://huggingface.co/docs/huggingface_hub/v0.22.2/en/guides/inference
-"""
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+model_id = "mattraj/curacel-transcription-1"
+COLORS = ['#4285f4', '#db4437', '#f4b400', '#0f9d58', '#e48ef1']
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = PaliGemmaForConditionalGeneration.from_pretrained(model_id).eval().to(device)
+processor = PaliGemmaProcessor.from_pretrained(model_id)
 
+def resize_and_pad(image, target_dim):
+    # Calculate the aspect ratio
+    scale_factor = 1
+    aspect_ratio = image.width / image.height
+    if aspect_ratio > 1:
+        # Width is greater than height
+        new_width = int(target_dim * scale_factor)
+        new_height = int((target_dim / aspect_ratio) * scale_factor)
+    else:
+        # Height is greater than width
+        new_height = int(target_dim * scale_factor)
+        new_width = int(target_dim * aspect_ratio * scale_factor)
 
-def respond(
-    message,
-    history: list[tuple[str, str]],
-    system_message,
-    max_tokens,
-    temperature,
-    top_p,
-):
-    messages = [{"role": "system", "content": system_message}]
+    resized_image = image.resize((new_width, new_height), Image.LANCZOS)
 
-    for val in history:
-        if val[0]:
-            messages.append({"role": "user", "content": val[0]})
-        if val[1]:
-            messages.append({"role": "assistant", "content": val[1]})
+    # Create a new image with the target dimensions and a white background
+    new_image = Image.new("RGB", (target_dim, target_dim), (255, 255, 255))
+    new_image.paste(resized_image, ((target_dim - new_width) // 2, (target_dim - new_height) // 2))
 
-    messages.append({"role": "user", "content": message})
+    return new_image
 
-    response = ""
 
-    for message in client.chat_completion(
-        messages,
-        max_tokens=max_tokens,
-        stream=True,
-        temperature=temperature,
-        top_p=top_p,
-    ):
-        token = message.choices[0].delta.content
+###### Transformers Inference
+@spaces.GPU
+def infer(
+        image: PIL.Image.Image,
+        text: str,
+        max_new_tokens: int
+) -> str:
+    inputs = processor(text=text, images=resize_and_pad(image), return_tensors="pt").to(device)
+    with torch.inference_mode():
+        generated_ids = model.generate(
+            **inputs,
+            max_new_tokens=max_new_tokens,
+            do_sample=False
+        )
+    result = processor.batch_decode(generated_ids, skip_special_tokens=True)
+    return result[0][len(text):].lstrip("\n")
 
-        response += token
-        yield response
 
+##### Parse segmentation output tokens into masks
+##### Also returns bounding boxes with their labels
+
+def parse_segmentation(input_image, input_text):
+    out = 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
+
+
+######## Demo
+
+INTRO_TEXT = """## Curacel Handwritten Arabic demo\n\n
+Finetuned from: google/paligemma-3b-pt-448
+
+
+Translation model demo at: https://prod.arabic-gpt.ai/
+
+Prompts:
+Translate the Arabic to English: {model output}
+
+The following is a diagnosis in Arabic from a medical billing form we need to translate to English. The transcriber is not necessariily accurate so one or more characters or words may be wrong. Given what is written, what is the most likely diagnosis. Think step by step, and think about similar words or mispellings in Arabic. Give multiple arabic diagnoses along with the translation in English for each, then finally select the diagnosis that makes the most sense given what was transcribed and print the English translation as your most likely final translation. Transcribed text:  {model output}
 """
-For information on how to customize the ChatInterface, peruse the gradio docs: https://www.gradio.app/docs/chatinterface
-"""
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="You are a friendly Chatbot.", label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.95,
-            step=0.05,
-            label="Top-p (nucleus sampling)",
-        ),
-    ],
+
+with gr.Blocks(css="style.css") as demo:
+    gr.Markdown(INTRO_TEXT)
+    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=20,
+                step=10,
+            )
+
+        chat_inputs = [
+            image,
+            text_input,
+            tokens
+        ]
+        chat_outputs = [
+            text_output
+        ]
+        chat_btn.click(
+            fn=infer,
+            inputs=chat_inputs,
+            outputs=chat_outputs,
+        )
+
+        examples = [["./diagnosis-1.jpg", "Transcribe the Arabic text."],
+                    ["./examples/sign.jpg", "Transcribe the Arabic text."]]
+        gr.Markdown("")
+
+        gr.Examples(
+            examples=examples,
+            inputs=chat_inputs,
+        )
+'''        
+    with gr.Tab("Segment/Detect"):
+        image = gr.Image(type="pil")
+        seg_input = gr.Text(label="Entities to Segment/Detect")
+        seg_btn = gr.Button("Submit")
+        annotated_image = gr.AnnotatedImage(label="Output")
+
+        examples = [["./diagnosis-1.jpg", "Transcribe the Arabic text."],
+                    ["./examples/sign.jpg", "Transcribe the Arabic text."]]
+        gr.Markdown(
+            "")
+        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,
+        )
+'''
+
+### 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.launch()
+    demo.queue(max_size=10).launch(debug=True)