|
| 1 | +# General imports |
| 2 | +import os |
| 3 | + |
| 4 | +import pandas as pd |
| 5 | +import numpy as np |
| 6 | + |
| 7 | +# ETS imports |
| 8 | +from traits.api import Directory, Event, HasStrictTraits, Instance |
| 9 | + |
| 10 | +# Local imports |
| 11 | +from .image_file import ImageFile, SUPPORTED_FORMATS |
| 12 | + |
| 13 | +FILENAME_COL = "filename" |
| 14 | + |
| 15 | +NUM_FACE_COL = "Num. faces" |
| 16 | + |
| 17 | + |
| 18 | +class ImageFolder(HasStrictTraits): |
| 19 | + """ Model to hold an image folder. |
| 20 | + """ |
| 21 | + path = Directory |
| 22 | + |
| 23 | + data = Instance(pd.DataFrame) |
| 24 | + |
| 25 | + data_updated = Event |
| 26 | + |
| 27 | + def __init__(self, **traits): |
| 28 | + # Don't forget this! |
| 29 | + super(ImageFolder, self).__init__(**traits) |
| 30 | + if not os.path.isdir(self.path): |
| 31 | + msg = f"Unable to create an ImageFolder from {self.path} since" \ |
| 32 | + f" it is not a valid directory." |
| 33 | + raise ValueError(msg) |
| 34 | + |
| 35 | + self.data = self.to_dataframe() |
| 36 | + |
| 37 | + def to_dataframe(self): |
| 38 | + if not self.path: |
| 39 | + return pd.DataFrame({FILENAME_COL: [], NUM_FACE_COL: []}) |
| 40 | + |
| 41 | + data = [] |
| 42 | + for filename in os.listdir(self.path): |
| 43 | + file_ext = os.path.splitext(filename)[1].lower() |
| 44 | + if file_ext in SUPPORTED_FORMATS: |
| 45 | + filepath = os.path.join(self.path, filename) |
| 46 | + img_file = ImageFile(filepath=filepath) |
| 47 | + file_data = {FILENAME_COL: filename, NUM_FACE_COL: np.nan} |
| 48 | + try: |
| 49 | + file_data.update(img_file.metadata) |
| 50 | + except Exception: |
| 51 | + pass |
| 52 | + data.append(file_data) |
| 53 | + |
| 54 | + return pd.DataFrame(data) |
| 55 | + |
| 56 | + def compute_num_faces(self, **kwargs): |
| 57 | + cols = list(self.data.columns) |
| 58 | + for i, filename in enumerate(self.data[FILENAME_COL]): |
| 59 | + print(filename) |
| 60 | + filepath = os.path.join(self.path, filename) |
| 61 | + img_file = ImageFile(filepath=filepath) |
| 62 | + faces = img_file.detect_faces(**kwargs) |
| 63 | + j = cols.index(NUM_FACE_COL) |
| 64 | + self.data.iloc[i, j] = len(faces) |
| 65 | + self.data_updated = True |
0 commit comments