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)