screen detectors

screen-detection/screen-detection-5-nn.py

@@ -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)

screen-detection/screen-detection-6-dataset-prep.py

@@ -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)