Created
September 8, 2023 06:53
-
-
Save khengyun/ca819304aab1ffe09174fa4f169200fc to your computer and use it in GitHub Desktop.
anylabeling to mask
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import json | |
import os | |
from PIL import Image, ImageDraw | |
############################################## | |
# - image_label: | |
# - 1.json | |
# - 1.jpg | |
# - 2.json | |
# - 2.jpg | |
# - <...>.json | |
############################################## | |
# Main folder from anylabeling | |
folder_path = "image_label" | |
# Create 'images' and 'masks' directories if they don't exist | |
images_dir = "images" | |
if not os.path.exists(images_dir): | |
os.makedirs(images_dir) | |
masks_dir = "masks" | |
if not os.path.exists(masks_dir): | |
os.makedirs(masks_dir) | |
# Loop through all JSON files in the folder | |
for filename in os.listdir(folder_path): | |
if filename.endswith(".json"): | |
json_file_path = os.path.join(folder_path, filename) | |
# Read data from the JSON file | |
with open(json_file_path, "r") as json_file: | |
data = json.load(json_file) | |
# Check if the image and mask paths exist | |
image_path = os.path.join(folder_path, data["imagePath"]) | |
if not os.path.exists(image_path): | |
print(f"Image not found for JSON file: {json_file_path}") | |
continue | |
# Copy the image to the 'images' folder | |
image = Image.open(image_path) | |
image.save(os.path.join(images_dir, os.path.basename(data["imagePath"]))) | |
# Create a mask image | |
mask = Image.new("RGB", image.size, (0, 0, 0)) | |
draw = ImageDraw.Draw(mask) | |
# Process all labels in the JSON | |
for shape in data.get("shapes", []): | |
if shape["label"] == "road": | |
color = (128, 0, 0) # Red for "road" | |
elif shape["label"] == "line": | |
color = (192, 128, 128) # Gray for "line" | |
# Convert points to the correct (x, y) format | |
points = [(point[0], point[1]) for point in shape["points"]] | |
draw.polygon(points, outline=color, fill=color) | |
mask.save(os.path.join(masks_dir, f"{os.path.splitext(filename)[0]}.png"), "PNG") | |
print("Images and masks saved successfully.") |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment