|
|
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 |
