diff --git a/files.py b/files.py index c3e3d10..1da3ba4 100644 --- a/files.py +++ b/files.py @@ -13,6 +13,7 @@ import hashlib import exifread import base64 import numpy +import cv2 ################################################################################ # Local Class imports @@ -76,6 +77,42 @@ class FileData(): except: return False + def generateVideoThumbnail(self, file): + #overall wrapper function for generating video thumbnails + frame = self.video_to_frames(file)[0] + fthumbnail = self.image_to_thumbs(frame)["160"] + return fthumbnail + + def video_to_frames(self, video_filename): + #Extract frames from video + cap = cv2.VideoCapture(video_filename) + video_length = int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) - 1 + frames = [] + if cap.isOpened() and video_length > 0: + frame_ids = [0] + if video_length >= 4: + frame_ids = [round(video_length * 0.25)] + count = 0 + success, image = cap.read() + while success: + if count in frame_ids: + frames.append(image) + success, image = cap.read() + count += 1 + return frames + + def image_to_thumbs(self, img): + #Create thumbs from image + height, width, channels = img.shape + thumbs = {"original": img} + sizes = [160] + for size in sizes: + if (width >= size): + r = (size + 0.0) / width + max_size = (size, int(height * r)) + thumbs[str(size)] = cv2.resize(img, max_size, interpolation=cv2.INTER_AREA) + return thumbs + ############################################################################## # HACK: At present this only handles one path (need to re-factor if we have # # multiple valid paths in import_path) # @@ -106,6 +143,7 @@ class FileData(): fthumbnail = self.getExif(file) elif self.isVideo(file): ftype = 'Video' + fthumbnail = self.generateVideoThumbnail(file) else: ftype = 'File'