Cycle gan webcam

https://github.com/junyanz/pytorch-CycleGAN-and-pix2pix

junyanz/pytorch-CycleGAN-and-pix2pix

Cyclegan을 사용하는 것은 test폴더에 이미지를 넣고 모델을 통해 결과물을 저장하는 방식으로 실행이 된다.

 

웹캠을 이용하여 웹캠으로 들어온 영상을 바로 변환하여 사용을 하고 싶을 수도 있다.

 

웹캠으로 받은 이미지를 사용하는 코드를 올리고자 한다.

 

import os
from options.test_options import TestOptions
from data import create_dataset
from models import create_model
from util.visualizer import save_images
# from util import html
import cv2
import numpy as np
import torch
from util import util
#import matplotlib.pyplot as plt
from PIL import Image
import time



if __name__ == '__main__':
    opt = TestOptions().parse()  # get test options
    # hard-code some parameters for test
    opt.num_threads = 0   # test code only supports num_threads = 1
    opt.batch_size = 1    # test code only supports batch_size = 1
    opt.serial_batches = True  # disable data shuffling; comment this line if results on randomly chosen images are needed.
    opt.no_flip = True    # no flip; comment this line if results on flipped images are needed.
    opt.display_id = -1   # no visdom display; the test code saves the results to a HTML file.
    dataset = create_dataset(opt)  # create a dataset given opt.dataset_mode and other options
    model = create_model(opt)      # create a model given opt.model and other options
    model.setup(opt)               # regular setup: load and print networks; create schedulers

    if opt.eval:
        model.eval()
    for i, data in enumerate(dataset):
        if i >= opt.num_test:  # only apply our model to opt.num_test images.
            break
    data = next(iter(dataset),None)
    
    capture = cv2.VideoCapture(0)
    capture.set(cv2.CAP_PROP_FRAME_WIDTH, 1280)
    capture.set(cv2.CAP_PROP_FRAME_HEIGHT, 720)
    # plt.ion()
    if (capture.isOpened() == False): 
        print("Unable to read camera feed")    
    while(True):
        ret, frame = capture.read()        
        if ret == True:            
            cv2.imshow("VideoFrame", frame)                       
            if cv2.waitKey(1) & 0xFF == ord('q'):
                break
        else :
            break
        frame = cv2.cvtColor(frame,cv2.COLOR_BGR2RGB)
        #plt.imshow(frame, interpolation='nearest')        
        frame = np.asarray([frame])
        frame = np.transpose(frame,(0,3,1,2,))
        #float to tnnsor을 통해 입력받은 frame를 dataloader로 전송해준다.
        data['A'] = torch.FloatTensor(frame)   
        start = time.time()        
        #dataloader을 통해 model에 input해줌
        model.set_input(data)  
        # run inference
        model.test()           
        # 실행시간 측정
        torch.cuda.synchronize()
        end = time.time() - start
        print(end)
        # get image results
        visuals = model.fake.data          
        #만들어진 tensor을 image로 변환해줌
        #함수는 util에 tensor2im으로 함수 선언이 되어있음.
        visuals = util.tensor2im(visuals)
        
        visuals = np.uint8(visuals)  
        #opencv를 사용해서 출력할 경우 BGR2RGB를 해주어야함
        visuals = cv2.cvtColor(np.array(visuals), cv2.COLOR_BGR2RGB)        
        visuals = cv2.resize(visuals, (512, 512))
        
        cv2.imshow("VideoFrame22", visuals)
    
    capture.release()
    cv2.destroyAllWindows()

 

test대신 사용을 실행하면 되고

 

test폴더에는 아무 이미지 한개만 준비하면된다.

 

물론 컴퓨팅 사양이 부족하면 fps가 많이 저하된다.

 

2080GTX를 기준으로 약 30fps가 나왔으니 그 밑의 사양의 gpu는 자연스럽게 보이지 않을 것 같다.