File size: 5,123 Bytes
032e687 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 | import argparse
import numpy as np
import random
import os
import torch
from PIL import Image
from huggingface_hub import hf_hub_download
from .ram.models import ram_plus
from .ram import get_transform
def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
best_ratio_diff = float('inf')
best_ratio = (1, 1)
area = width * height
for ratio in target_ratios:
target_aspect_ratio = ratio[0] / ratio[1]
ratio_diff = abs(aspect_ratio - target_aspect_ratio)
if ratio_diff < best_ratio_diff:
best_ratio_diff = ratio_diff
best_ratio = ratio
elif ratio_diff == best_ratio_diff:
if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
best_ratio = ratio
return best_ratio
def dynamic_preprocess(image, min_num=1, max_num=6, image_size=448, use_thumbnail=False):
orig_width, orig_height = image.size
aspect_ratio = orig_width / orig_height
# calculate the existing image aspect ratio
target_ratios = set(
(i, j) for n in range(min_num, max_num + 1) for i in range(1, n + 1) for j in range(1, n + 1) if
i * j <= max_num and i * j >= min_num)
target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
# find the closest aspect ratio to the target
target_aspect_ratio = find_closest_aspect_ratio(
aspect_ratio, target_ratios, orig_width, orig_height, image_size)
# calculate the target width and height
target_width = image_size * target_aspect_ratio[0]
target_height = image_size * target_aspect_ratio[1]
blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
# resize the image
resized_img = image.resize((target_width, target_height))
processed_images = []
for i in range(blocks):
box = (
(i % (target_width // image_size)) * image_size,
(i // (target_width // image_size)) * image_size,
((i % (target_width // image_size)) + 1) * image_size,
((i // (target_width // image_size)) + 1) * image_size
)
# split the image
split_img = resized_img.crop(box)
processed_images.append(split_img)
assert len(processed_images) == blocks
if use_thumbnail and len(processed_images) != 1:
thumbnail_img = image.resize((image_size, image_size))
processed_images.append(thumbnail_img)
return processed_images
def load_image(image_file, input_size=384, max_num=12, upscale=False):
image = Image.open(image_file)
if upscale:
image = image.resize((image.width * 2, image.height * 2), Image.BILINEAR)
images = dynamic_preprocess(image, image_size=input_size, use_thumbnail=True, max_num=max_num)
return images
class RAMPredictor(object):
def __init__(
self,
repo_id="xinyu1205/recognize-anything-plus-model",
checkpoint_file="ram_plus_swin_large_14m.pth",
image_size=384,
device='cuda',
):
super().__init__()
self.image_size = image_size
self.device = device
self.transform = get_transform(image_size=image_size)
if os.path.exists(checkpoint_file):
init_checkpoint = checkpoint_file
else:
init_checkpoint = hf_hub_download(repo_id=repo_id, filename=checkpoint_file)
self.model = ram_plus(pretrained=init_checkpoint, image_size=image_size, vit='swin_l', text_encoder_type="third_parts/recognize_anything/google-bert/bert-base-uncased")
self.model.eval()
self.model.to(device)
def run_on_image(
self,
image_file_path,
dynamic_resolution=False,
):
if dynamic_resolution:
images = load_image(image_file_path, input_size=self.image_size)
images = [self.transform(image) for image in images]
images = torch.stack(images).to(self.device)
else:
if not isinstance(image_file_path, str):
images = self.transform(image_file_path).unsqueeze(0).to(self.device)
else:
images = self.transform(Image.open(image_file_path)).unsqueeze(0).to(self.device)
res = self.model.generate_tag(images)
return res
def build_ram_predictor(override_ckpt_file="", device="cuda"):
repo_id="xinyu1205/recognize-anything-plus-model"
checkpoint_file="ram_plus_swin_large_14m.pth"
if os.path.exists(override_ckpt_file):
checkpoint_file = override_ckpt_file
ram_predictor = RAMPredictor(repo_id, checkpoint_file, image_size=384, device=device)
return ram_predictor
if __name__ == "__main__":
ram_predictor = build_ram_predictor(override_ckpt_file="xinyu1205/recognize-anything-plus-model/ram_plus_swin_large_14m.pth")
res = ram_predictor.run_on_image(image_file_path="sa_7963505.jpg", dynamic_resolution=True)
tag_list = []
for tag_string in res[0]:
tags = tag_string.split(' | ')
tag_list += tags
tags = list(set(tag_list))
print("Image Tags: ", tags)
|