123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129 |
- import os
- from functools import reduce
- import cv2
- import numpy as np
- from PIL import Image, ImageFont, ImageDraw
- import sys
- sys.path.append(os.path.dirname(os.path.abspath(__file__)) + "/../")
- def pil_resize(image_np, height, width):
- image_pil = Image.fromarray(cv2.cvtColor(image_np, cv2.COLOR_BGR2RGB))
- image_pil = image_pil.resize((int(width), int(height)), Image.BICUBIC)
- image_np = cv2.cvtColor(np.asarray(image_pil), cv2.COLOR_RGB2BGR)
- return image_np
- def get_best_predict_size(image_np, times=8, min_size=128, max_size=400):
- sizes = []
- for i in range(int(min_size/times), 100):
- if i*times <= max_size:
- sizes.append(i*times)
- sizes.sort(key=lambda x: x, reverse=True)
- min_len = 10000
- best_height = sizes[0]
- for height in sizes:
- if abs(image_np.shape[0] - height) < min_len:
- min_len = abs(image_np.shape[0] - height)
- best_height = height
- min_len = 10000
- best_width = sizes[0]
- for width in sizes:
- if abs(image_np.shape[1] - width) < min_len:
- min_len = abs(image_np.shape[1] - width)
- best_width = width
- if best_height > best_width:
- best_width = best_height
- else:
- best_height = best_width
- return best_height, best_width
- def letterbox_image(image, size):
- """resize image with unchanged aspect ratio using padding"""
- iw, ih = image.size
- w, h = size
- scale = min(w/iw, h/ih)
- nw = int(iw*scale)
- nh = int(ih*scale)
- image = image.resize((nw,nh), Image.BICUBIC)
- new_image = Image.new('RGB', size, (128,128,128))
- new_image.paste(image, ((w-nw)//2, (h-nh)//2))
- return new_image
- def compose(*funcs):
- """Compose arbitrarily many functions, evaluated left to right.
- Reference: https://mathieularose.com/function-composition-in-python/
- """
- if funcs:
- return reduce(lambda f, g: lambda *a, **kw: g(f(*a, **kw)), funcs)
- else:
- raise ValueError('Composition of empty sequence not supported.')
- def draw_boxes(image, out_boxes, out_classes, out_scores, class_names, colors):
- font = ImageFont.truetype(font=os.path.abspath(os.path.dirname(__file__))+'/font/FiraMono-Medium.otf',
- size=np.floor(3e-2 * image.size[1] + 0.5).astype('int32'))
- thickness = (image.size[0] + image.size[1]) // 300
- box_list = []
- for i, c in reversed(list(enumerate(out_classes))):
- predicted_class = class_names[c]
- box = out_boxes[i]
- score = out_scores[i]
- label = '{} {:.2f}'.format(predicted_class, score)
- draw = ImageDraw.Draw(image)
- label_size = draw.textsize(label, font)
- top, left, bottom, right = box
- top = max(0, np.floor(top + 0.5).astype('int32'))
- left = max(0, np.floor(left + 0.5).astype('int32'))
- bottom = min(image.size[1], np.floor(bottom + 0.5).astype('int32'))
- right = min(image.size[0], np.floor(right + 0.5).astype('int32'))
- # print(label, (left, top), (right, bottom))
- box_list.append([(left, top), (right, bottom)])
- if top - label_size[1] >= 0:
- text_origin = np.array([left, top - label_size[1]])
- else:
- text_origin = np.array([left, top + 1])
- # My kingdom for a good redistributable image drawing library.
- for i in range(thickness):
- draw.rectangle(
- [left + i, top + i, right - i, bottom - i],
- outline=colors[c])
- draw.rectangle(
- [tuple(text_origin), tuple(text_origin + label_size)],
- fill=colors[c])
- draw.text(text_origin, label, fill=(0, 0, 0), font=font)
- del draw
- return image
- def adjust_boxes(image, boxes, threshold=10):
- new_boxes = []
- for box in boxes:
- w, h = image.size
- top, left, bottom, right = box
- top = max(0, np.floor(top + 0.5).astype('int32'))
- left = max(0, np.floor(left + 0.5).astype('int32'))
- bottom = min(h, np.floor(bottom + 0.5).astype('int32'))
- right = min(w, np.floor(right + 0.5).astype('int32'))
- # 把框加大,有时候圈不全
- top = 0 if top - threshold < 0 else top - threshold
- bottom = h if bottom + threshold > h else bottom + threshold
- left = 0 if left - threshold < 0 else left - threshold
- right = w if right + threshold > w else right + threshold
- new_boxes.append([(left, top), (right, bottom)])
- return new_boxes
|