Detecting Kuriyama's glasses

kuriyama-san-megane.py

#!/usr/bin/env python
import sys
import os
import shutil
import matplotlib.pyplot as plt

import cv2
import numpy as np

def color_cluster(src_img, DEBUG):

    hsv_img = cv2.cvtColor(src_img, cv2.COLOR_BGR2HSV_FULL)             

    h, s, v = hsv_img[:,:,0], hsv_img[:,:,1], hsv_img[:,:,2]
    
    med_h = np.median(hsv_img[:,:,0])
    med_s = np.median(hsv_img[:,:,1])
    med_v = np.median(hsv_img[:,:,2])
    if DEBUG == 1:
        print ('med_hsv:', med_h, med_s, med_v)
    div = 180
    hight, width, chan = hsv_img.shape
    hClp = hight / div
    wClp = width / div
    FLG = 0
    for x in range(0,div):
        for y in range(0,div):
            hsv_Clp = \
                hsv_img[int(y * hClp):int((y + 1) * hClp), \
                         int(x * wClp):int((x + 1) * wClp)]
            h_c, s_c, v_c = hsv_Clp[:,:,0], hsv_Clp[:,:,1], hsv_Clp[:,:,2]
            med_h_c = np.median(hsv_Clp[:,:,0])
            med_s_c = np.median(hsv_Clp[:,:,1])
            med_v_c = np.median(hsv_Clp[:,:,2])
            if (med_h_c < 8 or med_h_c > 230) and \
               med_s_c < 148 and med_s_c > 128:
                FLG = 1 ; break
        if FLG == 1: break

    if DEBUG == 1:
       print('med_h_c med_s_c med_v_c ', med_h_c,  med_s_c,  med_v_c)
       print('x y ', str(x * wClp) , str(y * hClp))
       cv2.imshow("Clip", cv2.cvtColor(hsv_Clp, cv2.COLOR_HSV2BGR_FULL))
       cv2.waitKey(0)
       check_img = src_img.copy()
       cv2.rectangle(check_img, (int(x * wClp), int(y * hClp)), \
                     (int(x * wClp + wClp) , int(y * hClp + hClp)), (0,0,255), 1)
       cv2.imshow("check", check_img)
       cv2.waitKey(0)
    if FLG == 1:
        hl_save, hu_save = 0, 0
        hl = med_h_c - 16 ; hl = min(255, max(hl, 0))
        hu = med_h_c + 6
        if hu > 255:
            hl_save = hl - 8; hu_save = 255
            hl = 0 ; hu = hu -255 + 8
        hu = min(255, max(hu, 0))
        sl = med_s_c - 52 ; sl = min(255, max(sl, 0))
        su = med_s_c + 28 ; su = min(255, max(su, 0))
        vl = med_v_c - 20 ; vl = min(255, max(vl, 0))
        vu = med_v_c + 48 ; vu = min(255, max(vu, 0))
    else:
        return src_img

    if DEBUG == 1:
       print('hl sl vl ', hl , sl, vl )
       print('hu su vu ', hu , su, vu )

    lower = np.array([hl,sl,vl], dtype=np.uint8)
    upper = np.array([hu,su,vu], dtype=np.uint8)
    mask = cv2.inRange(hsv_img, lower, upper)
    if hl_save > 0:
        lower = np.array([hl_save,sl,vl], dtype=np.uint8)
        upper = np.array([hu_save,su,vu], dtype=np.uint8)
        mask = cv2.inRange(hsv_img, lower, upper)

    img_red = cv2.bitwise_and(src_img, src_img, mask= mask)

    if DEBUG == 1:
        cv2.imshow("mask_image", img_red)
        cv2.waitKey(0)
        cv2.imwrite("mask.jpg", img_red)

    contours, hierarchy = cv2.findContours(mask, cv2.RETR_EXTERNAL, \
                                           cv2.CHAIN_APPROX_SIMPLE)
    filtered_contour = []

    result_img = src_img.copy()

    for indx in range(len(contours)):
        cnt = contours[indx]
        if (len(cnt) > 40):
            if cv2.contourArea(cnt) > len(cnt):
                np_contour = \
                    np.array(cnt).reshape(len(cnt),2)
                left_x = min(np_contour[:,0])
                right_x = max(np_contour[:,0])
                top_y = min(np_contour[:,1])
                bottom_y = max(np_contour[:,1])
                ratio = (right_x - left_x) /  (bottom_y - top_y)
                if ratio > 0.5 and ratio < 4:
                    area = cv2.contourArea(cnt)
                    rec_area = (right_x - left_x) * (bottom_y - top_y)
                    area_ratio = area / rec_area
                    if area_ratio > 0.08:
                        filtered_contour.append(cnt)
                        cv2.rectangle(result_img, (left_x,top_y), \
                              (right_x, bottom_y), (0,0,255), 2)
                    if DEBUG == 1:
                        print('ID:', indx) ; print('len:', len(cnt))
                        print('Area', area) ; print('ratio', ratio)
                        print('area_ratio', area_ratio)
                        cv2.putText(result_img, '['+ str(indx) +']', \
                                    (left_x, top_y), \
                                    cv2.FONT_HERSHEY_PLAIN, \
                                    0.8, (255, 255, 255), 1, cv2.LINE_AA)
                pass
    result_img = cv2.drawContours(result_img, filtered_contour, -1, (0,255,0), 1)

    if DEBUG == 1:
        cv2.imshow("result",result_img)
        cv2.waitKey(0)
    return result_img


if __name__ == "__main__":

    IMG_FILE = sys.argv[1]
    if (len(sys.argv)) == 3:
        OUT_DIR = sys.argv[2]

    if os.path.isfile(IMG_FILE):
        DEBUG = 1
        src_img = cv2.imread(IMG_FILE)
        result_img = color_cluster(src_img, DEBUG)
        sys.exit()
        
    if os.path.isdir(OUT_DIR):
        shutil.rmtree(OUT_DIR)
        os.makedirs(OUT_DIR)

    if os.path.isdir(IMG_FILE):
        DEBUG = 0
        IMG_DIR = IMG_FILE
        files = os.listdir(IMG_DIR)
        files = sorted(files)

        for file in files:
            src_img_path = IMG_DIR + file    
            src_img = cv2.imread(src_img_path)    
            result_img = color_cluster(src_img, DEBUG)
            fname = file.split('.')[0]
            cv2.imwrite(OUT_DIR + 'megane-' + fname + '.jpg', result_img)