edmt.py

#coding:utf-8

import cv2
import os
import sys
import numpy as np
import matplotlib.pyplot as plt

from sklearn.decomposition import RandomizedPCA
from sklearn import svm
from PIL import Image

moe_path = "/home/enmt/pdc_LT/MOE/"
ero_path = "/home/enmt/pdc_LT/ERO/"
tmp_path = "/home/enmt/pdc_LT/TMP/"
ero_gray_path = "/home/enmt/pdc_LT/ERO_GRAY/"
moe_gray_path = "/home/enmt/pdc_LT/MOE_GRAY/"
ero_append_path = "/home/enmt/pdc_LT/EMO_APPEND/"
moe_append_path = "/home/enmt/pdc_LT/MOE_APPEND/"
kmeans_path = "/home/enmt/pdc_LT/KMEANS/"

cascade_file_path = "/home/enmt/pdc_LT/work/lbpcascade_animeface.xml"


def img_to_matrix(filename):
    img = Image.open(filename)
    imgArray = np.asarray(img)
    return imgArray  


def flatten_image(img):
    s = img.shape[0] * img.shape[1] * img.shape[2]
    #s = img.shape[0]*img.shape[1]
    img_wide = img.reshape(1, s)
    return img_wide[0]


def rgb2gray(input_path,output_path):
	files = os.listdir(input_path)
	for file in files:
                fullpath = input_path+file
                im = cv2.imread(fullpath)
                gray = cv2.cvtColor(im,cv2.COLOR_RGB2GRAY)
                fullpath = output_path+file
                cv2.imwrite(fullpath,gray)


def resize(input_path,output_path,pinpoint=False):

	if pinpoint is True:
		input_fullpath = input_path
		output_fullpath = output_path
		pic = cv2.imread(input_fullpath)
		cascade = cv2.CascadeClassifier(cascade_file_path)
		facerect = cascade.detectMultiScale(pic,scaleFactor=1.2,minNeighbors=1,minSize=(10,10))

		for (x,y,w,h) in facerect:
			#cv2.rectangle(im,(x,y),(x+w,y+h+300),(0,0,255),2)
 			dst = pic[y:y+h+200,x:x+w]
			cv2.imwrite(output_fullpath,dst)
		img = Image.open(output_fullpath,'r')
		resize_img = img.resize((150,300))
		resize_img.save(output_fullpath,'jpeg',quality=100)		

	else:

		files = os.listdir(input_path)
       		cascade = cv2.CascadeClassifier(cascade_file_path)
	
		for file in files:
        		input_fullpath = input_path+file
        		pic = cv2.imread(input_fullpath)
			facerect = cascade.detectMultiScale(pic,scaleFactor=1.2, minNeighbors=1,minSize=(10, 10))	
		
			for (x,y,w,h) in facerect:
				#cv2.rectangle(im,(x,y),(x+w,y+h+300),(0,0,255),2)
				dst = pic[y:y+h+200,x:x+w]

			output_fullpath = output_path+file
			cv2.imwrite(output_fullpath,dst)

			img = Image.open(output_fullpath,'r')
			resize_img = img.resize((150,300))
			resize_img.save(output_fullpath,'jpeg',quality=100)
			#cv2.imshow(,im)
			#cv2.waitKey(0)


def calc_hist(img):
	color = ('b','g','r')
	for i,col in enumerate(color):
		histr = cv2.calcHist([img],[i],None,[256],[0,256])
		plt.plot(histr,color=col)
		plt.xlim([0,256])
	plt.show()


def kmeans(input_path,output_path,pinpoint=False):
	
	if pinpoint is True:

		input_fullpath = input_path
		output_fullpath = output_path

		pic = cv2.imread(input_fullpath)
		z = pic.reshape((-1,3))
		z = np.float32(z)
		criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)
		k = 4
		ret,label,center = cv2.kmeans(z,k,criteria,10,0)
		center = np.uint8(center)
		res = center[label.flatten()]
		dst = res.reshape((pic.shape))

		cv2.imwrite(output_fullpath,dst)

	else:
		files = os.listdir(input_path)
		for file in files:
			input_fullpath = input_path+file
			pic = cv2.imread(input_fullpath)
			z = pic.reshape((-1,3))
			z = np.float32(z)
			criteria = (cv2.TERM_CRITERIA_EPS + cv2.TERM_CRITERIA_MAX_ITER, 10, 1.0)
			k = 4
			ret,label,center = cv2.kmeans(z,k,criteria,10,0)
		
			center = np.uint8(center)
			res = center[label.flatten()]
			dst = res.reshape((pic.shape))
		
			output_fullpath = output_path+file
			cv2.imwrite(output_fullpath,dst)
			#calc_hist(dst)
			#cv2.imshow('tesuya',dst)
			#cv2.waitKey(0)	


def pca(input_path,test_data_path):
	data = []
	files = os.listdir(input_path+"ERO/")

	for file in files:
		fullpath = input_path+"ERO/"+file
		img = img_to_matrix(fullpath)
		img = flatten_image(img)
		data.append(img)

	files = os.listdir(input_path+"MOE/")

	for file in files:
		fullpath = input_path+"MOE/"+file
		img = img_to_matrix(fullpath)
                img = flatten_image(img)
                data.append(img)

	img = img_to_matrix(test_data_path)
	img = flatten_image(img)
	data.append(img)

	data = np.array(data)	
	pca = RandomizedPCA(n_components=3)
	x = pca.fit_transform(data)
	np.savetxt("./result.csv",x,delimiter=",")
	return x
		
if __name__=='__main__':
	#resize(ero_append_path,tmp_path+"ERO/")
	#resize(moe_append_path,tmp_path+"MOE/")
	
	#kmeans(tmp_path+"ERO/",kmeans_path+"ERO/")
	#kmeans(tmp_path+"MOE/",kmeans_path+"MOE/")
	
	#rgb2gray(tmp_path+"ERO/",ero_gray_path)
	#rgb2gray(tmp_path+"MOE/",moe_gray_path)	

	argvs = sys.argv
	argc = len(argvs)

	if argc != 2:
		print "Usage: #python %s filepath"%(argvs[0])
		exit()
		
	print "EDMTの眼が起動しました."

	path = argvs[1]
	resize(path,path,True)	
	kmeans(path,path,True)
	
	x = pca(kmeans_path,path)
	X = []
	y = []

	for index in range(100):
		X.append(x[index])
		#print x[index]
	
	for index in range(100):
		if 0 <= index and index <= 49:
			y.append(1)
		else:
			y.append(0)				 
	
	clf = svm.SVC(kernel="linear")
	clf.fit(X,y)

	"""
	for index in range(100):
		print x[index]
		print clf.predict(x[index])
		print "\n"
	
	print clf.predict(x[100])
	"""	

	ans = clf.predict(x[100])
	if ans[0] == 1:
		print "エロ画像です."
	else:
		print "萌え画像です."