Spaces:
Build error
Build error
File size: 5,055 Bytes
3b82ac6 |
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 153 154 |
import os, ast
from glob import glob
from PIL import ImageFont, ImageDraw, Image
def process_txtfile(filename):
"""
Read txt annotations files (designed for YOLO xywh format)
Parameters:
filename(str): path of the txt annotation file.
Returns:
segments: list of bboxes in format xmin, ymin, xmax, ymax (as image ratio)
confs: list of confidences of the bboxes object detection
"""
segments = []
confs = []
with open(filename, 'r') as file:
for line in file:
# print(line)
line = line.strip().split(' ')
cls = int(line[0])
conf = line[5]
x, y, w, h = map(float, line[1:5])
x_min = x - (w / 2)
y_min = y - (h / 2)
x_max = x + (w / 2)
y_max = y + (h / 2)
segment = [x_min, y_min, x_max, y_max]
segments.append(segment)
confs.append(conf)
return segments, confs
def process_jsonfile(filename):
"""
Read json annotations files (designed for mmdetect dict format)
Parameters:
filename(str): path of the json annotation file.
Returns:
segments: bboxes in format xmin, ymin, xmax, ymax (as px coordinates)
confs: list of confidences of the bboxes object detection
"""
with open(filename, 'r') as file:
line = file.readline().strip()
dic = ast.literal_eval(line)
segments = dic['bboxes']
confs = dic['scores']
# labels = dic['labels']
return segments, confs
def lerp_color(color1, color2, t):
"""
Linearly interpolate between two RGB colors.
Parameters:
color1 (tuple): RGB tuple of the first color.
color2 (tuple): RGB tuple of the second color.
t (float): Interpolation factor between 0 and 1.
Returns:
tuple: Interpolated RGB color tuple.
"""
r = int(color1[0] + (color2[0] - color1[0]) * t)
g = int(color1[1] + (color2[1] - color1[1]) * t)
b = int(color1[2] + (color2[2] - color1[2]) * t)
return r, g, b
def generate_color_palette(start_color, end_color, steps):
"""
Generate an RGB color palette between two colors.
Parameters:
start_color (tuple): RGB tuple of the starting color.
end_color (tuple): RGB tuple of the ending color.
steps (int): Number of steps between the two colors.
Returns:
list: List of RGB tuples
"""
palette = []
for i in range(steps):
t = i / (steps - 1) # interpolation factor
color = lerp_color(start_color, end_color, t)
palette.append(color)
return palette
def draw_bbox(model_name, results_folder="./inference/results/", image_path="inptest.jpg"):
"""
Draw bounding boxes from mmdetect or yolo formats
"""
# annotations style
txt_color=(255, 255, 255)
yellow=(255, 255, 128)
black = (0, 0, 0)
steps = 11 # Step : 5%
# (255, 0, 0) # Red
# (0, 0, 255) # Blue
palette = generate_color_palette((255, 0, 0), (0, 0, 255), steps)
lw = 9
font = ImageFont.truetype(font="Pillow/Tests/fonts/FreeMono.ttf", size=48)
im = Image.open(image_path)
width, height = im.size
imdraw = ImageDraw.Draw(im)
exps = sorted(glob(f"inference/results/{model_name}_inference/*", recursive = True))
# print(exps)
if model_name[:4] == "yolo":
annot_file = glob(f"{exps[-1]}/labels/" + "*.txt")[0]
segments, confs = process_txtfile(annot_file)
else:
annot_file = glob(f"{exps[1]}/{image_path[:-4]}.json")[0]
segments, confs = process_jsonfile(annot_file)
# print("Result bboxes : " + annot_file)
for conf, box in zip(confs, segments):
conf_r = round(float(conf), 3) # round conf
if conf_r >= 0.5: # 0.5 threshold
bbox_c = palette[1] #
if conf_r <= 1.0: bbox_c = palette[-1]
if conf_r < 0.95: bbox_c = palette[-2]
if conf_r < 0.90: bbox_c = palette[-3]
if conf_r < 0.85: bbox_c = palette[-4]
if conf_r < 0.80: bbox_c = palette[-5]
if conf_r < 0.75: bbox_c = palette[-6]
if conf_r < 0.70: bbox_c = palette[-7]
if conf_r < 0.65: bbox_c = palette[-8]
if conf_r < 0.60: bbox_c = palette[-9]
if conf_r < 0.55: bbox_c = palette[-10]
if model_name[:4] == "yolo":
box = [box[0]*width, box[1]*height, box[2]*width, box[3]*height]
imdraw.rectangle(box, width=lw, outline=bbox_c) # box
# label
w, h = font.getbbox(str(conf_r))[2:4] # text w, h
imdraw.rectangle([box[0], box[1]-h, box[0]+w+1, box[1]+1], width=3, fill = black) # box
imdraw.text([box[0], box[1]-h], str(conf_r), fill=yellow, font=font)
im.save(f"{results_folder}{model_name}_inference/clean.jpg")
# count
count = len([i for i in confs if float(i) > 0.5])
return im, count
|