ImageRecovery.py

from abc import ABC, abstractmethod
from PIL import Image, ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True
from functools import reduce
from datetime import datetime
import shutil
import os

class ImageRecovery:
    def __init__(self, pathToRecoveryDirectory, pathToNewImagesDirectory, imageFilters):
        self.pathToRecoveryDirectory = pathToRecoveryDirectory
        self.pathToNewImagesDirectory = pathToNewImagesDirectory
        self.imageFilters = imageFilters

    def organizeRecoveredImages(self):
        images = self._getAllImagesFromDirectory(self.pathToRecoveryDirectory)
        filteredImages = self._filterOutImages(images)
        self._moveImagesToTheNewDirectory(filteredImages)

    def _getAllImagesFromDirectory(self, directoryPath):
        allImages = list()
        for entry in os.listdir(directoryPath):
            newPath = os.path.join(directoryPath, entry)
            if os.path.isdir(newPath):
                allImages += self._getAllImagesFromDirectory(newPath)
            elif self._isFileAnImage(newPath):
                allImages.append(newPath)

        return allImages

    def _isFileAnImage(self, file):
        fileLower = file.lower()
        return fileLower.endswith(".png") or \
               fileLower.endswith(".jpg") or \
               fileLower.endswith(".jpeg")

    def _filterOutImages(self, images):
        return list(reduce(self._filterReduces, self.imageFilters, images))

    def _filterReduces(self, images, imageFilter):
        return imageFilter.filter(images)

    def _moveImagesToTheNewDirectory(self, filteredImages):
        fileModificationAndImageMap = self._createFileModificationAndImage(filteredImages)
        for timestamp in fileModificationAndImageMap:
            self._moveImagesBasedOnTimestamp(timestamp, fileModificationAndImageMap[timestamp])

    def _createFileModificationAndImage(self, images):
        map = {}

        for imagePath in images:
            fileModificationDate = round(os.path.getmtime(imagePath), 1)

            if fileModificationDate not in map:
                map[fileModificationDate] = []

            map[fileModificationDate].append(imagePath)

        return map

    def _moveImagesBasedOnTimestamp(self, timestamp, images):
        newDirectoryPath = self._createNewDirectoryForFile(timestamp)
        for index, imagePath in enumerate(images):
            newImageName = self._createImageName(index, timestamp, imagePath)

            newImagePath = os.path.join(newDirectoryPath, newImageName)
            shutil.move(imagePath, newImagePath)

    def _createNewDirectoryForFile(self, fileModificationDate):
        newDirectoryName = self._getImageDirectoryName(fileModificationDate)
        newDirectoryPath = os.path.join(self.pathToNewImagesDirectory, newDirectoryName)
        if not os.path.exists(newDirectoryPath):
            os.makedirs(newDirectoryPath)

        return newDirectoryPath

    def _getImageDirectoryName(self, fileModificationDate):
        return datetime.fromtimestamp(fileModificationDate).strftime('%Y.%m')

    def _createImageName(self, index, fileModificationDate, imagePath):
        image = Image.open(imagePath)
        extension = "." + image.format.lower()

        newImageName = datetime.fromtimestamp(fileModificationDate).strftime('%Y%m%d_%H%M%S')
        if index != 0:
            newImageName += "_" + str(index)

        return newImageName + extension


class ImageFilter(ABC):
    @classmethod
    @abstractmethod
    def filter(cls, paths):
        pass


class AspectRationImageFilter(ImageFilter):
    def __init__(self):
        super().__init__()
        self.correctAspectRatios = ["4:3", "3:2", "5:3", "16:9", "16:10", "10:8", "7:5", "37:18"]

    def filter(self, paths):
        return list(filter(self._isAspectRatioCorrect, paths))

    def _isAspectRatioCorrect(self, path):
        image = Image.open(path)
        width, height = image.size

        divisor = self.gcd(width, height)
        aspectHeight = int(height / divisor)
        aspectWidth = int(width / divisor)

        firstPart = max(aspectHeight, aspectWidth)
        secondPart = min(aspectHeight, aspectWidth)
        aspectRatio = "{}:{}".format(firstPart, secondPart)

        return aspectRatio in self.correctAspectRatios

    def gcd(self, a, b):
        return a if b == 0 else self.gcd(b, a % b)


class SizeImageFilter(ImageFilter):
    def __init__(self, minSizeInBytes):
        super().__init__()
        self.minSizeInBytes = minSizeInBytes

    def filter(self, paths):
        return list(filter(self._isSizeOverMin, paths))

    def _isSizeOverMin(self, path):
        return os.path.getsize(path) > self.minSizeInBytes

def main():
    recoveryDirectory = os.path.join("Recovered")

    newImagesDirectory = os.path.join("New")
    if not os.path.exists(newImagesDirectory):
        os.makedirs(newImagesDirectory)

    filters = [AspectRationImageFilter(), SizeImageFilter(300000)]
    ir = ImageRecovery(recoveryDirectory, newImagesDirectory, filters)
    ir.organizeRecoveredImages()


if __name__ == '__main__':
    main()