Object detection app

app.py

@@ -0,0 +1,67 @@
+import torch
+import torch.nn as nn
+import cv2
+import gradio as gr
+import numpy as np
+from PIL import Image
+import transformers
+from transformers import RobertaModel, RobertaTokenizer
+import timm
+import pandas as pd 
+import matplotlib.pyplot as plt
+from timm.data import resolve_data_config
+from timm.data.transforms_factory import create_transform
+
+from model import Model
+import get_output
+
+
+# Use GPU if available
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+# Initialize models
+vit = timm.create_model('vit_base_patch16_224', pretrained=True, num_classes=0, global_pool='').to(device)
+tokenizer = RobertaTokenizer.from_pretrained('roberta-base', truncation=True, do_lower_case=True)
+roberta = RobertaModel.from_pretrained("roberta-base")
+model = Model(vit, roberta, tokenizer, device).to(device)
+model.eval()
+
+config = resolve_data_config({}, model=vit)
+config['no_aug'] = True
+config['interpolation'] = 'bilinear'
+transform = create_transform(**config)
+
+
+def query_image(input_img, query, binarize, eval_threshold):
+
+    PIL_image = Image.fromarray(input_img, "RGB")
+    img = transform(PIL_image)
+    img = torch.unsqueeze(img,0).to(device)
+
+    with torch.no_grad():
+        output = model(img, query)
+
+    img = visualize_output(img, output, binarize, eval_threshold)
+    return img
+
+
+description = """
+Gradio demo for an object detection architecture, 
+introduced in <a href="https://arxiv.org/abs/2205.06230">my bachelor thesis</a>. 
+\n\nLorem ipsum ....
+*"image of a shoe"*. Refer to the <a href="https://arxiv.org/abs/2103.00020">CLIP</a> paper to see the full list of text templates used to augment the training data.
+"""
+demo = gr.Interface(
+    query_image, 
+    inputs=[gr.Image(), "text", "checkbox", gr.Slider(0, 1, value=0.25)], 
+    outputs="image",
+    title="Object Detection Using Textual Queries",
+    description=description,
+    examples=[
+        ["examples/img1.jpeg", "Find a person.", True, 0.25],
+    ],
+    allow_flagging = "never",
+    cache_examples=False,
+)
+demo.launch(debug=True)
+

get_output.py

@@ -0,0 +1,46 @@
+import numpy as np
+import pandas as pd
+import matplotlib.pyplot as plt
+from PIL import Image
+
+
+def preprocess(image, output, binarize, threshold):
+   
+    image = image.cpu().detach().numpy().squeeze()
+    image = np.transpose(image,(1,2,0))
+    image = (image + 1) * 0.5
+    output = output.cpu().detach().numpy().squeeze()
+    
+    if binarize:
+        output = np.where(output > threshold, 1., 0.)
+        
+    return image, output
+
+
+def enlarge_array(output):
+    df = pd.DataFrame(np.reshape(output, (14,14)))
+    df = pd.DataFrame(np.repeat(df.values, 16, axis=0))
+    df = pd.DataFrame(np.repeat(df.values, 16, axis=1))
+    output = df.to_numpy()
+    
+    return output
+
+def visualize_output(image, output, binarize, threshold):
+    
+    image, output = preprocess(image, output, binarize, threshold) 
+    output = enlarge_array(output)
+    output_mask = Image.fromarray(output * 255)
+
+    fig = plt.figure(figsize = (6,6))
+    plt.axis('off')
+    plt.imshow(image)
+    if binarize:
+        plt.imshow(output_mask, alpha=.67)
+    else:
+        plt.imshow(output_mask, alpha=.8)
+    fig.tight_layout(pad=0)
+    fig.canvas.draw()
+    data = np.frombuffer(fig.canvas.tostring_rgb(), dtype=np.uint8)
+    data = data.reshape(fig.canvas.get_width_height()[::-1] + (3,))
+
+    return data

model.py

@@ -0,0 +1,40 @@
+import torch
+import torch.nn as nn
+
+
+class Model(nn.Module):
+    def __init__(self, vit, roberta, tokenizer, device):
+        super().__init__()
+        self.bertmap = nn.Conv1d(768, 768, 1)
+        self.vitmap = nn.Conv1d(768, 768, 1)
+        self.conv1d = nn.Conv1d(1, 1, 1)
+        self.add_module("vit", vit)
+        self.add_module("roberta", roberta)
+        self.tokenizer = tokenizer
+        self.conv1d.weight = torch.nn.Parameter(torch.tensor([[[1.]]]))
+        self.conv1d.bias = torch.nn.Parameter(torch.tensor([0.]))
+        self.device = device
+        
+    def forward(self, image, cats):
+        vit_out = self.vit(image)
+        vit_out = vit_out[:,1:vit_out.shape[1],:]
+        vit_out = torch.transpose(vit_out, 2,1)
+        vit_out = self.vitmap(vit_out)
+        vit_out = torch.transpose(vit_out, 2,1)
+        token_out = self.tokenizer.encode_plus(
+                        cats,
+                        padding=True,
+                        add_special_tokens=True,
+                        return_token_type_ids=True,
+                        return_tensors='pt'
+                    ).to(self.device)
+        bert_out = self.roberta(**token_out)
+        hidden_state = bert_out.last_hidden_state
+        hidden_state = torch.transpose(hidden_state, 2,1)
+        hidden_state = self.bertmap(hidden_state)
+        hidden_state = torch.transpose(hidden_state, 2,1)
+        pooled_bert_out = hidden_state[:, 0]
+        pooled_bert_out = torch.unsqueeze(pooled_bert_out, dim=2)
+        out = torch.matmul(vit_out, pooled_bert_out)
+        out = torch.transpose(out, 2,1)
+        return torch.squeeze(self.conv1d(out), dim=1)

requirements.txt

@@ -0,0 +1,10 @@
+# pip install -r requirements.txt
+
+numpy>=1.18.5
+torch>=1.7.0
+torchvision>=0.8.1
+git+https://github.com/huggingface/transformers.git
+opencv-python
+pandas
+matplotlib
+timm