import os

import cv2
import numpy as np

dirpath = 'BadApple_pro_frames'
newdir = 'BadApple_only_edges'

if not os.path.exists(newdir):
    os.makedirs(newdir)

order = [int(i.strip(".jpg")) for i in os.listdir(dirpath) if i.endswith(".jpg")]
jpglist = [f"{i}.jpg" for i in sorted(order)]  # 直接读取可能非顺序帧

for i, jpg in enumerate(jpglist):
    old = dirpath + f'/{jpg}'
    img = cv2.imread(old)  # 返回的是numpy.array对象

    edge = cv2.Canny(img, 20, 255)
    invert = np.abs((255 * np.ones((edge.shape)) - edge)).clip(0, 255)

    # cv2.imshow('img', invert)
    new = newdir + f'/{jpg}'
    cv2.imwrite(new, invert)
    print(f'{i + 1} / {len(jpglist)}')