|
import streamlit as st |
|
import yaml |
|
import torch |
|
import torch.nn.functional as F |
|
|
|
from transformers import DetrImageProcessor, DetrForObjectDetection |
|
|
|
from lib.IRRA.tokenizer import tokenize, SimpleTokenizer |
|
from lib.IRRA.image import prepare_images |
|
from lib.IRRA.model.build import build_model, IRRA |
|
from PIL import Image |
|
from pathlib import Path |
|
|
|
from easydict import EasyDict |
|
|
|
|
|
@st.cache_resource |
|
def get_model(): |
|
args = yaml.load(open('model/configs.yaml'), Loader=yaml.FullLoader) |
|
args = EasyDict(args) |
|
args['training'] = False |
|
|
|
model = build_model(args) |
|
|
|
return model |
|
|
|
|
|
@st.cache_resource |
|
def get_detr(): |
|
processor = DetrImageProcessor.from_pretrained( |
|
"facebook/detr-resnet-50", revision="no_timm") |
|
|
|
model = DetrForObjectDetection.from_pretrained( |
|
"facebook/detr-resnet-50", revision="no_timm") |
|
|
|
return model, processor |
|
|
|
|
|
def segment_images(model, processor, images: list[str]): |
|
segments = [] |
|
id = 0 |
|
|
|
p = Path('segments') |
|
p.mkdir(exist_ok=True) |
|
|
|
for image in images: |
|
image = Image.open(image) |
|
|
|
inputs = processor(images=image, return_tensors="pt") |
|
outputs = model(**inputs) |
|
|
|
target_sizes = torch.tensor([image.size[::-1]]) |
|
results = processor.post_process_object_detection( |
|
outputs, target_sizes=target_sizes, threshold=0.9)[0] |
|
|
|
for score, label, box in zip(results["scores"], results["labels"], results["boxes"]): |
|
box = [round(i, 2) for i in box.tolist()] |
|
label = model.config.id2label[label.item()] |
|
|
|
if box[2] - box[0] > 70 and box[3] - box[1] > 70: |
|
if label == 'person': |
|
file = p / f'img_{id}.jpg' |
|
image.crop(box).save(file) |
|
segments.append(file.as_posix()) |
|
|
|
id += 1 |
|
|
|
return segments |
|
|
|
|
|
def get_similarities(text: str, images: list[str], model: IRRA) -> torch.Tensor: |
|
tokenizer = SimpleTokenizer() |
|
|
|
txt = tokenize(text, tokenizer) |
|
imgs = prepare_images(images) |
|
|
|
image_feats = model.encode_image(imgs) |
|
text_feats = model.encode_text(txt.unsqueeze(0)) |
|
|
|
image_feats = F.normalize(image_feats, p=2, dim=1) |
|
text_feats = F.normalize(text_feats, p=2, dim=1) |
|
|
|
return text_feats @ image_feats.t() |
|
|
