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.05 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
9
from time import time
10
from math import floor
11 12
import numpy as np
import youtube_dl
Hachemin Pierre-Yves's avatar
Hachemin Pierre-Yves committed
13
import cv2
Hachemin Pierre-Yves's avatar
Hachemin Pierre-Yves committed
14
import os
15
# import ffmpeg : imageio.plugins.ffmpeg.download()
Hachemin Pierre-Yves's avatar
Hachemin Pierre-Yves committed
16 17 18 19
import moviepy.editor as mp

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

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

35 36 37 38 39 40 41 42
    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
43

44
def downloadTrailer(): #called by main()
45 46 47 48
    """Function downloading the trailer thanks to the youtube-dl library"""
    global countDownload
    global exceptDict
    ydl_opts = {'format': 'worst/worstvideo',
49 50
                'outtmpl': videoDir+'%(id)s.%(ext)s',
                'noplaylist': True,
51
                'nocheckcertificate':True,
52
                #'max_filesize' : 10000000,
SoleneDc's avatar
SoleneDc committed
53
                "nocheckcertificate": True,
54
                'save_path' : videoDir }
55
    for key in linkDict.keys():
56 57 58 59 60 61
        if not os.path.isfile(spectrumDir+linkDict[key][2]+'.jpg'):
            with youtube_dl.YoutubeDL(ydl_opts) as ydl:
                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)
62 63
                    max_filesize=info_dict.get('filesize',None)
                    #print("max_filesize : {}".format(max_filesize))
64 65 66 67 68 69 70 71 72 73
                    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")
    				f.write(str(exceptDict))
    				f.close()"""
            print("{}/{} downloads".format(countDownload, len(linkDict)))

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

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

77
def createSpectrum(): #called by main()
78 79 80 81 82
    """Function creating the spectrum from the full queue"""
    global countSpectrum
    global spectrumDir
    while (not queue.empty()) or follow:
        item = queue.get()
83
        imgSpectrum(videoDir+item, spectrumDir + item[:-4] + '.jpg')
84 85 86 87
        print(item[:-4])
        global countDownload
        print("{}/{} spectrum".format(countSpectrum, len(linkDict)))
        countSpectrum += 1
Hachemin Pierre-Yves's avatar
Hachemin Pierre-Yves committed
88
        os.remove(videoDir+item)
89 90
        queue.task_done()

91 92 93 94 95 96 97 98 99 100 101 102 103
def imgSpectrum(vidPath, spectrumOut): #called by createSpectrum()
    """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("--------------------------------------------------------")
    print("Duration for {} : {} s".format(vidPath[:-3],time() - start))
    print("--------------------------------------------------------")
Hachemin Pierre-Yves's avatar
Hachemin Pierre-Yves committed
104

105 106
def resizeVideo(vidIn, vidOut): #called by imgSpectrum()
    """Function to resize the video in case it is too big (mp4 format)"""
Hachemin Pierre-Yves's avatar
Hachemin Pierre-Yves committed
107
    try:
108 109 110
        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
111
    except Exception as ex:
112 113
        print(str(ex))
        return vidOut
Hachemin Pierre-Yves's avatar
Hachemin Pierre-Yves committed
114 115
    return vidOut

116
def tableSpectrum(video): #called by imgSpectrum()
117 118 119 120
    """Function creating the array for the spectrum"""
    count = 0
    ratioFrame = 6
    countFrame = 0
121
    vidcrop=cv2.VideoCapture(video)
122 123

    """Find if the trailer is in 4:3 or 21:9 --> borders to crop"""
124 125 126 127 128
    success, image = vidcrop.read()
    countcrop=0
    while success and countcrop<240:
        success,image=vidcrop.read()
        countcrop+=1
Hachemin Pierre-Yves's avatar
Hachemin Pierre-Yves committed
129

130
    cropdim=crop_image_dim(image,0)
131 132
    print("image size : {}".format(image.shape)) 
    print("cropped dim : {}".format(cropdim))
133

134
    """Read every images of the trailer and out the mean color of each"""
135 136 137 138 139 140 141 142 143
    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
144
            res[0][count] = meanImage(image,cropdim)
145 146 147 148 149 150 151 152 153 154
            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


155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179
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)
Hachemin Pierre-Yves's avatar
Hachemin Pierre-Yves committed
180

181 182 183 184 185 186 187 188 189 190 191 192 193
def meanImage(img,cropdim): #called by tableSpectrum()
    """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)
    for i in range(cropdim[0],cropdim[1]):
        for j in range(cropdim[2],cropdim[3]):
            res0 = res0 + lab[i][j]
    res0 = res0 / ((cropdim[1]-cropdim[0]+1)*(cropdim[3]-cropdim[2]+1))
    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
194 195


196 197 198 199 200 201
def reading(path): #Called by main()
    """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
202 203

if __name__ == '__main__':
204
    main()