File size: 4,087 Bytes
1060621 |
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 |
import os
import json
import numpy as np
import cv2
import shutil
from tqdm import tqdm
def labelme2mask_single_img(img_path, labelme_json_path, class_info):
'''
Convert a single image's LabelMe annotation to a mask.
'''
img_bgr = cv2.imread(img_path)
img_mask = np.zeros(img_bgr.shape[:2], dtype=np.uint8) # Create an empty mask image (0 - background)
with open(labelme_json_path, 'r', encoding='utf-8') as f:
labelme = json.load(f)
for one_class in class_info: # Iterate over each class in class_info
for each in labelme['shapes']: # Iterate over all shapes in the annotation
if each['label'] == one_class['label']:
if one_class['type'] == 'polygon': # Handle polygon annotation
points = [np.array(each['points'], dtype=np.int32).reshape((-1, 1, 2))] # Ensure correct shape
img_mask = cv2.fillPoly(img_mask, points, color=one_class['color'])
elif one_class['type'] == 'line' or one_class['type'] == 'linestrip': # Handle line annotation
points = [np.array(each['points'], dtype=np.int32).reshape((-1, 1, 2))]
img_mask = cv2.polylines(img_mask, points, isClosed=False, color=one_class['color'],
thickness=one_class.get('thickness', 1))
elif one_class['type'] == 'circle': # Handle circle annotation
points = np.array(each['points'], dtype=np.int32)
center_x, center_y = points[0][0], points[0][1]
edge_x, edge_y = points[1][0], points[1][1]
radius = int(np.linalg.norm([center_x - edge_x, center_y - edge_y]))
img_mask = cv2.circle(img_mask, (center_x, center_y), radius, one_class['color'], -1)
else:
print('Unknown annotation type:', one_class['type'])
return img_mask
def convert_labelme_to_mask(Dataset_Path):
'''
Convert all LabelMe annotations in the dataset to mask images.
'''
# Dataset directories
img_dir = os.path.join(Dataset_Path, 'images')
ann_dir = os.path.join(Dataset_Path, 'labelme_jsons')
# Class information for mask generation
class_info = [
{'label': 'panicle', 'type': 'polygon', 'color': 1}
]
# Create target directories
images_target_dir = os.path.join(Dataset_Path, 'img_dir')
ann_target_dir = os.path.join(Dataset_Path, 'ann_dir')
# Create target directories if they do not exist
os.makedirs(images_target_dir, exist_ok=True)
os.makedirs(ann_target_dir, exist_ok=True)
# Process each image in the images directory
for img_name in tqdm(os.listdir(img_dir), desc="Converting images to masks"):
try:
img_path = os.path.join(img_dir, img_name)
labelme_json_path = os.path.join(ann_dir, f'{os.path.splitext(img_name)[0]}.json')
if os.path.exists(labelme_json_path):
# Convert LabelMe annotations to mask
img_mask = labelme2mask_single_img(img_path, labelme_json_path, class_info)
# Save the mask to the target directory
mask_path = os.path.join(ann_target_dir, f'{os.path.splitext(img_name)[0]}.png')
cv2.imwrite(mask_path, img_mask)
# Move the image to the target directory
shutil.move(img_path, os.path.join(images_target_dir, img_name))
else:
print(f"Annotation file missing for {img_name}")
except Exception as e:
print(f"Failed to convert {img_name}: {e}")
# Optionally remove the original directories if they are empty
shutil.rmtree(img_dir, ignore_errors=True)
shutil.rmtree(ann_dir, ignore_errors=True)
print("Conversion completed.")
if __name__ == '__main__':
Dataset_Path = 'CVRP' # Update this to the path of your dataset
convert_labelme_to_mask(Dataset_Path)
|