app.py

import gradio as gr

import torch
import io
from PIL import Image
from transformers import (
    AutoImageProcessor,
    AutoTokenizer,
    AutoModelForCausalLM,
)
import numpy as np
model_root = "qihoo360/fg-clip-base"

model = AutoModelForCausalLM.from_pretrained(model_root,trust_remote_code=True)
device = model.device
tokenizer = AutoTokenizer.from_pretrained(model_root)
image_processor = AutoImageProcessor.from_pretrained(model_root)

import math
import matplotlib
matplotlib.use('Agg') 
import matplotlib.pyplot as plt

def postprocess_result(probs, labels):
    pro_output = {labels[i]: probs[i] for i in range(len(labels))}

    return pro_output


def Retrieval(image, candidate_labels):
    """
    Takes an image and a comma-separated string of candidate labels,
    and returns the classification scores.
    """
    image_size=224
    image = image.convert("RGB")
    image = image.resize((image_size,image_size))
    image_input = image_processor.preprocess(image, return_tensors='pt')['pixel_values'].to(device)
    walk_short_pos = True

    caption_input = torch.tensor(tokenizer(candidate_labels, max_length=77, padding="max_length", truncation=True).input_ids, dtype=torch.long, device=device)

    with torch.no_grad():
        image_feature = model.get_image_features(image_input)
        text_feature = model.get_text_features(caption_input,walk_short_pos=walk_short_pos)
        image_feature = image_feature / image_feature.norm(p=2, dim=-1, keepdim=True)
        text_feature = text_feature / text_feature.norm(p=2, dim=-1, keepdim=True)

        logits_per_image = image_feature @ text_feature.T
        logits_per_image = model.logit_scale.exp() * logits_per_image
        probs = logits_per_image.softmax(dim=1) 
    results = probs[0].tolist()
    return results


def Get_Densefeature(image, candidate_labels):
    """
    Takes an image and a comma-separated string of candidate labels,
    and returns the classification scores.
    """
    candidate_labels = [label.lstrip(" ") for label in candidate_labels.split(",") if label !=""]
    # print(candidate_labels)

    image_size=224
    image = image.convert("RGB")
    image = image.resize((image_size,image_size))
    image_input = image_processor.preprocess(image, return_tensors='pt')['pixel_values'].to(device)

    with torch.no_grad():
        dense_image_feature = model.get_image_dense_features(image_input)
        captions = candidate_labels
        caption_input = torch.tensor(tokenizer(captions, max_length=77, padding="max_length", truncation=True).input_ids, dtype=torch.long, device=device)
        text_feature = model.get_text_features(caption_input,walk_short_pos=True)
        text_feature = text_feature / text_feature.norm(p=2, dim=-1, keepdim=True)
        dense_image_feature = dense_image_feature / dense_image_feature.norm(p=2, dim=-1, keepdim=True)

    similarity = dense_image_feature.squeeze() @ text_feature.squeeze().T
    similarity = similarity.cpu().numpy()
    patch_size = int(math.sqrt(similarity.shape[0]))


    original_shape = (patch_size, patch_size)
    show_image = similarity.reshape(original_shape) 
    # normalized = (show_image - show_image.min()) / (show_image.max() - show_image.min())

    # def viridis_colormap(x):

    #     r = np.clip(1.1746 * x - 0.1776, 0, 1)
    #     g = np.clip(2.0 * x - 0.7, 0, 1)
    #     b = np.clip(-2.0 * x + 1.7, 0, 1)
    #     return np.stack([r, g, b], axis=-1)

    # color_mapped = viridis_colormap(normalized)

    # color_mapped_uint8 = (color_mapped * 255).astype(np.uint8)

    # pil_img = Image.fromarray(color_mapped_uint8)
    # pil_img = pil_img.resize((512,512))

    fig = plt.figure(figsize=(6, 6))
    plt.imshow(show_image)
    plt.title('similarity Visualization')
    plt.axis('off')  

    buf = io.BytesIO()
    plt.savefig(buf, format='png')
    buf.seek(0)  
    plt.close(fig)

    pil_img = Image.open(buf)
    # buf.close()
    return  pil_img




def infer(image, candidate_labels):
    candidate_labels = [label.lstrip(" ") for label in candidate_labels.split(",") if label !=""]
    fg_probs = Retrieval(image, candidate_labels)
    return postprocess_result(fg_probs,candidate_labels)


with gr.Blocks() as demo:
    gr.Markdown("# FG-CLIP Retrieval")
    gr.Markdown(
        "This app uses the FG-CLIP model (qihoo360/fg-clip-base) for retrieval on CPU :"
    )
    with gr.Row():
        with gr.Column():
            image_input = gr.Image(type="pil")
            text_input = gr.Textbox(label="Input a list of labels (comma seperated)")
            run_button = gr.Button("Run Retrieval", visible=True)
            dfs_button = gr.Button("Run Densefeature", visible=True)
        with gr.Column():
            fg_output = gr.Label(label="FG-CLIP Output", num_top_classes=11)
            dfs_output = gr.Image(label="Similarity Visualization", type="pil")
            
    examples = [
        # ["./baklava.jpg", "dessert on a plate, a serving of baklava, a plate and spoon"],
        # ["./dog.jpg", "A light brown wood stool, A bucket with a body made of dark brown plastic, A black velvet back cover for a cellular telephone, A green ball with a perforated pattern, A light blue plastic helmet made of plastic, A grey slipper made of wool, A newspaper with white and black perforated printed on a paper texture, A blue dog with a white colored head, A yellow sponge with a dark green rough surface, A book with white, dark orange and brown pages made of paper, A black ceramic scarf with a body made of fabric."],
        ["./Landscape.jpg", "red grass, yellow grass, green grass"],
        ["./cat.jpg", "two sleeping cats, two cats playing, three cats laying down"],
        ["./cat_dfclor.jpg", "white cat,"],
    ]
    gr.Examples(
        examples=examples,
        inputs=[image_input, text_input],
        # outputs=fg_output,
        # fn=infer,
    )
    run_button.click(fn=infer, inputs=[image_input, text_input], outputs=fg_output)
    dfs_button.click(fn=Get_Densefeature, inputs=[image_input, text_input], outputs=dfs_output)
demo.launch()