78 lines
3.1 KiB
Python
78 lines
3.1 KiB
Python
import os
|
|
import json
|
|
import cv2
|
|
import numpy as np
|
|
from dataclasses import dataclass
|
|
import random
|
|
|
|
class CommonUtils:
|
|
@staticmethod
|
|
def creat_dirs(path):
|
|
"""
|
|
Ensure the given path exists. If it does not exist, create it using os.makedirs.
|
|
|
|
:param path: The directory path to check or create.
|
|
"""
|
|
try:
|
|
if not os.path.exists(path):
|
|
os.makedirs(path, exist_ok=True)
|
|
print(f"Path '{path}' did not exist and has been created.")
|
|
else:
|
|
print(f"Path '{path}' already exists.")
|
|
except Exception as e:
|
|
print(f"An error occurred while creating the path: {e}")
|
|
|
|
@staticmethod
|
|
def draw_masks_and_box(raw_image_path, mask_path, json_path, output_path):
|
|
CommonUtils.creat_dirs(output_path)
|
|
raw_image_name_list = os.listdir(raw_image_path)
|
|
raw_image_name_list.sort()
|
|
for raw_image_name in raw_image_name_list:
|
|
image_path = os.path.join(raw_image_path, raw_image_name)
|
|
image = cv2.imread(image_path)
|
|
if image is None:
|
|
raise FileNotFoundError("Image file not found.")
|
|
# load mask
|
|
mask_npy_path = os.path.join(mask_path, "mask_"+raw_image_name.split(".")[0]+".npy")
|
|
mask = np.load(mask_npy_path)
|
|
# color map
|
|
unique_ids = np.unique(mask)
|
|
colors = {uid: CommonUtils.random_color() for uid in unique_ids}
|
|
colors[0] = (0, 0, 0) # background color
|
|
|
|
# apply mask to image
|
|
colored_mask = np.zeros_like(image)
|
|
for uid in unique_ids:
|
|
colored_mask[mask == uid] = colors[uid]
|
|
alpha = 0.5 # 调整 alpha 值以改变透明度
|
|
output_image = cv2.addWeighted(image, 1 - alpha, colored_mask, alpha, 0)
|
|
|
|
|
|
file_path = os.path.join(json_path, "mask_"+raw_image_name.split(".")[0]+".json")
|
|
with open(file_path, 'r') as file:
|
|
json_data = json.load(file)
|
|
# Draw bounding boxes and labels
|
|
for obj_id, obj_item in json_data["labels"].items():
|
|
# Extract data from JSON
|
|
x1, y1, x2, y2 = obj_item["x1"], obj_item["y1"], obj_item["x2"], obj_item["y2"]
|
|
instance_id = obj_item["instance_id"]
|
|
class_name = obj_item["class_name"]
|
|
|
|
# Draw rectangle
|
|
cv2.rectangle(output_image, (x1, y1), (x2, y2), (0, 255, 0), 2)
|
|
|
|
# Put text
|
|
label = f"{instance_id}: {class_name}"
|
|
cv2.putText(output_image, label, (x1, y1 - 10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 2)
|
|
|
|
# Save the modified image
|
|
output_image_path = os.path.join(output_path, raw_image_name)
|
|
cv2.imwrite(output_image_path, output_image)
|
|
|
|
print(f"Annotated image saved as {output_image_path}")
|
|
|
|
@staticmethod
|
|
def random_color():
|
|
"""random color generator"""
|
|
return (random.randint(0, 255), random.randint(0, 255), random.randint(0, 255))
|