Ce serveur Gitlab sera éteint le 30 juin 2020, pensez à migrer vos projets vers les serveurs gitlab-research.centralesupelec.fr et gitlab-student.centralesupelec.fr !

Extract_spectrum.py 7.59 KB
Newer Older
Hachemin Pierre-Yves's avatar
Hachemin Pierre-Yves committed
1 2 3 4 5 6 7 8
#! /usr/bin/env python3
# coding: utf-8
# ydl1.py
from __future__ import unicode_literals
from threading import Thread
from queue import Queue
from ast import literal_eval
from skimage.transform import resize
SoleneDc's avatar
SoleneDc committed
9
import datetime
10
from time import time
11
from math import floor
12 13
import numpy as np
import youtube_dl
Hachemin Pierre-Yves's avatar
Hachemin Pierre-Yves committed
14
import cv2
Hachemin Pierre-Yves's avatar
Hachemin Pierre-Yves committed
15
import os
16
# import ffmpeg : imageio.plugins.ffmpeg.download()
Hachemin Pierre-Yves's avatar
Hachemin Pierre-Yves committed
17 18 19 20
import moviepy.editor as mp

follow = True
queue = Queue()
SoleneDc's avatar
SoleneDc committed
21
linkFile = './link-dictionaries/link-dictionary2015-1.txt' # Input dict of trailers #TODO; change path
22 23
linkDict = {}
exceptDict = {}
24
videoDir = './video/' # Folder to store temporarely the videos
SoleneDc's avatar
SoleneDc committed
25
spectrumDir = './spectrumImages/spectrumImages2015/' # Output folder to store the spectrums #TODO; change path
26 27
countDownload = 1
countSpectrum = 1
Hachemin Pierre-Yves's avatar
Hachemin Pierre-Yves committed
28

SoleneDc's avatar
SoleneDc committed
29

30 31 32 33 34 35
def main():
    global videoDir
    global spectrumDir
    global linkDict
    global linkFile
    global follow
Hachemin Pierre-Yves's avatar
Hachemin Pierre-Yves committed
36

37 38 39 40 41 42 43 44
    linkDict = reading(linkFile)
    TRAILER = Thread(target=downloadTrailer)
    SPECTRUM = Thread(target=createSpectrum)
    SPECTRUM.start()
    TRAILER.start()
    TRAILER.join()
    follow = False
    SPECTRUM.join()
Hachemin Pierre-Yves's avatar
Hachemin Pierre-Yves committed
45

SoleneDc's avatar
SoleneDc committed
46 47

def downloadTrailer():  # called by main()
48 49 50 51
    """Function downloading the trailer thanks to the youtube-dl library"""
    global countDownload
    global exceptDict
    ydl_opts = {'format': 'worst/worstvideo',
SoleneDc's avatar
SoleneDc committed
52
                'outtmpl': videoDir + '%(id)s.%(ext)s',
53
                'noplaylist': True,
SoleneDc's avatar
SoleneDc committed
54 55 56
                'nocheckcertificate': True,
                # 'max_filesize' : 10000000,
                'save_path': videoDir}
57
    for key in linkDict.keys():
SoleneDc's avatar
SoleneDc committed
58
        if not os.path.isfile(spectrumDir + linkDict[key][2] + '.jpg'):
SoleneDc's avatar
SoleneDc committed
59
            #time1 = datetime.datetime.now()
60
            with youtube_dl.YoutubeDL(ydl_opts) as ydl:
SoleneDc's avatar
SoleneDc committed
61
                #while (datetime.datetime.now() - time1).seconds < 190:
62 63 64 65
                try:
                    ydl.download(['https://www.youtube.com/watch?v=' + linkDict[key][2]])
                    info_dict = ydl.extract_info('https://www.youtube.com/watch?v=' + linkDict[key][2], download=False)
                    video_ext = info_dict.get("ext", None)
SoleneDc's avatar
SoleneDc committed
66 67
                    max_filesize = info_dict.get('filesize', None)
                    # print("max_filesize : {}".format(max_filesize))
68 69 70 71
                    queue.put(linkDict[key][2] + '.' + video_ext)
                except Exception as e:  # catch *all* exceptions
                    exceptDict[linkDict[key][2]] = str(e)
                    """f = open("Exception-dictionnary.txt", "w")
SoleneDc's avatar
SoleneDc committed
72 73
                    f.write(str(exceptDict))
                    f.close()"""
74 75 76
            print("{}/{} downloads".format(countDownload, len(linkDict)))

        else:
SoleneDc's avatar
SoleneDc committed
77
            print(spectrumDir + linkDict[key][2] + '.jpg already exists !')
78
        countDownload += 1
Hachemin Pierre-Yves's avatar
Hachemin Pierre-Yves committed
79

Hachemin Pierre-Yves's avatar
Hachemin Pierre-Yves committed
80

SoleneDc's avatar
SoleneDc committed
81
def createSpectrum():  # called by main()
82 83 84 85 86
    """Function creating the spectrum from the full queue"""
    global countSpectrum
    global spectrumDir
    while (not queue.empty()) or follow:
        item = queue.get()
SoleneDc's avatar
SoleneDc committed
87
        imgSpectrum(videoDir + item, spectrumDir + item[:-4] + '.jpg')
88 89 90 91
        print(item[:-4])
        global countDownload
        print("{}/{} spectrum".format(countSpectrum, len(linkDict)))
        countSpectrum += 1
SoleneDc's avatar
SoleneDc committed
92 93 94 95 96
        try:
            os.remove(videoDir + item)

        except FileNotFoundError:
            pass
97 98
        queue.task_done()

SoleneDc's avatar
SoleneDc committed
99 100

def imgSpectrum(vidPath, spectrumOut):  # called by createSpectrum()
101 102 103 104 105 106 107 108 109 110
    """Function saving the image thanks to the spectrum array """
    start = time()
    if vidPath[-3:] != '3gp':
        res2 = tableSpectrum(resizeVideo(vidPath, vidPath))
    else:
        res2 = tableSpectrum(vidPath)
    output = cv2.cvtColor(res2, cv2.COLOR_Lab2BGR)
    cv2.imwrite(spectrumOut, output)
    print()
    print("--------------------------------------------------------")
SoleneDc's avatar
SoleneDc committed
111
    print("Duration for {} : {} s".format(vidPath[:-3], time() - start))
112
    print("--------------------------------------------------------")
Hachemin Pierre-Yves's avatar
Hachemin Pierre-Yves committed
113

SoleneDc's avatar
SoleneDc committed
114 115

def resizeVideo(vidIn, vidOut):  # called by imgSpectrum()
116
    """Function to resize the video in case it is too big (mp4 format)"""
Hachemin Pierre-Yves's avatar
Hachemin Pierre-Yves committed
117
    try:
118 119 120
        clip = mp.VideoFileClip(vidIn)
        clip_resized = clip.resize(height=36)  # Make the height 36px
        clip_resized.write_videofile(vidOut)
Hachemin Pierre-Yves's avatar
Hachemin Pierre-Yves committed
121
    except Exception as ex:
122 123
        print(str(ex))
        return vidOut
Hachemin Pierre-Yves's avatar
Hachemin Pierre-Yves committed
124 125
    return vidOut

SoleneDc's avatar
SoleneDc committed
126 127

def tableSpectrum(video):  # called by imgSpectrum()
128 129 130 131
    """Function creating the array for the spectrum"""
    count = 0
    ratioFrame = 6
    countFrame = 0
SoleneDc's avatar
SoleneDc committed
132
    vidcrop = cv2.VideoCapture(video)
133 134

    """Find if the trailer is in 4:3 or 21:9 --> borders to crop"""
135
    success, image = vidcrop.read()
SoleneDc's avatar
SoleneDc committed
136 137 138 139 140 141 142
    if success:
        cropdim = (0, image.shape[0], 0, image.shape[1])
        countcrop = 0
    while success and countcrop < 240:
        success, image = vidcrop.read()
        countcrop += 1
    if success == True:
143
        try:
SoleneDc's avatar
SoleneDc committed
144
            cropdim = crop_image_dim(image, 0)
145 146
        except cv2.error:
            cropdim = (0, image.shape[0], 0, image.shape[1])
SoleneDc's avatar
SoleneDc committed
147
        print("image size : {}".format(image.shape))
148
        print("cropped dim : {}".format(cropdim))
149

150
    """Read every images of the trailer and out the mean color of each"""
151 152 153 154 155 156 157 158 159
    vidcap = cv2.VideoCapture(video)
    length = int(vidcap.get(cv2.CAP_PROP_FRAME_COUNT) / ratioFrame)
    width = int(length * 3 / 16)
    res = np.zeros([width, length, 3], dtype=np.uint8)
    success, image = vidcap.read()

    while success:
        if countFrame == ratioFrame:
            countFrame = 0
SoleneDc's avatar
SoleneDc committed
160
            res[0][count] = meanImage(image, cropdim)
161 162 163 164 165 166 167 168 169 170
            for i in range(1, res.shape[0]):
                res[i][count] = res[0][count]
            count += 1
            if cv2.waitKey(10) == 27:
                break
        success, image = vidcap.read()
        countFrame += 1
    return res


SoleneDc's avatar
SoleneDc committed
171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198
def crop_image_dim(img, tol):  # called by tableSpectrum()
    """Find the dimensions cropped if this is a 4:3 or 21:9 trailer"""
    # tol is the tolerance : 0 to find black pixels
    img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    mask = img_gray > tol
    midV = floor(mask.shape[0] / 2)
    midH = floor(mask.shape[1] / 2)
    upBorder = 0
    downBorder = mask.shape[0]
    leftBorder = 0
    rightBorder = mask.shape[1]
    for i in range(1, mask.shape[0]):
        if mask[i - 1][midH] != mask[i][midH] and mask[i][midH] == True:
            if upBorder == mask[0][midH]:
                upBorder = i - 1
        if mask[i - 1][midH] != mask[i][midH] and mask[i][midH] == False:
            downBorder = i
    for j in range(1, mask.shape[1]):
        if mask[midV][j - 1] != mask[midV][j] and mask[midV][j] == True:
            if leftBorder == mask[midV][0]:
                leftBorder = j - 1
        if mask[midV][j - 1] != mask[midV][j] and mask[midV][j] == False:
            rightBorder = j

    return (upBorder, downBorder, leftBorder, rightBorder)


def meanImage(img, cropdim):  # called by tableSpectrum()
199 200 201 202
    """Return the mean LAB array only with pixels in cropped image"""
    res0 = [0, 0, 0]
    res = [0, 0, 0]
    lab = cv2.cvtColor(img, cv2.COLOR_BGR2LAB)
SoleneDc's avatar
SoleneDc committed
203 204
    for i in range(cropdim[0], cropdim[1]):
        for j in range(cropdim[2], cropdim[3]):
205
            res0 = res0 + lab[i][j]
206
    sumpix=(cropdim[1]-cropdim[0]+1)*(cropdim[3]-cropdim[2]+1)
SoleneDc's avatar
SoleneDc committed
207 208 209 210 211
    if sumpix == 0:
        sumpix = 1
    res0[0] = res0[0] / sumpix
    res0[1] = res0[1] / sumpix
    res0[2] = res0[2] / sumpix
212 213 214 215
    res[0] = int(res0[0])
    res[1] = int(res0[1])
    res[2] = int(res0[2])
    return res
Hachemin Pierre-Yves's avatar
Hachemin Pierre-Yves committed
216 217


SoleneDc's avatar
SoleneDc committed
218
def reading(path):  # Called by main()
219 220 221 222 223
    """Function to create the dictionnary from the text file"""
    with open(path, 'r') as f:
        s = f.read()
        whip = literal_eval(s)
    return whip
Hachemin Pierre-Yves's avatar
Hachemin Pierre-Yves committed
224

SoleneDc's avatar
SoleneDc committed
225

Hachemin Pierre-Yves's avatar
Hachemin Pierre-Yves committed
226
if __name__ == '__main__':
227
    main()