Spaces:
Runtime error
Runtime error
File size: 3,027 Bytes
8e5dadf |
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 95 96 97 98 99 100 |
import cv2
import numpy as np
import supervision as sv
from PIL import Image
MIN_AREA = 100
def load_image(image_path: str):
return Image.open(image_path).convert("RGB")
def draw_image(image_rgb, masks, xyxy, probs, labels):
box_annotator = sv.BoxCornerAnnotator()
label_annotator = sv.LabelAnnotator()
mask_annotator = sv.MaskAnnotator()
# Create class_id for each unique label
unique_labels = list(set(labels))
class_id_map = {label: idx for idx, label in enumerate(unique_labels)}
class_id = [class_id_map[label] for label in labels]
# Add class_id to the Detections object
detections = sv.Detections(
xyxy=xyxy,
mask=masks.astype(bool),
confidence=probs,
class_id=np.array(class_id),
)
annotated_image = box_annotator.annotate(scene=image_rgb.copy(), detections=detections)
annotated_image = label_annotator.annotate(scene=annotated_image, detections=detections, labels=labels)
annotated_image = mask_annotator.annotate(scene=annotated_image, detections=detections)
return annotated_image
def get_contours(mask):
if len(mask.shape) > 2:
mask = np.squeeze(mask, 0)
mask = mask.astype(np.uint8)
mask *= 255
contours, _ = cv2.findContours(mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
effContours = []
for c in contours:
area = cv2.contourArea(c)
if area > MIN_AREA:
effContours.append(c)
return effContours
def contour_to_points(contour):
pointsNum = len(contour)
contour = contour.reshape(pointsNum, -1).astype(np.float32)
points = [point.tolist() for point in contour]
return points
def generate_labelme_json(binary_masks, labels, image_size, image_path=None):
"""Generate a LabelMe format JSON file from binary mask tensor.
Args:
binary_masks: Binary mask tensor of shape [N, H, W].
labels: List of labels for each mask.
image_size: Tuple of (height, width) for the image size.
image_path: Path to the image file (optional).
Returns:
A dictionary representing the LabelMe JSON file.
"""
num_masks = binary_masks.shape[0]
binary_masks = binary_masks.numpy()
json_dict = {
"version": "4.5.6",
"imageHeight": image_size[0],
"imageWidth": image_size[1],
"imagePath": image_path,
"flags": {},
"shapes": [],
"imageData": None,
}
# Loop through the masks and add them to the JSON dictionary
for i in range(num_masks):
mask = binary_masks[i]
label = labels[i]
effContours = get_contours(mask)
for effContour in effContours:
points = contour_to_points(effContour)
shape_dict = {
"label": label,
"line_color": None,
"fill_color": None,
"points": points,
"shape_type": "polygon",
}
json_dict["shapes"].append(shape_dict)
return json_dict
|