ImageAnnotator.py 2.1 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758
  1. import ipywidgets as widgets
  2. from IPython.display import display
  3. from py.Session import Session
  4. from py.ImageClassifier import AbstractImageClassifier
  5. import numpy as np
  6. # Old annotator script using IPython widgets, which are very slow and sometimes buggy.
  7. # It is preferred to use quick_label.py.
  8. class ImageAnnotator():
  9. def __init__(self, classifier: AbstractImageClassifier, session: Session, initial_scores = [], initial_annotations = [], load_from = None):
  10. self.scores = initial_scores
  11. self.annotations = initial_annotations
  12. self.score = -1
  13. self.classifier = classifier
  14. self.session = session
  15. if load_from is not None:
  16. data = np.load(load_from, allow_pickle=True)
  17. self.annotations = data[0]
  18. self.scores = data[1]
  19. normal_btn = widgets.Button(description = "Normal")
  20. anomalous_btn = widgets.Button(description = "Anomalous")
  21. self.button_box = widgets.HBox([normal_btn, anomalous_btn])
  22. self.output = widgets.Output(layout={"height": "400px"})
  23. display(self.button_box, self.output)
  24. normal_btn.on_click(self.mark_as_normal)
  25. anomalous_btn.on_click(self.mark_as_anomalous)
  26. self.next_image()
  27. # Click on normal button
  28. def mark_as_normal(self, _):
  29. with self.output:
  30. print("Marking as normal...")
  31. self.annotations.append(True)
  32. self.scores.append(self.score)
  33. self.next_image()
  34. # Click on anomalous button
  35. def mark_as_anomalous(self, _):
  36. with self.output:
  37. print("Marking as anomalous...")
  38. self.annotations.append(False)
  39. self.scores.append(self.score)
  40. self.next_image()
  41. # Show next image
  42. def next_image(self):
  43. img = self.session.get_random_motion_image(day_only=True)
  44. self.score = self.classifier.evaluate(img)
  45. self.output.clear_output()
  46. with self.output:
  47. display(img.to_ipython_image())
  48. print(f"score = {self.score}")
  49. # Save annotation data to file
  50. def save(self, filename: str):
  51. np.save(filename, [self.annotations, self.scores])