116 lines
3.1 KiB
Python
116 lines
3.1 KiB
Python
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)
|