anonymizer.py 2.7 KB

1234567891011121314151617181920212223242526272829303132333435363738394041424344454647484950515253545556575859606162636465666768697071
  1. import json
  2. from pathlib import Path
  3. import numpy as np
  4. from PIL import Image
  5. from tqdm import tqdm
  6. def load_np_image(image_path):
  7. image = Image.open(image_path).convert('RGB')
  8. np_image = np.array(image)
  9. return np_image
  10. def save_np_image(image, image_path):
  11. pil_image = Image.fromarray((image).astype(np.uint8), mode='RGB')
  12. pil_image.save(image_path)
  13. def save_detections(detections, detections_path):
  14. json_output = []
  15. for box in detections:
  16. json_output.append({
  17. 'y_min': box.y_min,
  18. 'x_min': box.x_min,
  19. 'y_max': box.y_max,
  20. 'x_max': box.x_max,
  21. 'score': box.score,
  22. 'kind': box.kind
  23. })
  24. with open(detections_path, 'w') as output_file:
  25. json.dump(json_output, output_file, indent=2)
  26. class Anonymizer:
  27. def __init__(self, detectors, obfuscator):
  28. self.detectors = detectors
  29. self.obfuscator = obfuscator
  30. def anonymize_image(self, image, detection_thresholds):
  31. assert set(self.detectors.keys()) == set(detection_thresholds.keys()),\
  32. 'Detector names must match detection threshold names'
  33. detected_boxes = []
  34. for kind, detector in self.detectors.items():
  35. new_boxes = detector.detect(image, detection_threshold=detection_thresholds[kind])
  36. detected_boxes.extend(new_boxes)
  37. return self.obfuscator.obfuscate(image, detected_boxes), detected_boxes
  38. def anonymize_images(self, input_path, output_path, detection_thresholds, file_types, write_json):
  39. print(f'Anonymizing images in {input_path} and saving the anonymized images to {output_path}...')
  40. Path(output_path).mkdir(exist_ok=True)
  41. assert Path(output_path).is_dir(), 'Output path must be a directory'
  42. files = []
  43. for file_type in file_types:
  44. files.extend(list(Path(input_path).glob(f'**/*.{file_type}')))
  45. for input_image_path in tqdm(files):
  46. # Create output directory
  47. relative_path = input_image_path.relative_to(input_path)
  48. (Path(output_path) / relative_path.parent).mkdir(exist_ok=True, parents=True)
  49. output_image_path = Path(output_path) / relative_path
  50. output_detections_path = (Path(output_path) / relative_path).with_suffix('.json')
  51. # Anonymize image
  52. image = load_np_image(str(input_image_path))
  53. anonymized_image, detections = self.anonymize_image(image=image, detection_thresholds=detection_thresholds)
  54. save_np_image(image=anonymized_image, image_path=str(output_image_path))
  55. if write_json:
  56. save_detections(detections=detections, detections_path=str(output_detections_path))