| | from PIL import Image, ImageDraw, ImageColor |
| | import json |
| | import os |
| |
|
| | def dump_args_to_json(model_config, data_processor, model_args, data_args, training_args, output_dir): |
| | def is_json_serializable(v): |
| | try: |
| | json.dumps(v) |
| | return True |
| | except: |
| | return False |
| |
|
| | save_path = f"{output_dir}/args.json" |
| | if not os.path.exists(save_path): |
| | with open(save_path, "w") as f: |
| | json.dump({ |
| | "model_config": {k: v for k, v in model_config.__dict__.items() if is_json_serializable(v)}, |
| | "data_processor_config": {k: v for k, v in data_processor.__dict__.items() if is_json_serializable(v)}, |
| | "image_processor_config": {k: v for k, v in data_processor.image_processor.__dict__.items() if is_json_serializable(v)}, |
| | "model_args": {k: v for k, v in model_args.__dict__.items() if is_json_serializable(v)}, |
| | "data_args": {k: v for k, v in data_args.__dict__.items() if is_json_serializable(v)}, |
| | "training_args": {k: v for k, v in training_args.__dict__.items() if is_json_serializable(v)}, |
| | }, f, indent=4) |
| |
|
| | def draw_point(image: Image.Image, point: list, color=None): |
| | if isinstance(color, str): |
| | try: |
| | color = ImageColor.getrgb(color) |
| | color = color + (128,) |
| | except ValueError: |
| | color = (255, 0, 0, 128) |
| | else: |
| | color = (255, 0, 0, 128) |
| |
|
| | overlay = Image.new('RGBA', image.size, (255, 255, 255, 0)) |
| | overlay_draw = ImageDraw.Draw(overlay) |
| | radius = 14 |
| | x, y = point |
| |
|
| | overlay_draw.rectangle( |
| | [x - radius, y - radius, x + radius, y + radius], |
| | fill=color |
| | ) |
| | |
| | center_radius = radius * 0.1 |
| | overlay_draw.ellipse( |
| | [(x - center_radius, y - center_radius), |
| | (x + center_radius, y + center_radius)], |
| | fill=(0, 255, 0, 255) |
| | ) |
| |
|
| | image = image.convert('RGBA') |
| | combined = Image.alpha_composite(image, overlay) |
| |
|
| | return combined.convert('RGB') |
| |
|
| | def draw_bbox(image: Image.Image, bbox: list, color=None): |
| | """bbox is in the format of [x1, y1, x2, y2]""" |
| | if isinstance(color, str): |
| | try: |
| | color = ImageColor.getrgb(color) |
| | color = color + (128,) |
| | except ValueError: |
| | color = (255, 0, 0, 128) |
| | else: |
| | color = (255, 0, 0, 128) |
| | |
| | overlay = Image.new('RGBA', image.size, (255, 255, 255, 0)) |
| | overlay_draw = ImageDraw.Draw(overlay) |
| | overlay_draw.rectangle(bbox, fill=color) |
| | return Image.alpha_composite(image, overlay).convert('RGB') |
| |
|
| | def do_boxes_overlap(box1, box2): |
| | """ |
| | Check if two boxes overlap. |
| | |
| | Each box is represented as a tuple: (x1, y1, x2, y2) |
| | Where (x1, y1) is the top-left and (x2, y2) is the bottom-right corner. |
| | """ |
| | |
| | x1_min, y1_min, x1_max, y1_max = box1 |
| | x2_min, y2_min, x2_max, y2_max = box2 |
| |
|
| | |
| | if x1_max < x2_min or x2_max < x1_min: |
| | return False |
| | if y1_max < y2_min or y2_max < y1_min: |
| | return False |
| |
|
| | return True |