AnnotatedImageView.py 3.2 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293
  1. #!/usr/bin/env python
  2. """AnnotatedImageView: Displays images with bounding boxes etc."""
  3. from PyQt5 import QtWidgets, QtGui, QtCore
  4. from ..utils import Video
  5. class AnnotatedImageView(QtWidgets.QGraphicsView):
  6. def __init__(self, parent):
  7. super(AnnotatedImageView, self).__init__(parent=parent)
  8. self._scene = QtWidgets.QGraphicsScene()
  9. self.setScene(self._scene)
  10. self._rect = QtCore.QRectF(0, 0, 0, 0)
  11. self.setRenderHint(QtGui.QPainter.Antialiasing)
  12. self._image = None
  13. self._current_video = None
  14. def display(self, prediction):
  15. self._scene.clear()
  16. if prediction['filetype'] == 'image':
  17. self._image = QtGui.QImage(prediction['filename'])
  18. item = QtWidgets.QGraphicsPixmapItem(QtGui.QPixmap.fromImage(self._image))
  19. self._scene.addItem(item)
  20. self._scene.setSceneRect(item.boundingRect())
  21. self._rect = item.boundingRect()
  22. if 'prediction' in prediction.keys():
  23. self._overlay(prediction['prediction'])
  24. elif prediction['filetype'] == 'video':
  25. if 'cap' in prediction.keys():
  26. cap = prediction['cap']
  27. else:
  28. cap = Video(prediction['filename'])
  29. if 'frame' in prediction.keys():
  30. frame = prediction['frame']
  31. else:
  32. frame = 0
  33. video_frame = cap.get_frame(frame)
  34. if video_frame is not None:
  35. self._image = QtGui.QImage(video_frame,
  36. video_frame.shape[1], video_frame.shape[0],
  37. video_frame.shape[2] * video_frame.shape[1], QtGui.QImage.Format_RGB888)
  38. item = QtWidgets.QGraphicsPixmapItem(QtGui.QPixmap.fromImage(self._image))
  39. self._scene.addItem(item)
  40. self._scene.setSceneRect(item.boundingRect())
  41. self._rect = item.boundingRect()
  42. if 'prediction-by-frame' in prediction.keys():
  43. self._overlay(prediction['prediction-by-frame'][frame])
  44. self.refit_display()
  45. def _overlay(self, prediction):
  46. if 'faces' in prediction.keys():
  47. self._overlay_faces(prediction['faces'])
  48. def _overlay_faces(self, faces):
  49. box_background = QtGui.QColor(255, 255, 255, 51)
  50. box_border = QtGui.QColor(0, 0, 0, 1)
  51. width = float(self._image.width())
  52. height = float(self._image.height())
  53. for face in faces:
  54. self._scene.addRect(face['x'] * width, face['y'] * height, face['w'] * width, face['h'] * height,
  55. QtGui.QPen(QtGui.QBrush(box_border), 3.0),
  56. QtGui.QBrush(box_background))
  57. def refit_display(self):
  58. if self._rect is None:
  59. self.resetTransform()
  60. else:
  61. self.fitInView(self._rect, QtCore.Qt.KeepAspectRatio)
  62. def resizeEvent(self, event: QtGui.QResizeEvent):
  63. self.refit_display()
  64. QtWidgets.QGraphicsView.resizeEvent(self, event)
  65. def reset(self):
  66. self._scene.clear()
  67. self._image = None
  68. self._current_video = None
  69. self._rect = None
  70. # TODO display context action