Initial deployment with model and dependencies

app.py

@@ -1,7 +1,56 @@
+"""
+Interface for HuggingFace deployment
+"""
+
 import gradio as gr
+import numpy as np
+from src.model import AffordanceModel
+from src.utils.argument_utils import get_yaml_config
+import cv2
+
+print("Loading config...")
+config = get_yaml_config("checkpoints/gemini/config.yaml")
+print("Building model...")
+model = AffordanceModel(config)
+print("Model built successfully!")
+
+def predict(image, text):
+    """
+    Gradio inference function
+    Args:
+        image: PIL Image (Gradio's default image input format)
+        text: str
+    Returns:
+        visualization of the heatmap
+    """
+    # Convert PIL image to numpy array
+    image = np.array(image)
+    
+    # Run model inference
+    heatmap = model.inference(image, text)  # Returns (H, W) array
+    
+    # Visualize heatmap (convert to RGB for display)
+    # Scale to 0-255 and apply colormap
+    heatmap_vis = (heatmap * 255).astype(np.uint8)
+    heatmap_colored = cv2.applyColorMap(heatmap_vis, cv2.COLORMAP_JET)
+    heatmap_colored = cv2.cvtColor(heatmap_colored, cv2.COLOR_BGR2RGB)
+    
+    return heatmap_colored
 
-def greet(name):
-    return "Hello " + name + "!!"
+# Create Gradio interface
+demo = gr.Interface(
+    fn=predict,
+    inputs=[
+        gr.Image(type="pil", label="Input Image"),  # Accepts uploaded images
+        gr.Textbox(label="Text Query", placeholder="Enter text description...")
+    ],
+    outputs=gr.Image(label="Affordance Heatmap"),
+    title="Affordance Detection",
+    description="Upload an image and provide a text query to detect affordances.",
+    examples=[
+        ["test.png", "rim"]  # Add your test image and query
+    ]
+)
 
-demo = gr.Interface(fn=greet, inputs="text", outputs="text")
-demo.launch()
+if __name__ == "__main__":
+    demo.launch()

checkpoints/gemini/checkpoint.pth

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

checkpoints/gemini/config.yaml

@@ -0,0 +1,10 @@
+model:
+  in_channels: 1024
+  filters: [256, 64, 1]
+  kernel_sizes: [3, 3, 1]
+  strides: [1, 1, 1]
+  norm: None
+  activation: lrelu
+  lang_emb_dim: 1024
+  film_mode: zero
+checkpoint_path: checkpoints/gemini/checkpoint.pth

requirements.txt

@@ -0,0 +1,10 @@
+torch
+torchvision
+numpy
+opencv-python-headless  # using headless version for server deployment
+Pillow  # for PIL
+openai  # for OpenAI API
+gradio
+pyyaml  # for yaml
+matplotlib
+IPython

src/model.py

@@ -0,0 +1,92 @@
+"""
+Affordance model definition
+"""
+import torch
+import torch.nn.functional as F
+
+import os
+import sys
+import os.path as osp
+
+import numpy as np
+import cv2
+import matplotlib.pyplot as plt
+
+from .network import Conv2DFiLMNet
+from .utils.img_utils import *
+from .utils.argument_utils import get_yaml_config
+
+use_cuda = True
+device = torch.device("cuda" if (use_cuda and torch.cuda.is_available()) else "cpu")
+
+def build_network(config):
+    """
+    Build model
+    """
+    net = Conv2DFiLMNet(**config)
+
+    net.build()  
+    net = net.to(device) 
+
+    return net
+
+class AffordanceModel:
+
+    def __init__(self, config):
+
+        print("============ Building network and loading checkpoint =============")
+        # build network
+        self.model = build_network(config.model)
+
+        # load checkpoint
+        self.load_checkpoint(config.checkpoint_path)
+
+        print("============ Building DINO model =============")
+        torch_path = osp.join(osp.dirname(osp.dirname(__file__)), "data/torch_home")
+        self.dinov2 = load_pretrained_dino(torch_path=torch_path)
+
+    def load_checkpoint(self, checkpoint_path):
+        """
+        Load checkpoint
+        """
+        checkpoint = torch.load(checkpoint_path, map_location=device)
+        self.model.load_state_dict(checkpoint['model_state_dict'])
+        self.model.eval()
+
+    @torch.no_grad()
+    def inference(self, img, text, keep_orig_size=True, temp=1.0):
+        """
+        Inference model output on query image and text.
+        img: np.ndarray, (H, W, C)
+        text: str
+
+        Returns:
+        out: np array, (H, W)
+        """
+        self.model.eval()
+
+        # prepare input for model
+        img = rescale_img(img, max_size=672) # rescale image in case it is too large
+        processed_img, lang_emb = preprocess_data(img, text, self.model._lang_emb_dim)
+        img = torch.stack(processed_img, dim=0).to(device)
+        img_feat = get_dino_features(self.dinov2, img, repeat_to_orig_size=keep_orig_size).permute(0, 3, 1, 2)
+
+        lang_emb = lang_emb.to(device)
+
+        # forward pass
+        out = self.model(img_feat, lang_emb) # (1, 1, h, w)
+        out = out.squeeze()
+
+        # post-process output
+        out = torch.sigmoid(out)
+        out = out.detach().cpu().numpy()
+        return out
+    
+
+if __name__ == "__main__":
+    config_path = osp.join(osp.dirname(osp.dirname(__file__)), "checkpoints/gemini/config.yaml")
+    config = get_yaml_config(config_path)
+
+    affordance_model = AffordanceModel(config)
+
+    from IPython import embed; embed(); exit(0)

src/network.py

@@ -0,0 +1,271 @@
+"""
+Defines basic network building blocks and network architecture
+Some code adapted from PerAct: https://github.com/peract/peract
+"""
+
+import torch
+import torch.nn as nn
+
+from typing import List
+
+########################################
+### layers
+########################################
+LRELU_SLOPE = 0.02
+
+def act_layer(act):
+    if act == 'relu':
+        return nn.ReLU()
+    elif act == 'lrelu':
+        return nn.LeakyReLU(LRELU_SLOPE)
+    elif act == 'elu':
+        return nn.ELU()
+    elif act == 'tanh':
+        return nn.Tanh()
+    elif act == 'prelu':
+        return nn.PReLU()
+    else:
+        raise ValueError('%s not recognized.' % act)
+
+def norm_layer2d(norm, channels):
+    if norm == 'batch':
+        return nn.BatchNorm2d(channels)
+    elif norm == 'instance':
+        return nn.InstanceNorm2d(channels, affine=True)
+    elif norm == 'layer':
+        return nn.GroupNorm(1, channels, affine=True)
+    elif norm == 'group':
+        return nn.GroupNorm(4, channels, affine=True)
+    else:
+        raise ValueError('%s not recognized.' % norm)
+
+
+########################################
+### network blocks
+########################################
+
+class FiLMBlockRand(nn.Module):
+    """
+    FiLM block with random init gamma and beta. 
+    x = gamma * x + beta
+    Adapted from PerAct (and original FiLM paper)
+    """
+    def __init__(self, lang_emb_dim, num_channels):
+        super(FiLMBlockRand, self).__init__()
+
+        self.fc_gamma = nn.Linear(lang_emb_dim, num_channels)
+        self.fc_beta = nn.Linear(lang_emb_dim, num_channels)
+
+    def forward(self, x, lang_emb):
+        gamma = self.fc_gamma(lang_emb)
+        beta = self.fc_beta(lang_emb)
+
+        beta = beta.view(x.size(0), x.size(1), 1, 1)
+        gamma = gamma.view(x.size(0), x.size(1), 1, 1)
+
+        x = gamma * x + beta
+
+        return x
+    
+
+class FiLMBlockZero(nn.Module):
+    """
+    FiLM block with zero init gamma and beta.
+    x = (1 + gamma) * x + beta
+    Adapted from RT-1 https://github.com/google-research/robotics_transformer/blob/master/film_efficientnet/film_conditioning_layer.py
+    """
+    def __init__(self, lang_emb_dim, num_channels):
+        super(FiLMBlockZero, self).__init__()
+
+        self.fc_gamma = nn.Linear(lang_emb_dim, num_channels)
+        self.fc_beta = nn.Linear(lang_emb_dim, num_channels)
+
+        nn.init.zeros_(self.fc_gamma.weight)
+        nn.init.zeros_(self.fc_gamma.bias)
+        nn.init.zeros_(self.fc_beta.weight)
+        nn.init.zeros_(self.fc_beta.bias)
+
+    def forward(self, x, lang_emb):
+        gamma = self.fc_gamma(lang_emb)
+        beta = self.fc_beta(lang_emb)
+
+        beta = beta.view(x.size(0), x.size(1), 1, 1)
+        gamma = gamma.view(x.size(0), x.size(1), 1, 1)
+
+        x = (1 + gamma) * x + beta
+
+        return x
+    
+
+class Conv2DBlock(nn.Module):
+
+    def __init__(self, in_channels, out_channels, kernel_sizes, strides,
+                 norm=None, activation=None, padding_mode='replicate'):
+        super(Conv2DBlock, self).__init__()
+        padding = kernel_sizes // 2 if isinstance(kernel_sizes, int) else (
+            kernel_sizes[0] // 2, kernel_sizes[1] // 2)
+        self.conv2d = nn.Conv2d(
+            in_channels, out_channels, kernel_sizes, strides, padding=padding,
+            padding_mode=padding_mode)
+
+        if activation is None:
+            nn.init.xavier_uniform_(self.conv2d.weight,
+                                    gain=nn.init.calculate_gain('linear'))
+            nn.init.zeros_(self.conv2d.bias)
+        elif activation == 'tanh':
+            nn.init.xavier_uniform_(self.conv2d.weight,
+                                    gain=nn.init.calculate_gain('tanh'))
+            nn.init.zeros_(self.conv2d.bias)
+        elif activation == 'lrelu':
+            nn.init.kaiming_uniform_(self.conv2d.weight, a=LRELU_SLOPE,
+                                     nonlinearity='leaky_relu')
+            nn.init.zeros_(self.conv2d.bias)
+        elif activation == 'relu':
+            nn.init.kaiming_uniform_(self.conv2d.weight, nonlinearity='relu')
+            nn.init.zeros_(self.conv2d.bias)
+        else:
+            raise ValueError()
+
+        self.activation = activation
+        self.norm = norm
+        if norm is not None:
+            self.norm = norm_layer2d(norm, out_channels)
+        if activation is not None:
+            self.activation = act_layer(activation)
+
+    def forward(self, x):
+        x = self.conv2d(x)
+        x = self.norm(x) if self.norm is not None else x
+        x = self.activation(x) if self.activation is not None else x
+        return x
+    
+class Conv2DFiLMBlock(Conv2DBlock):
+
+    def __init__(self, in_channels, out_channels, kernel_sizes, strides,
+                 lang_emb_dim,
+                 norm=None, activation=None, padding_mode='replicate', 
+                 film_mode='rand', film_place='after'
+                 ):
+        super(Conv2DFiLMBlock, self).__init__(
+            in_channels, out_channels, kernel_sizes, strides, norm, activation,
+            padding_mode)
+        
+    
+        self.film_place = film_place    
+        if film_place == 'after':
+            film_channels = out_channels
+        elif film_place == 'before':
+            film_channels = in_channels
+        else:
+            raise ValueError(f"film_place {film_place} not recognized")
+        
+        if film_mode == 'rand':
+            self.film = FiLMBlockRand(lang_emb_dim, film_channels)
+        elif film_mode == 'zero':
+            self.film = FiLMBlockZero(lang_emb_dim, film_channels)
+        else:
+            raise ValueError(f"film_mode {film_mode} not recognized")
+
+    def forward(self, x, lang_emb):
+        if self.film_place == 'before':
+            x = self.film(x, lang_emb)
+            x = self.conv2d(x)
+            x = self.norm(x) if self.norm is not None else x
+            x = self.activation(x) if self.activation is not None else x
+
+        elif self.film_place == 'after':
+            x = self.conv2d(x) # (B, C, H, W)
+            x = self.norm(x) if self.norm is not None else x
+            x = self.film(x, lang_emb) # lang_emb: (B, lang_emb_dim), output: (B, C, H, W)
+            x = self.activation(x) if self.activation is not None else x
+
+        else:
+            raise ValueError(f"film_place {self.film_place} not recognized")
+        
+        return x
+    
+
+##############################################
+#### Network
+##############################################
+
+class Conv2DFiLMNet(nn.Module):
+
+    def __init__(self,
+                 in_channels: int,
+                 filters: List[int], # num of output channels for each Conv2D layer
+                 kernel_sizes: List[int], 
+                 strides: List[int], 
+                 norm: str = None, 
+                 activation: str = 'relu',
+
+                 lang_emb_dim: int = 256, 
+                 film_mode: str = 'zero',
+                 film_place: str = 'after'
+                 ):
+        super(Conv2DFiLMNet, self).__init__()
+
+        self._in_channels = in_channels
+        self._filters = filters
+        self._kernel_sizes = kernel_sizes
+        self._strides = strides
+        self._norm = norm
+        self._activation = activation
+
+        self._lang_emb_dim = lang_emb_dim
+        self._film_mode = film_mode
+        self._film_place = film_place
+
+    def build(self):
+        self.conv_blocks = nn.ModuleList()
+        for i in range(len(self._filters)):
+            in_channels = self._in_channels if i == 0 else self._filters[i-1]
+            out_channels = self._filters[i]
+            kernel_sizes = self._kernel_sizes[i]
+            strides = self._strides[i]
+            norm = self._norm 
+            activation = self._activation if i < len(self._filters) - 1 else None # no activation for the last layer
+            conv_block = Conv2DFiLMBlock(
+                in_channels, out_channels, kernel_sizes, strides,
+                self._lang_emb_dim,
+                norm=norm, activation=activation,
+                film_mode=self._film_mode,
+                film_place=self._film_place
+            )
+            self.conv_blocks.append(conv_block)
+
+    def forward(self, x, lang_emb):
+        """
+        Args:
+            x: (B, C, H, W)
+            lang_emb: (B, lang_emb_dim)
+        """
+        for conv_block in self.conv_blocks:
+            x = conv_block(x, lang_emb)
+        return x
+
+
+if __name__ == "__main__":
+
+    use_cuda = False
+    device = torch.device("cuda" if (use_cuda and torch.cuda.is_available()) else "cpu")
+    # from IPython import embed; embed(); exit(0)
+
+    # Test Conv2DFiLMNet
+    in_channels = 1024
+    filters = [256, 64, 1]
+    kernel_sizes = [3, 3, 1]
+    strides = [1, 1, 1]
+    norm = None
+    activation = 'lrelu'
+    lang_emb_dim = 1536
+    film_mode = 'zero'
+
+    net = Conv2DFiLMNet(
+        in_channels, filters, kernel_sizes, strides, norm, activation,
+        lang_emb_dim, film_mode
+    )
+
+    net.build()
+
+    from IPython import embed; embed(); exit(0)

src/utils/argument_utils.py

@@ -0,0 +1,128 @@
+"""load and save YAML config file. Originally from VoxPoser"""
+import os
+import sys
+import yaml
+import json
+
+class ConfigDict(dict):
+    def __init__(self, config):
+        """recursively build config"""
+        # self.config = config
+        for key, value in config.items():
+            if isinstance(value, str) and value.lower() == 'none':
+                value = None
+            if isinstance(value, dict):
+                self[key] = ConfigDict(value)
+            else:
+                self[key] = value
+
+    def __getattr__(self, key):
+        if key in self:
+            return self[key]
+        elif key == 'config':
+            return self
+        else:
+            raise AttributeError(f"'{self.__class__.__name__}' object has no attribute '{key}'")
+    def __setattr__(self, key, value):
+        self[key] = value
+    def __delattr__(self, key):
+        del self[key]
+    def __getstate__(self):
+        return self.config
+    def __setstate__(self, state):
+        self.config = state
+        self.__init__(state)
+
+    def update(self, config):
+        """update with another dict"""
+        if isinstance(config, ConfigDict):
+            config = config.convert_to_dict()
+        for key, value in config.items():
+            if isinstance(value, dict):
+                self[key].update(value)
+            else:
+                self[key] = value
+
+    def convert_to_dict(self):
+        """convert to dict"""
+        config = {}
+        for key, value in self.items():
+            if isinstance(value, ConfigDict):
+                config[key] = value.convert_to_dict()
+            else:
+                config[key] = value
+        return config
+
+def load_yaml_config(config_path):
+    with open(config_path, 'r') as f:
+        config = yaml.load(f, Loader=yaml.FullLoader)
+    return config
+
+def get_yaml_config(config_path=None):
+    assert config_path and os.path.exists(config_path), f'config file does not exist ({config_path})'
+    config = load_yaml_config(config_path)
+    config = ConfigDict(config)
+    return config
+
+def eval_str_to_lst(query_str):
+    """
+    Parse a string in format [a, b, c] to a list
+    """
+    query_str = query_str.replace('[', '').replace(']', '')
+    query_lst = query_str.split(',')
+    query_lst = [q.strip() for q in query_lst]
+    return query_lst
+
+def get_command_line_args(argv, to_config_dict=True):
+    """
+    Utility function to parse all command line arguments and return them as a dictionary.
+    If argument is in format 'key1.key2=value', it will be parsed as a nested dictionary.
+    """
+    args_dict = {}
+    for arg in argv[1:]:  # Skip the first argument (script name)
+        if '=' in arg:
+            key, value = arg.split('=', 1)
+            key = key.lstrip('-')  # Remove leading dashes
+
+            # Try to convert value to int, float, bool, or leave as string
+            try:
+                value = json.loads(value)
+            except json.JSONDecodeError:
+                pass
+
+            if isinstance(value, str) and '[' in value and ']' in value:
+                value = eval_str_to_lst(value)
+            
+            # Check for hierarchy in the key (indicated by '.')
+            if '.' in key:
+                sub_keys = key.split('.')
+                current_dict = args_dict
+                # Iterate through sub_keys to create nested dictionaries
+                for sub_key in sub_keys[:-1]:
+                    if sub_key not in current_dict:
+                        current_dict[sub_key] = {}
+                    current_dict = current_dict[sub_key]
+                current_dict[sub_keys[-1]] = value
+            else:
+                args_dict[key] = value
+
+    if to_config_dict:
+        args_dict = ConfigDict(args_dict)
+    return args_dict
+
+def save_config(config, config_path):
+    if not isinstance(config, dict):
+        raise ValueError("config must be a dictionary")
+    if type(config) != dict:
+        print("Converting config to dict")
+        config = config.convert_to_dict()
+    with open(config_path, 'w') as f:
+        yaml.dump(config, f, default_flow_style=False)
+
+# def main():
+#     config = get_yaml_config(config_path='./configs/sim_env/empty_scene_fetch.yaml')
+#     from IPython import embed; embed()
+
+if __name__ == '__main__':
+    from IPython import embed; embed(); exit(0)
+    # main()

src/utils/img_utils.py

@@ -0,0 +1,136 @@
+"""
+Image related util functions
+"""
+
+import cv2
+import numpy as np
+import os
+
+import torch
+import torch.nn.functional as F
+import torchvision.transforms as T
+from PIL import Image
+
+from openai import OpenAI
+
+##############################################
+### Image processing
+##############################################
+
+def rescale_img(img, max_size=448):
+    """
+    Rescale image such that largest dimension is max_size
+    img: np.ndarray, (H, W, C)
+    """
+    h, w = img.shape[:2]
+    if max(h, w) <= max_size:
+        return img
+    
+    scale = max_size / max(h, w)
+    new_h, new_w = int(h*scale), int(w*scale)
+    img = cv2.resize(img, (new_w, new_h), interpolation=cv2.INTER_AREA)
+    return img
+
+
+##############################################
+### DINO features
+##############################################
+
+def load_pretrained_dino(model_type='dinov2_vitl14', device=None):
+    if device is None:
+        device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    
+    dinov2 = torch.hub.load('facebookresearch/dinov2', model_type).eval()
+    dinov2 = dinov2.to(device)
+    return dinov2
+
+def get_dino_features(dinov2, imgs, repeat_to_orig_size=False):
+    """
+    Get features from DINO model
+    ::param dinov2:: DINO model
+    ::param imgs:: tensor of shape (bs, C, H, W)
+    """
+    bs, C, H, W = imgs.shape
+    patch_h = H // 14
+    patch_w = W // 14
+
+    with torch.no_grad():
+        features_dict = dinov2.forward_features(imgs)
+        features = features_dict['x_norm_patchtokens']
+        features = features.reshape(bs, patch_h, patch_w, -1)
+    
+    if not repeat_to_orig_size:
+        return features # (bs, patch_h, patch_w, n_features)
+    else:
+        # repeat on batched dims to original size
+        ratio = H // (patch_h*2)
+        features = F.interpolate(features.permute(0, 3, 1, 2), scale_factor=ratio, mode='bilinear').permute(0, 2, 3, 1)
+        return features
+
+def transform_imgs_for_dino(imgs, blur=True):
+    """
+    Transform image before passing to DINO model
+    ::param imgs:: np.array of shape (H, W, C) or (bs, H, W, C)
+    ::param blur:: bool, whether to apply Gaussian blur before resizing
+
+    ::return:: list of transformed images
+    """
+    # handles both single image and batch of images
+    if len(imgs.shape) == 3:
+        H, W, C = imgs.shape
+        imgs = imgs[None, ...]
+        bs = 1
+    else:
+        bs, H, W, C = imgs.shape
+
+    H *= 2
+    W *= 2
+
+    patch_h = H // 14
+    patch_w = W // 14
+
+    if blur:
+        transform_lst = [T.GaussianBlur(9, sigma=(1.0, 2.0))]
+    else:
+        transform_lst = []
+    transform_lst += [
+        T.Resize((patch_h * 14, patch_w * 14)),
+        T.CenterCrop((patch_h * 14, patch_w * 14)),
+        T.ToTensor(),
+        T.Normalize(mean=(0.485, 0.456, 0.406), std=(0.229, 0.224, 0.225)),
+    ]
+
+    transform = T.Compose(transform_lst)
+    
+    transformed_imgs = []
+    for i in range(bs):
+        temp = imgs[i].copy()
+        if temp.max() <= 1.1: # handle images with values in [0, 1]
+            temp = (temp * 255)
+        temp = temp.astype(np.uint8).clip(0, 255)
+        transformed_imgs.append(transform(Image.fromarray(temp)))
+    
+    return transformed_imgs
+
+
+##############################################
+### Preparing data for model
+##############################################
+
+def get_text_embedding(text, model="text-embedding-3-large", dim=1024):
+    """
+    Get text embedding with specified dimension
+    """
+    client = OpenAI(api_key=os.getenv('OPENAI_API_KEY'))
+    embedding = np.array(client.embeddings.create(input=[text], model=model, dimensions=dim).data[0].embedding)
+    return embedding
+
+def preprocess_data(img, text, lang_emb_dim=1024):
+    if len(img.shape) == 3 and img.shape[2] == 4:
+        img = img[:, :, :3]
+    elif len(img.shape) == 2:
+        img = np.stack([img] * 3, axis=-1)
+    processed_img = transform_imgs_for_dino(img)
+    lang_emb = torch.tensor(get_text_embedding(text, dim=lang_emb_dim)).to(torch.float32)
+
+    return processed_img, lang_emb