Commit 791e580d authored by Gihan Jayatilaka's avatar Gihan Jayatilaka

screen detectors

parent 0b2977e3
......@@ -15,13 +15,7 @@ from keras.models import load_model
from keras.layers import *
from keras import Model, Sequential
from keras.activations import relu,sigmoid
def conv2d_bn(x,
filters,
num_row,
num_col,
padding='same',
strides=(1, 1),
name=None):
def conv2d_bn(x,filters,num_row,num_col,padding='same',strides=(1, 1),name=None):
if name is not None:
bn_name = name + '_bn'
conv_name = name + '_conv'
......@@ -127,7 +121,7 @@ def load(fileName):
X=data['X']
temp=np.ndarray(shape=(X.shape[0],299,299,3),dtype=np.uint8)
for i in range(X.shape[0]):
temp[i]=cv2.resize(X[i],dsize=(299,299))
temp[i]=cv2.resize(X[i],dsize=(299,299),interpolation=cv2.INTER_LINEAR)
X=temp
Y=data['Y']
data=None
......@@ -135,6 +129,8 @@ def load(fileName):
def transformX(X):
if len(X.shape)==3:
X=np.reshape(X,(1,299,299,3))
X=(X.astype(np.float32)/128.0)-1
return X
......@@ -143,104 +139,94 @@ def inverseTransformX(X):
return ((X+1.0)*128.0).astype(np.uint8)
if __name__ == '__main__':
print("python screen-detetcion-5-nn.py trainData.npz 10 testVideo.mp4 modelSave.h5 0")
TRAINING_DATA=sys.argv[1]
EPOCHS=int(sys.argv[2])
print("python screen-detetcion-5-nn.py {test|train} model.h5 {train.npz|test.npz|test.mp4} 0 100")
print("args[0] args[1] args[2] args[3] [4][5]")
OPERATION=sys.argv[1]
MODEL_FILE_LOCATION=sys.argv[2]
TRAINING_DATA=sys.argv[3]
TEST_VIDEO=sys.argv[3]
MODEL_FILE_LOCATION=sys.argv[4]
CUDA_DEV=sys.argv[5]
CUDA_DEV=sys.argv[4]
X,Y=load(TRAINING_DATA)
os.environ["CUDA_VISIBLE_DEVICES"]="{}".format(CUDA_DEV)
os.environ["CUDA_VISIBLE_DEVICES"]="{}".format(CUDA_DEV)
if OPERATION=='train':
EPOCHS=int(sys.argv[5])
if DEBUG: print("DEBUG: Resized X shape={}".format(X.shape))
if DEBUG: print("Non-normalized X=[{},{}] Y=[{},{}]".format(np.min(X),np.max(X),np.min(Y),np.max(Y)))
X,Y=load(TRAINING_DATA)
X=transformX(X)
Y,yParams=transformY(Y)
if DEBUG: print("DEBUG: Resized X shape={}".format(X.shape))
if DEBUG: print("Non-normalized X=[{},{}] Y=[{},{}]".format(np.min(X),np.max(X),np.min(Y),np.max(Y)))
if DEBUG: print("Shapes X={} Y={}".format(X.shape,Y.shape))
if DEBUG: print("Normalized X=[{},{}] Y=[{},{}]".format(np.min(X),np.max(X),np.min(Y),np.max(Y)))
X=transformX(X)
Y,yParams=transformY(Y)
if DEBUG: print("Shapes X={} Y={}".format(X.shape,Y.shape))
if DEBUG: print("Normalized X=[{},{}] Y=[{},{}]".format(np.min(X),np.max(X),np.min(Y),np.max(Y)))
mod=None
try:
mod = load_model(MODEL_FILE_LOCATION)
except:
mod=getModel(X.shape[1:],Y.shape[1:])
mod.fit(x=X,y=Y,batch_size=4,epochs=EPOCHS,validation_split=0.0,shuffle=True)
mod.save(MODEL_FILE_LOCATION)
print("DEBUG: FINISHED SAVING MODEL as {}".format(MODEL_FILE_LOCATION))
yPred=mod.predict(X)
yPred=inverseTransformY(yPred,yParams)
Y=inverseTransformY(Y,yParams)
X=inverseTransformX(X)
#<<<<<<<<<<<<<<Inverting over
# if DEBUG: print("Y shape {}. yPred shape {}".format(yPred.shape,Y.shape))
# yPred=Y
for i in range(yPred.shape[0]):
if DEBUG:
print("Y, yPred")
print(Y[i])
print(yPred[i])
frame=np.copy(X[i])
for j in range(yPred.shape[1]):
#Predicted
cv2.circle(frame,center=(yPred[i,j,0],yPred[i,j,1]),radius=3,color=[255,0,0],thickness=3)
cv2.line(frame,pt1=(yPred[i,j,0],yPred[i,j,1]),pt2=(yPred[i,(j+1)%4,0],yPred[i,(j+1)%4,1]),color=[255,0,0],thickness=1)
#True
cv2.circle(frame,center=(Y[i,j,0],Y[i,j,1]),radius=3,color=[255,255,255],thickness=2)
cv2.line(frame,pt1=(Y[i,j,0],Y[i,j,1]),pt2=(Y[i,(j+1)%4,0],Y[i,(j+1)%4,1]),color=[255,255,255],thickness=1)
cv2.imshow("frame",frame)
cv2.waitKey(10)
vidIn=cv2.VideoCapture(TEST_VIDEO)
while True:
ret,frame=vidIn.read()
if not ret: break
frame=cv2.resize(frame,dsize=(299,299))
frame=np.reshape(frame,newshape=(1,299,299,3))
frame=transformX(frame)
yPred=mod.predict(frame)
yPred=inverseTransformY(yPred,yParams)
frame=inverseTransformX(frame)
frame=np.reshape(frame,newshape=(299,299,3))
if OPERATION=='test':
mod = load_model(MODEL_FILE_LOCATION)
if TEST_VIDEO.split(".")[-1]=='npz':
X,Y=load(TEST_VIDEO)
X=X[:250]
Y=Y[:250]
if DEBUG: print("Non-normalized X=[{},{}] Y=[{},{}]".format(np.min(X),np.max(X),np.min(Y),np.max(Y)))
for j in range(yPred.shape[1]):
cv2.circle(frame,center=(yPred[0,j,0],yPred[0,j,1]),radius=3,color=[255,0,0],thickness=3)
cv2.line(frame,pt1=(yPred[0,j,0],yPred[0,j,1]),pt2=(yPred[0,(j+1)%4,0],yPred[0,(j+1)%4,1]),color=[255,255,255],thickness=1)
cv2.imshow("frame",frame)
cv2.waitKey(10)
X=transformX(X)
Y,yParams=transformY(Y)
if DEBUG: print("Normalized X=[{},{}] Y=[{},{}]".format(np.min(X),np.max(X),np.min(Y),np.max(Y)))
yPred=mod.predict(X)
Y=inverseTransformY(Y,[1080.0,1920.0])
yPred=inverseTransformY(yPred,[1080.0,1920.0])
X=inverseTransformX(X)
for i in range(yPred.shape[0]):
frame=X[i]
for j in range(yPred.shape[1]):
#Predicted
cv2.circle(frame,center=(yPred[i,j,0],yPred[i,j,1]),radius=3,color=[255,0,0],thickness=3)
cv2.line(frame,pt1=(yPred[i,j,0],yPred[i,j,1]),pt2=(yPred[i,(j+1)%4,0],yPred[i,(j+1)%4,1]),color=[255,0,0],thickness=1)
#True
cv2.circle(frame,center=(Y[i,j,0],Y[i,j,1]),radius=3,color=[255,255,255],thickness=2)
cv2.line(frame,pt1=(Y[i,j,0],Y[i,j,1]),pt2=(Y[i,(j+1)%4,0],Y[i,(j+1)%4,1]),color=[255,255,255],thickness=1)
cv2.imshow("frame",frame)
cv2.waitKey(100)
elif TEST_VIDEO.split(".")[-1]=='mp4':
vidIn=cv2.VideoCapture(TEST_VIDEO)
while True:
ret,frame=vidIn.read()
if not ret: break
frame=cv2.resize(frame,dsize=(299,299),interpolation=cv2.INTER_LINEAR)
frame=transformX(frame)
yPred=mod.predict(frame)
yPred=inverseTransformY(yPred,[1920.0,1080.0])
frame=inverseTransformX(frame)
frame=np.reshape(frame,newshape=(299,299,3))
for j in range(yPred.shape[1]):
cv2.circle(frame,center=(yPred[0,j,0],yPred[0,j,1]),radius=3,color=[255,0,0],thickness=3)
cv2.line(frame,pt1=(yPred[0,j,0],yPred[0,j,1]),pt2=(yPred[0,(j+1)%4,0],yPred[0,(j+1)%4,1]),color=[255,255,255],thickness=1)
cv2.imshow("frame",frame)
cv2.waitKey(10)
......@@ -16,36 +16,66 @@ GAUSS = 0
DELTA = 5
print("python screen-detection-6-dataset-prep.py ./video/video-and-green.mp4 ./dataset-screen-detection.npz")
print("python screen-detection-6-dataset-prep.py ./video/video-and-green.mp4 ./dataset-screen-detection.npz red fast")
INPUT_FILE_NAME = sys.argv[1]
SAVE_FILE_NAME= sys.argv[2]
COLOR=sys.argv[3]
FAST=False
MAX_FRAMES=999999
try:
FAST=(sys.argv[4]=='fast')
MAX_FRAMES=int(sys.argv[5])
except:
print("Slow execution")
video = cv2.VideoCapture(INPUT_FILE_NAME)
corners=[]
isVideo=[]
isGreen=[]
mask=None
ret, frame = video.read()
xMid=int(frame.shape[1]/2)
yMid=int(frame.shape[0]/2)
maxGreen = -1
while video.isOpened():
if len(isVideo)==MAX_FRAMES:
break
ret, frame = video.read()
if ret:
lower = np.array([50, 80, 100])
upper = np.array([70, 255, 255])
hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
mask = cv2.inRange(hsv, lower, upper)
if COLOR=='green':
lower = np.array([50, 80, 100])
upper = np.array([70, 255, 255])
hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
mask = cv2.inRange(hsv, lower, upper)
elif COLOR=='red':
lower1 = np.array([0, 110, 150])
upper1 = np.array([10, 255, 255])
lower2 = np.array([170, 110, 150])
upper2 = np.array([180, 255, 255])
hsv = cv2.cvtColor(frame, cv2.COLOR_BGR2HSV)
mask1 = cv2.inRange(hsv, lower1, upper1)
mask2 = cv2.inRange(hsv, lower2, upper2)
mask = np.bitwise_or(mask1, mask2)
contours, hierarchy = cv2.findContours(mask, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)
areas = [cv2.contourArea(c) for c in contours]
if len(contours)>0:
max_cnt=contours[np.argmax(areas)]
m=np.reshape(max_cnt,(max_cnt.shape[0],2))
maxArea=np.max(areas)
isVideo.append(maxArea<frame.shape[0]*frame.shape[1]/16.0)
isGreen.append(maxArea>frame.shape[0]*frame.shape[1]/4.0)
isVideo.append(maxArea<frame.shape[0]*frame.shape[1]/32.0)
isGreen.append((maxArea>frame.shape[0]*frame.shape[1]/8.0) and (maxArea/(np.finfo(float).eps+maxGreen) > 0.98))
maxGreen = maxArea
if isGreen[-1]:
M = cv2.moments(max_cnt)
xMid = int(M["m10"] / M["m00"])
......@@ -60,59 +90,68 @@ while video.isOpened():
corners.append(np.array([m[lt],m[rt],m[rd],m[ld]]))
cv2.circle(frame,center=(xMid,yMid),radius=20,color=[255,0,0],thickness=10)
if not FAST: cv2.circle(frame,center=(xMid,yMid),radius=20,color=[255,0,0],thickness=10)
else:
corners.append(np.zeros((4,2)))
if isVideo[-1]:
cv2.drawContours(frame, max_cnt, -1, (0,0,255), 5)
cv2.circle(frame,center=(xMid,yMid),radius=300,color=[0,0,255],thickness=6)
if not FAST: cv2.drawContours(frame, max_cnt, -1, (0,0,255), 5)
if not FAST: cv2.circle(frame,center=(xMid,yMid),radius=300,color=[0,0,255],thickness=6)
elif isGreen[-1]:
cv2.drawContours(frame, max_cnt, -1, (255,0,0), 5)
if not FAST: cv2.drawContours(frame, max_cnt, -1, (255,0,0), 5)
else:
cv2.drawContours(frame, max_cnt, -1, (0,255,0), 5)
if not FAST: cv2.drawContours(frame, max_cnt, -1, (0,255,0), 5)
else:
corners.append(np.zeros((4,2)))
isVideo.append(True)
isGreen.append(False)
cv2.circle(frame,center=(xMid,yMid),radius=300,color=[0,0,255],thickness=6)
if not FAST: cv2.circle(frame,center=(xMid,yMid),radius=300,color=[0,0,255],thickness=6)
cv2.imshow("frame", cv2.resize(frame, (500,500)))
cv2.waitKey(20)
# if cv2.waitKey() & 0xFF == ord('q'):
# break
if not FAST: cv2.imshow("frame", cv2.resize(frame, (500,500),interpolation=cv2.INTER_LINEAR))
if not FAST: cv2.waitKey(2)
else:
corners=np.array(corners).astype(np.uint)
isVideo=np.array(isVideo)
isGreen=np.array(isGreen)
break
if FAST:
if len(isVideo) % 100 == 0:
print(len(isVideo), "frames processed")
video.release()
cv2.destroyAllWindows()
video = cv2.VideoCapture(INPUT_FILE_NAME)
X=[]
Y=[]
ret,frame=video.read()
ret,frame=video.read()
for i in range(len(corners)-1):
ret,frame=video.read()
if isGreen[i] and isVideo[i+1]:
X.append(frame)
X.append(cv2.resize(frame, (299, 299),interpolation=cv2.INTER_LINEAR))
Y.append(corners[i])
frameCopy=np.copy(frame)
yPred=corners[i]
for j in range(4):
cv2.circle(frameCopy,center=(yPred[j,0],yPred[j,1]),radius=3,color=[255,0,0],thickness=3)
cv2.line(frameCopy,pt1=(yPred[j,0],yPred[j,1]),pt2=(yPred[(j+1)%4,0],yPred[(j+1)%4,1]),color=[255,255,255],thickness=1)
if not FAST:
frameCopy=np.copy(frame)
yPred=corners[i]
for j in range(4):
cv2.circle(frameCopy,center=(yPred[j,0],yPred[j,1]),radius=3,color=[255,0,0],thickness=3)
cv2.line(frameCopy,pt1=(yPred[j,0],yPred[j,1]),pt2=(yPred[(j+1)%4,0],yPred[(j+1)%4,1]),color=[255,255,255],thickness=1)
cv2.imshow("frame",cv2.resize(frameCopy, (500, 500),interpolation=cv2.INTER_LINEAR))
cv2.waitKey(100)
if FAST:
if len(X)%100 == 0:
print(len(X), "frames picked")
cv2.imshow("frame",cv2.resize(frameCopy, (500, 500)))
cv2.waitKey(100)
X=np.array(X)
Y=np.array(Y)
np.savez(SAVE_FILE_NAME,X=X,Y=Y)
......
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment