Add Stage3.0

Author: siddhantwahal

stage3.0_first_views/test_traited_face_detect.py

@@ -0,0 +1,37 @@
+from os.path import dirname, join
+from unittest import TestCase
+
+import numpy as np
+
+import ets_tutorial
+from traited_face_detect import ImageFile
+
+TUTORIAL_DIR = dirname(ets_tutorial.__file__)
+
+SAMPLE_IMG_DIR = join(TUTORIAL_DIR, "..", "sample_images")
+
+SAMPLE_IMG1 = join(SAMPLE_IMG_DIR, "IMG-0311_xmas_2020.JPG")
+
+
+class TestImageFile(TestCase):
+    def test_no_image_file(self):
+        img = ImageFile()
+        self.assertEqual(img.metadata, {})
+        data = img.to_array()
+        self.assertIsInstance(data, np.ndarray)
+        self.assertEqual(data.shape, (0,))
+
+    def test_image_metadata(self):
+        img = ImageFile(filepath=SAMPLE_IMG1)
+        self.assertNotEqual(img.metadata, {})
+        for key in ['ExifVersion', 'ExifImageWidth', 'ExifImageHeight']:
+            self.assertIn(key, img.metadata.keys())
+        data = img.to_array()
+        expected_shape = (img.metadata['ExifImageHeight'],
+                          img.metadata['ExifImageWidth'], 3)
+        self.assertEqual(data.shape, expected_shape)
+
+    def test_detects_faces(self):
+        img = ImageFile(filepath=SAMPLE_IMG1)
+        self.assertEqual(len(img.faces), 5)
+        self.assertEqual(img.metadata["Number of faces"], 5)

stage3.0_first_views/traited_face_detect.py

@@ -0,0 +1,100 @@
+# General imports
+from os.path import join
+
+import PIL.Image
+import matplotlib.pyplot as plt
+import numpy as np
+from matplotlib import patches
+from PIL.ExifTags import TAGS
+from skimage import data
+from skimage.feature import Cascade
+
+# ETS imports
+from traits.api import (
+    Dict,
+    File,
+    HasStrictTraits,
+    List,
+    observe,
+)
+from traitsui.api import OKButton, UItem, View
+
+
+class ImageFile(HasStrictTraits):
+    """ Model to hold an image file.
+    """
+    filepath = File
+
+    faces = List
+
+    metadata = Dict
+
+    def to_array(self):
+        if not self.filepath:
+            return np.array([])
+
+        with PIL.Image.open(self.filepath) as img:
+            return np.asarray(img)
+
+    @observe("filepath")
+    def _update_faces_and_metadata(self, event):
+        self.metadata = {}
+        self._update_metadata_with_exif()
+        self._detect_faces()
+        print(self.metadata)
+        print(f"Number of faces: {self.metadata['Number of faces']}")
+
+    def _update_metadata_with_exif(self):
+        if not self.filepath:
+            return
+        with PIL.Image.open(self.filepath) as img:
+            exif = img._getexif()
+        if not exif:
+            return
+        self.metadata.update(
+            {TAGS[k]: v for k, v in exif.items() if k in TAGS}
+        )
+
+    def _detect_faces(self):
+        self.faces = []
+        if not self.filepath:
+            return
+        # Load the trained file from the module root.
+        trained_file = data.lbp_frontal_face_cascade_filename()
+        # Initialize the detector cascade.
+        detector = Cascade(trained_file)
+        faces = detector.detect_multi_scale(
+            img=self.to_array(),
+            scale_factor=1.2,
+            step_ratio=1,
+            min_size=(60, 60),
+            max_size=(600, 600)
+        )
+        self.faces.extend(faces)
+        self.metadata['Number of faces'] = len(self.faces)
+
+
+if __name__ == '__main__':
+    img = ImageFile()
+    view = View(
+        UItem(name='filepath'),
+        buttons=[OKButton],
+        resizable=True,
+        width=640
+    )
+    img.configure_traits(view=view)
+
+    plt.imshow(img.to_array())
+    img_desc = plt.gca()
+    for patch in img.faces:
+        img_desc.add_patch(
+            patches.Rectangle(
+                (patch['c'], patch['r']),
+                patch['width'],
+                patch['height'],
+                fill=False,
+                color='r',
+                linewidth=2
+            )
+        )
+    plt.show()