import cv2
import numpy as np
import glob


def clean():
    filenames = []
    for i in range(300):
        if i<10:
            filenames.append(r'0000'+str(i))
        elif i<100:
            filenames.append(r'000'+str(i))
        else:
            filenames.append(r'00'+str(i))
            
    for number in range(1,5):
        video = "video" + str(number)
        imageadr = "C:\\Users\\Dargon\\Desktop\\SegTrackv2 changed\\SegTrackv2\\JPEGImages\\"+ video +"\\"
        img_array = []
        for i in filenames:
            img = cv2.imread(imageadr+i+".jpg")
            if img is None:
               print("konec")
               break
            height, width, layers = img.shape
            size = (width,height)
            img_array.append(img)
         
         
        out = cv2.VideoWriter("legit" + video+'.avi',cv2.VideoWriter_fourcc(*'DIVX'), 10, size)
         
        for i in range(len(img_array)):
            out.write(img_array[i])
        out.release()



def segmented():
    filenames = []
    for i in range(300):
        if i<10:
            filenames.append(r'0000'+str(i))
        elif i<100:
            filenames.append(r'000'+str(i))
        else:
            filenames.append(r'00'+str(i))

    for number in range(1,5):
        img_array = []
        video = "video" + str(number)
        imageadr = "C:\\Users\\Dargon\\Desktop\\SegTrackv2 changed\\SegTrackv2\\JPEGImages\\"+ video +"\\"
        maskadr = "C:\\Users\\Dargon\\Desktop\\segtrackSTCN\\output\\"+ video +"\\"
        
        for i in filenames:
            img = cv2.imread(imageadr+i+".jpg")
            #print(imageadr+i+".jpg")
            mask = cv2.imread(maskadr+i+".png")
            #cv2.imshow("asa",img)
            if mask is None:
                print("konec")
                break
            segmented = crop(img,mask)
            img_array.append(segmented)
            
        height, width, layers = img_array[0].shape
        size = (width,height)
        

        out = cv2.VideoWriter(video+'.avi',cv2.VideoWriter_fourcc(*'DIVX'), 6, size)
         
        for i in range(len(img_array)):
            out.write(img_array[i])
        out.release()

def crop(image = cv2.imread("00049.jpg"),mask = cv2.imread("00049.png"),threshold=20):
    mask=cv2.cvtColor(mask,cv2.COLOR_BGR2GRAY)
    mask = cv2.threshold(mask, threshold, 255, cv2.THRESH_BINARY)[1]
    mask=cv2.cvtColor(mask,cv2.COLOR_GRAY2BGR)
    mask_out=cv2.subtract(mask,image)
    mask_out=cv2.subtract(mask,mask_out)
    return mask_out
