File size: 5,910 Bytes
6c73b1f |
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 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 |
"""
Simple tool for annotating tumors in medical images
"""
import os
import cv2
import numpy as np
import argparse
import json
from glob import glob
class TumorAnnotator:
def __init__(self, images_dir, annotations_dir):
self.images_dir = images_dir
self.annotations_dir = annotations_dir
self.image_files = sorted(glob(os.path.join(images_dir, "*.png")))
self.current_idx = 0
self.drawing = False
self.roi_pt1 = None
self.roi_pt2 = None
self.annotations = {}
# Create annotations directory if it doesn't exist
os.makedirs(annotations_dir, exist_ok=True)
# Load existing annotations if available
self.load_annotations()
def load_annotations(self):
"""Load existing annotations"""
for image_file in self.image_files:
image_name = os.path.basename(image_file)
ann_file = os.path.join(self.annotations_dir, f"{os.path.splitext(image_name)[0]}.json")
if os.path.exists(ann_file):
with open(ann_file, 'r') as f:
self.annotations[image_name] = json.load(f)
def save_annotation(self, image_name):
"""Save annotation for current image"""
if image_name not in self.annotations or not self.annotations[image_name]:
return
ann_file = os.path.join(self.annotations_dir, f"{os.path.splitext(image_name)[0]}.json")
with open(ann_file, 'w') as f:
json.dump(self.annotations[image_name], f)
def mouse_callback(self, event, x, y, flags, param):
"""Mouse callback for drawing bounding boxes"""
if event == cv2.EVENT_LBUTTONDOWN:
self.drawing = True
self.roi_pt1 = (x, y)
self.roi_pt2 = (x, y)
elif event == cv2.EVENT_MOUSEMOVE:
if self.drawing:
self.roi_pt2 = (x, y)
elif event == cv2.EVENT_LBUTTONUP:
self.drawing = False
self.roi_pt2 = (x, y)
# Add bounding box to annotations
image_name = os.path.basename(self.image_files[self.current_idx])
if image_name not in self.annotations:
self.annotations[image_name] = []
x1, y1 = min(self.roi_pt1[0], self.roi_pt2[0]), min(self.roi_pt1[1], self.roi_pt2[1])
x2, y2 = max(self.roi_pt1[0], self.roi_pt2[0]), max(self.roi_pt1[1], self.roi_pt2[1])
# Store as [x, y, width, height]
bbox = [x1, y1, x2 - x1, y2 - y1]
self.annotations[image_name].append({
"bbox": bbox,
"label": "tumor"
})
def run(self):
"""Run the annotation tool"""
cv2.namedWindow("Tumor Annotator")
cv2.setMouseCallback("Tumor Annotator", self.mouse_callback)
while True:
if self.current_idx >= len(self.image_files) or self.current_idx < 0:
break
image_path = self.image_files[self.current_idx]
image_name = os.path.basename(image_path)
image = cv2.imread(image_path)
display_image = image.copy()
# Draw existing annotations
if image_name in self.annotations:
for ann in self.annotations[image_name]:
bbox = ann["bbox"]
cv2.rectangle(
display_image,
(bbox[0], bbox[1]),
(bbox[0] + bbox[2], bbox[1] + bbox[3]),
(0, 255, 0),
2
)
# Draw current selection
if self.drawing:
cv2.rectangle(
display_image,
self.roi_pt1,
self.roi_pt2,
(0, 0, 255),
2
)
# Display image with annotations
cv2.imshow("Tumor Annotator", display_image)
# Display instructions
print("\nImage:", image_name)
print("Controls:")
print(" Left-click and drag: Draw bounding box")
print(" n: Next image")
print(" p: Previous image")
print(" d: Delete last annotation")
print(" s: Save annotations")
print(" q: Quit")
key = cv2.waitKey(1) & 0xFF
if key == ord('n'): # Next image
self.save_annotation(image_name)
self.current_idx = min(self.current_idx + 1, len(self.image_files) - 1)
elif key == ord('p'): # Previous image
self.save_annotation(image_name)
self.current_idx = max(self.current_idx - 1, 0)
elif key == ord('d'): # Delete last annotation
if image_name in self.annotations and self.annotations[image_name]:
self.annotations[image_name].pop()
elif key == ord('s'): # Save annotations
self.save_annotation(image_name)
print("Annotations saved!")
elif key == ord('q'): # Quit
self.save_annotation(image_name)
break
cv2.destroyAllWindows()
if __name__ == "__main__":
parser = argparse.ArgumentParser(description='Tumor annotation tool')
parser.add_argument('--images', type=str, required=True, help='Directory containing images')
parser.add_argument('--annotations', type=str, default='./annotations', help='Directory to save annotations')
args = parser.parse_args()
annotator = TumorAnnotator(args.images, args.annotations)
annotator.run() |