Spaces:
Build error
Build error
| 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 | |