Files
LPR-OCR/annotate_project.py
AtHeartEngineer 85ac8e08d2 init for sharing
2025-07-30 13:44:13 -04:00

340 lines
13 KiB
Python

#!/usr/bin/env python3
"""
Interactive license plate annotation tool for projects.
Click and drag to select license plate regions for training the detection algorithm.
"""
import cv2
import numpy as np
import json
import os
import argparse
from pathlib import Path
class LicensePlateAnnotator:
def __init__(self, image_path, project_dir, load_existing=True):
self.image_path = image_path
self.project_dir = Path(project_dir)
self.image = cv2.imread(image_path)
self.original = self.image.copy()
self.annotations = []
self.current_rect = None
self.drawing = False
self.start_point = None
self.scale_factor = 1.0
# Scale image for display if too large
h, w = self.image.shape[:2]
if w > 1200 or h > 800:
self.scale_factor = min(1200/w, 800/h)
new_w = int(w * self.scale_factor)
new_h = int(h * self.scale_factor)
self.display_image = cv2.resize(self.image, (new_w, new_h))
else:
self.display_image = self.image.copy()
self.window_name = f"Annotate License Plates - {Path(image_path).name}"
# Load existing annotations if requested
if load_existing:
self.load_existing_annotations()
def load_existing_annotations(self):
"""Load existing annotations for this image."""
image_name = Path(self.image_path).stem
annotation_file = self.project_dir / 'annotations' / f"{image_name}.json"
if annotation_file.exists():
try:
with open(annotation_file, 'r') as f:
data = json.load(f)
self.annotations = data.get('annotations', [])
print(f"Loaded {len(self.annotations)} existing annotations")
except Exception as e:
print(f"Warning: Could not load existing annotations: {e}")
self.annotations = []
def mouse_callback(self, event, x, y, flags, param):
orig_x = int(x / self.scale_factor)
orig_y = int(y / self.scale_factor)
if event == cv2.EVENT_LBUTTONDOWN:
self.drawing = True
self.start_point = (orig_x, orig_y)
elif event == cv2.EVENT_MOUSEMOVE:
if self.drawing:
temp_display = self.display_image.copy()
display_start = (int(self.start_point[0] * self.scale_factor),
int(self.start_point[1] * self.scale_factor))
cv2.rectangle(temp_display, display_start, (x, y), (0, 255, 0), 2)
cv2.imshow(self.window_name, temp_display)
elif event == cv2.EVENT_LBUTTONUP:
if self.drawing:
self.drawing = False
end_point = (orig_x, orig_y)
x1, y1 = self.start_point
x2, y2 = end_point
if abs(x2 - x1) > 10 and abs(y2 - y1) > 5:
annotation = {
'top_left': (min(x1, x2), min(y1, y2)),
'bottom_right': (max(x1, x2), max(y1, y2)),
'width': abs(x2 - x1),
'height': abs(y2 - y1),
'aspect_ratio': abs(x2 - x1) / abs(y2 - y1)
}
self.annotations.append(annotation)
print(f"Added annotation: {annotation['width']}x{annotation['height']}, AR: {annotation['aspect_ratio']:.2f}")
self.update_display()
def update_display(self):
self.display_image = cv2.resize(self.original,
(int(self.original.shape[1] * self.scale_factor),
int(self.original.shape[0] * self.scale_factor)))
for i, annotation in enumerate(self.annotations):
tl = annotation['top_left']
br = annotation['bottom_right']
display_tl = (int(tl[0] * self.scale_factor), int(tl[1] * self.scale_factor))
display_br = (int(br[0] * self.scale_factor), int(br[1] * self.scale_factor))
cv2.rectangle(self.display_image, display_tl, display_br, (0, 255, 0), 2)
cv2.putText(self.display_image, f"LP{i+1}",
(display_tl[0], display_tl[1] - 10),
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2)
cv2.imshow(self.window_name, self.display_image)
def run(self):
cv2.namedWindow(self.window_name, cv2.WINDOW_NORMAL)
cv2.setMouseCallback(self.window_name, self.mouse_callback)
print(f"\nAnnotating: {self.image_path}")
print("Instructions:")
print("- Click and drag to select license plate regions")
print("- Press 'u' to undo last annotation")
print("- Press 's' to save annotations")
print("- Press 'n' to go to next image")
print("- Press 'q' to quit program")
print("- Press 'r' to reset all annotations")
self.update_display()
while True:
key = cv2.waitKey(1) & 0xFF
if key == ord('q'):
cv2.destroyWindow(self.window_name)
return None # Signal to quit program
elif key == ord('n'):
cv2.destroyWindow(self.window_name)
return self.annotations # Signal to go to next image
elif key == ord('u') and self.annotations:
self.annotations.pop()
print(f"Undid last annotation. {len(self.annotations)} remaining.")
self.update_display()
elif key == ord('r'):
self.annotations = []
print("Reset all annotations.")
self.update_display()
elif key == ord('s'):
self.save_annotations()
cv2.destroyWindow(self.window_name)
return self.annotations
def save_annotations(self):
if not self.annotations:
print("No annotations to save.")
return
# Save to project annotations folder
image_name = Path(self.image_path).stem
output_file = self.project_dir / 'annotations' / f"{image_name}.json"
annotation_data = {
'image_path': str(self.image_path),
'image_shape': self.original.shape,
'annotations': self.annotations,
'statistics': self.get_statistics()
}
with open(output_file, 'w') as f:
json.dump(annotation_data, f, indent=2)
print(f"Saved {len(self.annotations)} annotations to {output_file}")
def get_statistics(self):
if not self.annotations:
return {}
widths = [a['width'] for a in self.annotations]
heights = [a['height'] for a in self.annotations]
aspect_ratios = [a['aspect_ratio'] for a in self.annotations]
return {
'count': len(self.annotations),
'avg_width': np.mean(widths),
'avg_height': np.mean(heights),
'avg_aspect_ratio': np.mean(aspect_ratios),
'width_range': [min(widths), max(widths)],
'height_range': [min(heights), max(heights)],
'aspect_ratio_range': [min(aspect_ratios), max(aspect_ratios)]
}
def analyze_project_annotations(project_dir):
"""Analyze all annotations in a project and generate detection parameters."""
project_dir = Path(project_dir)
annotations_dir = project_dir / 'annotations'
if not annotations_dir.exists():
print("No annotations directory found.")
return
annotation_files = list(annotations_dir.glob('*.json'))
if not annotation_files:
print("No annotation files found.")
return
all_annotations = []
all_stats = []
print(f"\nAnalyzing {len(annotation_files)} annotation files...")
for file in annotation_files:
with open(file, 'r') as f:
data = json.load(f)
all_annotations.extend(data['annotations'])
all_stats.append(data['statistics'])
print(f" {file.name}: {data['statistics']['count']} plates")
if not all_annotations:
print("No annotations found in files.")
return
# Calculate statistics
widths = [a['width'] for a in all_annotations]
heights = [a['height'] for a in all_annotations]
aspects = [a['aspect_ratio'] for a in all_annotations]
areas = [a['width'] * a['height'] for a in all_annotations]
print(f"\n=== DETECTION PARAMETERS ===")
print(f"Based on {len(all_annotations)} manually annotated license plates:")
print(f"")
print(f"Width range: {min(widths):3.0f} - {max(widths):3.0f} pixels (avg: {np.mean(widths):.0f})")
print(f"Height range: {min(heights):3.0f} - {max(heights):3.0f} pixels (avg: {np.mean(heights):.0f})")
print(f"Aspect ratio: {min(aspects):.2f} - {max(aspects):.2f} (avg: {np.mean(aspects):.2f})")
print(f"Area range: {min(areas):4.0f} - {max(areas):4.0f} pixels²")
print(f"")
print(f"GENERATED PARAMETERS:")
print(f" min_width = {int(np.percentile(widths, 10))}")
print(f" max_width = {int(np.percentile(widths, 90)) * 2}")
print(f" min_height = {int(np.percentile(heights, 10))}")
print(f" max_height = {int(np.percentile(heights, 90)) * 2}")
print(f" min_aspect_ratio = {np.percentile(aspects, 5):.2f}")
print(f" max_aspect_ratio = {np.percentile(aspects, 95):.2f}")
print(f" min_area = {int(np.percentile(areas, 10))}")
print(f" max_area = {int(np.percentile(areas, 90)) * 3}")
# Save parameters to project
params = {
'min_width': int(np.percentile(widths, 10)),
'max_width': int(np.percentile(widths, 90)) * 2,
'min_height': int(np.percentile(heights, 10)),
'max_height': int(np.percentile(heights, 90)) * 2,
'min_aspect_ratio': np.percentile(aspects, 5),
'max_aspect_ratio': np.percentile(aspects, 95),
'min_area': int(np.percentile(areas, 10)),
'max_area': int(np.percentile(areas, 90)) * 3
}
params_file = project_dir / 'debug' / 'detection_parameters.json'
with open(params_file, 'w') as f:
json.dump(params, f, indent=2)
print(f"\nParameters saved to {params_file}")
return params
def main():
parser = argparse.ArgumentParser(description='License Plate Annotation Tool')
parser.add_argument('--project-id', type=int, required=True, help='Project ID')
parser.add_argument('--image', help='Specific image to annotate')
parser.add_argument('--analyze', action='store_true', help='Analyze all annotations')
parser.add_argument('--reannotate', action='store_true', help='Reannotate images that already have annotations')
args = parser.parse_args()
project_dir = Path(f"projects/{args.project_id:03d}")
if not project_dir.exists():
print(f"Project {args.project_id:03d} does not exist. Create it first.")
return
if args.analyze:
analyze_project_annotations(project_dir)
return
raw_dir = project_dir / 'raw'
if not raw_dir.exists():
print(f"Raw directory not found: {raw_dir}")
return
# Get list of images
image_files = list(raw_dir.glob('*'))
image_files = [f for f in image_files if f.suffix.lower() in ['.jpg', '.jpeg', '.png', '.bmp']]
if not image_files:
print(f"No image files found in {raw_dir}")
return
if args.image:
# Annotate specific image
target_image = raw_dir / args.image
if not target_image.exists():
print(f"Image {args.image} not found in {raw_dir}")
return
image_files = [target_image]
# Filter images based on existing annotations
if not args.reannotate:
filtered_images = []
for image_file in image_files:
image_name = image_file.stem
annotation_file = project_dir / 'annotations' / f"{image_name}.json"
if not annotation_file.exists():
filtered_images.append(image_file)
else:
print(f"Skipping {image_file.name} (already annotated)")
image_files = filtered_images
if not image_files:
print("All images already annotated. Use --reannotate flag to reannotate.")
return
print(f"Found {len(image_files)} images to annotate in project {args.project_id:03d}")
for i, image_file in enumerate(image_files):
print(f"\n{'='*60}")
print(f"Annotating: {image_file.name} ({i+1}/{len(image_files)})")
print(f"{'='*60}")
# Load existing annotations if reannotating
load_existing = args.reannotate
annotator = LicensePlateAnnotator(str(image_file), project_dir, load_existing=load_existing)
annotations = annotator.run()
if annotations is None:
print("Annotation session cancelled by user.")
break
elif annotations:
annotator.save_annotations()
print(f"Completed annotation of {image_file.name}")
print(f"Found {len(annotations)} license plate regions")
else:
print(f"No annotations saved for {image_file.name}")
print(f"\nAnnotation session completed for project {args.project_id:03d}")
print("Run with --analyze flag to generate detection parameters.")
if __name__ == '__main__':
main()