博客

import numpy as np
import cv2 as cv
cap = cv.VideoCapture('/root/autodl-tmp/code/BackgroundMattingV2/src.mp4')

while cap.isOpened():
    ret, frame = cap.read()
    # if frame is read correctly ret is True
    if not ret:
        print("Can't receive frame (stream end?). Exiting ...")
        break
    # cv2.imshow("frame", frame)
    # frame = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
    # cv2.imshow("frame", frame)
    cv.imwrite("src.jpg", frame)
    break 
    # oneframereader = OneFrameReader(frame, transform=ToTensor())
    # with torch.no_grad():
    #     for src in DataLoader(oneframereader):                     # RGB tensor normalized to 0 ~ 1.
    #         # print(torch.max(src))
    #         # print(src.shape) # torch.Size([1, 3, 1280, 720])
    #         fgr, pha, *rec = model(src.cuda(), *rec, downsample_ratio)  # Cycle the recurrent states.
    #         com = fgr * pha + bgr * (1 - pha)              # Composite to green background. 
    #         # print(torch.max(com))
    #         # print(com.shape)
    #         com2 = com.mul(255).byte()
    #         com2 = torch.squeeze(com2, 0)
    #         # print(com2.shape)
    #         com2 = torch.permute(com2, (1,2,0))
    #         com2 = com2.detach().cpu().numpy()
    #         com2 = com2[..., ::-1] # h w c rgb2bgr
    #         # print(com2.shape)
    #         # print(type(com2))
    #         # print(np.max(com2))
    #         # print(com2.dtype)
    #         # print(com2.shape)
    #         # cv2.imshow("com2", com2)
    #         # cv2.waitKey(0)

    #         # plt.imshow(com2)
    #         writer.write(com)                              # Write frame.
# plt.show()
cap.release()
cv.destroyAllWindows()