File size: 3,128 Bytes
cb2428f |
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 |
import colorsys
import itertools
import os
from copy import deepcopy
from typing import Any, List, Literal
import requests
from modelscope.hub.utils.utils import get_cache_dir
from PIL import Image, ImageDraw, ImageFont
def _shuffle_colors(nums: List[Any]) -> List[Any]:
if len(nums) == 1:
return nums
mid = len(nums) // 2
left = nums[:mid]
right = nums[mid:]
left = _shuffle_colors(left)
right = _shuffle_colors(right)
new_nums = []
for x, y in zip(left, right):
new_nums += [x, y]
new_nums += left[len(right):] or right[len(left):]
return new_nums
def generate_colors():
vs_combinations = [(v, s) for v, s in itertools.product([0.7, 0.3, 1], [0.7, 0.3, 1])]
colors = [colorsys.hsv_to_rgb(i / 16, s, v) for v, s in vs_combinations for i in _shuffle_colors(list(range(16)))]
colors = [(int(r * 255), int(g * 255), int(b * 255)) for r, g, b in colors]
return _shuffle_colors(colors)
def download_file(url: str) -> str:
url = url.rstrip('/')
file_name = url.rsplit('/', 1)[-1]
cache_dir = os.path.join(get_cache_dir(), 'files')
os.makedirs(cache_dir, exist_ok=True)
req = requests.get(url)
file_path = os.path.join(cache_dir, file_name)
with open(file_path, 'wb') as f:
f.write(req.content)
return file_path
colors = generate_colors()
color_mapping = {}
def _calculate_brightness(image, region: List[int]):
cropped_image = image.crop(region)
grayscale_image = cropped_image.convert('L')
pixels = list(grayscale_image.getdata())
average_brightness = sum(pixels) / len(pixels)
return average_brightness
def draw_bbox(image: Image.Image,
ref: List[str],
bbox: List[List[int]],
norm_bbox: Literal['norm1000', 'none'] = 'norm1000'):
bbox = deepcopy(bbox)
font_path = 'https://modelscope.cn/models/Qwen/Qwen-VL-Chat/resolve/master/SimSun.ttf'
# norm bbox
for i, box in enumerate(bbox):
for i in range(len(box)):
box[i] = int(box[i])
if norm_bbox == 'norm1000':
box[0] = box[0] / 1000 * image.width
box[2] = box[2] / 1000 * image.width
box[1] = box[1] / 1000 * image.height
box[3] = box[3] / 1000 * image.height
draw = ImageDraw.Draw(image)
# draw bbox
assert len(ref) == len(bbox), f'len(refs): {len(ref)}, len(bboxes): {len(bbox)}'
for (left, top, right, bottom), box_ref in zip(bbox, ref):
if box_ref not in color_mapping:
color_mapping[box_ref] = colors[len(color_mapping) % len(colors)]
color = color_mapping[box_ref]
draw.rectangle([(left, top), (right, bottom)], outline=color, width=3)
# draw text
file_path = download_file(font_path)
font = ImageFont.truetype(file_path, 20)
for (left, top, _, _), box_ref in zip(bbox, ref):
brightness = _calculate_brightness(
image, [left, top, min(left + 100, image.width),
min(top + 20, image.height)])
draw.text((left, top), box_ref, fill='white' if brightness < 128 else 'black', font=font)
|