image_comp2.py

#!/usr/bin/env python
import os
import sys
import shutil
import numpy as np
import cv2

def main():
    IMG_DIR = sys.argv[1]
    OUT_DIR = sys.argv[2]
    SAME_DIR = OUT_DIR + '/' + 'same' + '/'
    COMP1_DIR = OUT_DIR + '/' + 'comp1' + '/'
    COMP2_DIR = OUT_DIR + '/' + 'comp2' + '/'
    STAT_TXTF = OUT_DIR + '/' + 'stat.txt'
    FPS = 24
    frame_ct = 0
    comp_ct = 0
    comp1_ct = 0
    comp2_ct = 0
    diff_cnt_acc = 0
    im_diff_acc = 0
    
    if os.path.exists(OUT_DIR):
            shutil.rmtree(OUT_DIR)
    os.mkdir(OUT_DIR)
    os.mkdir(SAME_DIR)
    os.mkdir(COMP1_DIR)
    os.mkdir(COMP2_DIR)
    
    files = os.listdir(IMG_DIR)
    files = sorted(files)

    img_file = IMG_DIR + '/'+ files[0]    
    prev_img = cv2.imread(img_file)

    for file in files:
        img_file = IMG_DIR + '/'+ file
        frame_ct = frame_ct + 1
        fname = file.split('.')[0]
        img = cv2.imread(img_file)
        if np.array_equal(prev_img, img):
            cv2.imwrite(SAME_DIR + fname + '-same' + '.jpg',img)
            same_img = img.copy()
            cv2.putText(same_img, 'same', (0, 50), \
                       cv2.FONT_HERSHEY_PLAIN, 4, (255, 255, 255), 5, cv2.LINE_AA)
            cv2.imwrite(OUT_DIR + fname + '-same' + '.jpg', same_img)
        else:
            comp_ct = comp_ct + 1
            cv2.imwrite(OUT_DIR + fname + '-comp' + '.jpg',img)
            diff_cnt = np.count_nonzero(img != prev_img)
            diff_cnt_acc = diff_cnt_acc + diff_cnt
            if diff_cnt > img.size / 2:
                    comp1_ct = comp1_ct + 1
                    cv2.imwrite(COMP1_DIR + fname + '-comp1' + '.jpg',img)
            if diff_cnt > img.size / 16:
                    comp2_ct = comp2_ct + 1
                    cv2.imwrite(COMP2_DIR + fname + '-comp2' + '.jpg',img)
            im_diff = img.astype(int) - prev_img.astype(int)
            im_diff_abs = np.abs(im_diff)
            im_diff_acc = im_diff_acc + im_diff_abs.max()
            im_diff_abs_norm = im_diff_abs / im_diff_abs.max() * 255
            cv2.putText(im_diff_abs_norm, str(diff_cnt), (0, 50), \
                       cv2.FONT_HERSHEY_PLAIN, 4, (255, 255, 255), 5, cv2.LINE_AA)
            cv2.imwrite(OUT_DIR + fname + '-diff' + '.jpg', im_diff_abs_norm)
            prev_img = img.copy()

    comp_ratio = comp_ct / frame_ct * FPS
    comp1_ratio = comp1_ct / frame_ct * FPS
    comp2_ratio = comp2_ct / frame_ct * FPS
    diff_cnt_acc_ratio = diff_cnt_acc / frame_ct * FPS
    im_diff_acc_ratio = im_diff_acc / frame_ct * FPS
    score = comp1_ratio * diff_cnt_acc_ratio / 1000000
    f = open(STAT_TXTF, 'w')
    
    f.write('frame_ct:' + str(frame_ct) + '\n')
    f.write('comp_ct:' + str(comp_ct) + '\n')
    f.write('comp_ratio:' + str(comp_ratio) + '/sec\n')
    f.write('comp1_ct:' + str(comp1_ct) + '\n')
    f.write('comp1_ratio:' + str(comp1_ratio) + '/sec\n')
    f.write('comp2_ct:' + str(comp2_ct) + '\n')
    f.write('comp2_ratio:' + str(comp2_ratio) + '/sec\n')
    f.write('diff_cnt_acc:' + str(diff_cnt_acc) + '\n')
    f.write('diff_cnt_acc_ratio:' + str(diff_cnt_acc_ratio) + '/sec\n')
    f.write('im_diff_acc:' + str(im_diff_acc) + '\n')
    f.write('im_diff_acc_ratio:' + str(im_diff_acc_ratio) + '/sec\n')
    f.write('Score:' + str(score) + '\n')

    f.close()

if __name__ == '__main__':
        main()