| import gradio as gr |
| import numpy as np |
| from PIL import Image |
| from transformers import CLIPProcessor, CLIPModel, DetrFeatureExtractor, DetrForObjectDetection, AutoFeatureExtractor, AutoModelForObjectDetection |
| import torch |
|
|
| feature_extractor = AutoFeatureExtractor.from_pretrained("nielsr/detr-resnet-50") |
| dmodel = AutoModelForObjectDetection.from_pretrained("nielsr/detr-resnet-50") |
|
|
| model = CLIPModel.from_pretrained("openai/clip-vit-base-patch32") |
| processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch32") |
|
|
| i1 = gr.inputs.Image(type="pil", label="Input image") |
| i2 = gr.inputs.Textbox(label="Input text") |
| i3 = gr.inputs.Number(default=0.96, label="Threshold percentage score") |
| o1 = gr.outputs.Image(type="pil", label="Cropped part") |
| o2 = gr.outputs.Textbox(label="Similarity score") |
|
|
| def extract_image(image, text, prob, num=1): |
| |
| inputs = feature_extractor(images=image, return_tensors="pt") |
| outputs = dmodel(**inputs) |
| |
| |
| logits = outputs.logits |
| bboxes = outputs.pred_boxes |
| probas = outputs.logits.softmax(-1)[0, :, :-1] |
| |
| keep = probas.max(-1).values > prob |
| outs = feature_extractor.post_process(outputs, torch.tensor(image.size[::-1]).unsqueeze(0)) |
| bboxes_scaled = outs[0]['boxes'][keep].detach().numpy() |
| labels = outs[0]['labels'][keep].detach().numpy() |
| scores = outs[0]['scores'][keep].detach().numpy() |
| |
| images_list = [] |
| for i,j in enumerate(bboxes_scaled): |
| |
| xmin = int(j[0]) |
| ymin = int(j[1]) |
| xmax = int(j[2]) |
| ymax = int(j[3]) |
| |
| im_arr = np.array(image) |
| roi = im_arr[ymin:ymax, xmin:xmax] |
| roi_im = Image.fromarray(roi) |
| |
| images_list.append(roi_im) |
| |
| inpu = processor(text = [text], images=images_list , return_tensors="pt", padding=True) |
| output = model(**inpu) |
| logits_per_image = output.logits_per_text |
| probs = logits_per_image.softmax(-1) |
| l_idx = np.argsort(probs[-1].detach().numpy())[::-1][0:num] |
| |
| final_ims = [] |
| for i,j in enumerate(images_list): |
| json_dict = {} |
| if i in l_idx: |
| json_dict['image'] = images_list[i] |
| json_dict['score'] = probs[-1].detach().numpy()[i] |
| |
| final_ims.append(json_dict) |
| |
| fi = sorted(final_ims, key=lambda item: item.get("score"), reverse=True) |
| return fi[0]['image'], fi[0]['score'] |
|
|
| title = "ClipnCrop" |
| description = "Extract sections of images from your image by using OpenAI's CLIP and Facebooks Detr implemented on HuggingFace Transformers" |
| examples=[['ex3.jpg', 'black bag', 0.96],['ex2.jpg', 'man in red dress', 0.85]] |
| article = "<p style='text-align: center'><a href='https://github.com/Vishnunkumar/clipcrop' target='_blank'>clipcrop</a></p>" |
| gr.Interface(fn=extract_image, inputs=[i1, i2, i3], outputs=[o1, o2], title=title, description=description, article=article, examples=examples, enable_queue=True).launch() |
|
|
