Update annotation_tool.py

annotation_tool.py

@@ -0,0 +1,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()