vjloop-organiser/app.py

import cv2
import os
import math
from colors import colorslist
import glob
import shutil

def takescreenshots(video):
    print(video)
    cap = cv2.VideoCapture(videopath + video)
    framerate = cap.get(cv2.CAP_PROP_FPS)
    while cap.isOpened():
        frameId = cap.get(1)
        ret, frame = cap.read()
        if not ret:
            break
        if frameId % math.floor(framerate) == 0:
            filename = imagepath + "/screenshot_" + str(int(frameId)) + ".jpg"
            cv2.imwrite(filename, frame)


def getcolors(image):
    img = cv2.imread(imagepath + image)
    hsv = cv2.cvtColor(img, cv2.COLOR_BGR2HSV)
    return hsv


def getcolormask(hsv,lower, upper):
    colormask = cv2.inRange(hsv, lower, upper)
    return colormask


def getcolorpercentage(mask):
    percentage = round(((mask > 0).mean()) * 100, 3)
    return percentage


def calcaverage(color):
    return sum(color)/len(color)


def gettopcolor(colors):
    color = max(colors, key=colors.get)
    result = {color: results.pop(color)}
    return result

def movetofolder(folder):
    if not os.path.exists('/media/vm/HDD/organisedloops/' + folder):
        os.makedirs('/media/vm/HDD/organisedloops/' + folder)
    shutil.move(videopath + video, '/media/vm/HDD/organisedloops/' + folder)


# Create the temp dir
if not os.path.exists('data'):
    os.makedirs('data')

# loop through video files
videopath = '/media/vm/HDD/organise/'
imagepath = 'data/'
videos = os.listdir(videopath)
for video in videos:
    percentagedict = {}
    for color in colorslist:
        percentagedict[color['color']] = []

    takescreenshots(video)

    # analyze each screenshot for colors
    images = os.listdir(imagepath)
    for image in images:
        hsv = getcolors(image)

        for color in colorslist:
            mask = getcolormask(hsv, color['lowerlimit'], color['upperlimit'])
            percentagedict[color['color']].append(getcolorpercentage(mask))

    results = {}
    for color in colorslist:
        results[color['color']] = calcaverage(percentagedict[color['color']])
    results['red'] = results['red1'] + results['red2']
    del results['red1'], results['red2']

    primarycolor = gettopcolor(results)
    print(primarycolor)
    secondarycolor = gettopcolor(results)
    print(secondarycolor)
    tertiarycolor = gettopcolor(results)
    print(tertiarycolor)

    # If color detection fails
    if next(iter(primarycolor.values())) <= 0.001:
        folder = 'nofolder'
        movetofolder(folder)


    # if one color dominates
    elif next(iter(secondarycolor.values())) <= next(iter(primarycolor.values())) * 0.05:
        folder = next(iter(primarycolor.keys()))
        movetofolder(folder)

    else:
        # if two colors dominate
        if next(iter(secondarycolor.values())) >= next(iter(tertiarycolor.values())) * 4:
            folder = next(iter(primarycolor.keys())) + ' + ' + next(iter(secondarycolor.keys()))
            movetofolder(folder)

        # if three colors dominate
        else:
            folder = 'multicolor'
            movetofolder(folder)

    # clear data folder
    files = glob.glob('data/*')
    for f in files:
        os.remove(f)