app.py

import os
from typing import Callable

from clu import checkpoint
from flax import linen as nn
import gradio as gr
from huggingface_hub import snapshot_download
import jax
import jax.numpy as jnp
import numpy as np
from PIL import Image

from invariant_slot_attention.configs.clevr_with_masks.equiv_transl_scale import get_config
from invariant_slot_attention.lib import utils


def load_model(config, checkpoint_dir):
  rng = jax.random.PRNGKey(42)

  # Initialize model
  model = utils.build_model_from_config(config.model)

  def init_model(rng):
    rng, init_rng, model_rng, dropout_rng = jax.random.split(rng, num=4)

    init_conditioning = None
    init_inputs = jnp.ones([1, 1, 128, 128, 3], jnp.float32)
    initial_vars = model.init(
        {"params": model_rng, "state_init": init_rng, "dropout": dropout_rng},
        video=init_inputs, conditioning=init_conditioning,
        padding_mask=jnp.ones(init_inputs.shape[:-1], jnp.int32))
 
    # Split into state variables (e.g. for batchnorm stats) and model params.
    # Note that `pop()` on a FrozenDict performs a deep copy.
    state_vars, initial_params = initial_vars.pop("params")  # pytype: disable=attribute-error

    # Filter out intermediates (we don't want to store these in the TrainState).
    state_vars = utils.filter_key_from_frozen_dict(
        state_vars, key="intermediates")
    return state_vars, initial_params

  state_vars, initial_params = init_model(rng)

  opt_state = None
  state = utils.TrainState(
      step=1, opt_state=opt_state, params=initial_params, rng=rng,
      variables=state_vars)

  ckpt = checkpoint.Checkpoint(checkpoint_dir)
  state = ckpt.restore(state, checkpoint=checkpoint_dir + "/ckpt-0")

  return model, state, rng


def load_image(name):
  img = Image.open(f"images/{name}.png")
  img = img.crop((64, 29, 64 + 192, 29 + 192))
  img = img.resize((128, 128))
  img = np.array(img)[:, :, :3] / 255.
  img = jnp.array(img, dtype=jnp.float32)
  return img


download_path = snapshot_download(repo_id="ondrejbiza/isa")
checkpoint_dir = os.path.join(download_path, "clevr_isa_ts")

model, state, rng = load_model(get_config(), checkpoint_dir)

rng, init_rng = jax.random.split(rng, num=2)


class DecoderWrapper(nn.Module):
    decoder: Callable[[], nn.Module]
    @nn.compact
    def __call__(self, slots, train=False):
        return self.decoder()(slots, train)
decoder_model = DecoderWrapper(decoder=model.decoder)

with gr.Blocks() as demo:

    local_slots = gr.State(np.zeros((11, 64), dtype=np.float32))

    local_orig_pos = gr.State(np.ones((11, 2), dtype=np.float32))
    local_orig_scale = gr.State(np.ones((11, 2), dtype=np.float32))

    local_pos = gr.State(np.zeros((11, 2), dtype=np.float32))
    local_scale = gr.State(np.zeros((11, 2), dtype=np.float32))

    local_probs = gr.State(np.zeros((11, 128, 128), dtype=np.float32))

    with gr.Row():

        gr_choose_image = gr.Dropdown(
            [f"img{i}" for i in range(1, 9)], label="CLEVR Image", info="Start by a picking an image from the CLEVR dataset."
        )

    with gr.Row():

      with gr.Column():

        with gr.Row():
          with gr.Column():
            gr_image_1 = gr.Image(type="numpy", shape=(112, 112), source="canvas", label="Decoding")
          with gr.Column():
            gr_image_2 = gr.Image(type="numpy", shape=(112, 112), source="canvas", label="Segmentation")

      with gr.Column():
        gr_slot_slider = gr.Slider(1, 11, value=1, step=1, label="Slot Index",
                                   info="Change slot index too see the segmentation mask, position and scale of each slot.")

        gr_y_slider = gr.Slider(-1, 1, value=0, step=0.01, label="X")
        gr_x_slider = gr.Slider(-1, 1, value=0, step=0.01, label="Y")
        gr_sy_slider = gr.Slider(0.5, 1.5, value=1., step=0.1, label="Width Multiplier")
        gr_sx_slider = gr.Slider(0.5, 1.5, value=1., step=0.1, label="Height Multiplier")

        with gr.Row():
          with gr.Column():
            gr_button_render = gr.Button("Render", variant="primary", info="Render a new image with altered positions and scales.")
          with gr.Column():
            gr_button_reset = gr.Button("Reset", info="Reset slot statistics.")

    def update_image_and_segmentation(name, idx):
      idx = idx - 1

      img_input = load_image(name)
      out = model.apply(
        {"params": state.params, **state.variables},
        video=img_input[None, None],
        rngs={"state_init": init_rng},
        train=False)
      
      probs = np.array(nn.softmax(out["outputs"]["segmentation_logits"][0, 0, :, :, :, 0], axis=0))
      img = np.array(out["outputs"]["video"][0, 0])
      img = np.clip(img, 0, 1)

      slots_ = np.array(out["states"])
      slots = slots_[0, 0, :, :-4]
      pos = slots_[0, 0, :, -4: -2]
      scale = slots_[0, 0, :, -2:]

      return (img * 255).astype(np.uint8), (probs[idx] * 255).astype(np.uint8), float(pos[idx, 0]), \
             float(pos[idx, 1]), probs, slots, pos, np.ones((11, 2), dtype=np.float32), pos, scale

    gr_choose_image.change(
       fn=update_image_and_segmentation,
       inputs=[gr_choose_image, gr_slot_slider],
       outputs=[gr_image_1, gr_image_2, gr_x_slider, gr_y_slider, local_probs,
                local_slots, local_pos, local_scale, local_orig_pos, local_orig_scale]
    )

    def update_sliders(idx, local_probs, local_pos, local_scale):
      idx = idx - 1  # 1-indexing to 0-indexing
      return (local_probs[idx] * 255).astype(np.uint8), float(local_pos[idx, 0]), \
             float(local_pos[idx, 1]), float(local_scale[idx, 0]), float(local_scale[idx, 1])

    gr_slot_slider.change(
      fn=update_sliders,
      inputs=[gr_slot_slider, local_probs, local_pos, local_scale],
      outputs=[gr_image_2, gr_x_slider, gr_y_slider, gr_sx_slider, gr_sy_slider]
    )

    def update_pos_x(idx, val, local_pos):
       local_pos = np.copy(local_pos)
       local_pos[idx - 1, 0] = val
       return local_pos

    def update_pos_y(idx, val, local_pos):
       local_pos = np.copy(local_pos)
       local_pos[idx - 1, 1] = val
       return local_pos

    def update_scale_x(idx, val, local_scale):
       local_scale = np.copy(local_scale)
       local_scale[idx - 1, 0] = val
       return local_scale

    def update_scale_y(idx, val, local_scale):
       local_scale = np.copy(local_scale)
       local_scale[idx - 1, 1] = val
       return local_scale

    gr_x_slider.change(
       fn=update_pos_x,
       inputs=[gr_slot_slider, gr_x_slider, local_pos],
       outputs=local_pos
    )
    gr_y_slider.change(
       fn=update_pos_y,
       inputs=[gr_slot_slider, gr_y_slider, local_pos],
       outputs=local_pos
    )
    gr_sx_slider.change(
       fn=update_scale_x,
       inputs=[gr_slot_slider, gr_sx_slider, local_scale],
       outputs=local_scale
    )
    gr_sy_slider.change(
       fn=update_scale_y,
       inputs=[gr_slot_slider, gr_sy_slider, local_scale],
       outputs=local_scale
    )

    def render(idx, local_slots, local_pos, local_scale, local_orig_scale):
      idx = idx - 1

      slots = np.concatenate([local_slots, local_pos, local_scale * local_orig_scale], axis=-1)
      slots = jnp.array(slots)

      out = decoder_model.apply(
        {"params": state.params, **state.variables},
        slots=slots[None, None],
        train=False
      )

      probs = np.array(nn.softmax(out["segmentation_logits"][0, 0, :, :, :, 0], axis=0))
      image = np.array(out["video"][0, 0])
      image = np.clip(image, 0, 1)
      return (image * 255).astype(np.uint8), (probs[idx] * 255).astype(np.uint8), probs

    gr_button_render.click(
        fn=render,
        inputs=[gr_slot_slider, local_slots, local_pos, local_scale, local_orig_scale],
        outputs=[gr_image_1, gr_image_2, local_probs]
    )

    def reset(idx, local_orig_pos):
       idx = idx - 1
       return np.copy(local_orig_pos), np.ones((11, 2), dtype=np.float32), float(local_orig_pos[idx, 0]), \
              float(local_orig_pos[idx, 1]), 1., 1.

    gr_button_reset.click(
        fn=reset,
        inputs=[gr_slot_slider, local_orig_pos],
        outputs=[local_pos, local_scale, gr_x_slider, gr_y_slider, gr_sx_slider, gr_sy_slider]
    )

demo.launch()