Spaces:
Running
Running
| import os | |
| import random | |
| import numpy as np | |
| from multiprocessing import Process, Queue | |
| from . import image_processing as impro | |
| from . import mosaic,data | |
| class VideoLoader(object): | |
| """docstring for VideoLoader | |
| Load a single video(Converted to images) | |
| How to use: | |
| 1.Init VideoLoader as loader | |
| 2.Get data by loader.ori_stream | |
| 3.loader.next() to get next stream | |
| """ | |
| def __init__(self, opt, video_dir, test_flag=False): | |
| super(VideoLoader, self).__init__() | |
| self.opt = opt | |
| self.test_flag = test_flag | |
| self.video_dir = video_dir | |
| self.t = 0 | |
| self.n_iter = self.opt.M -self.opt.S*(self.opt.T+1) | |
| self.transform_params = data.get_transform_params() | |
| self.ori_load_pool = [] | |
| self.mosaic_load_pool = [] | |
| self.previous_pred = None | |
| feg_ori = impro.imread(os.path.join(video_dir,'origin_image','00001.jpg'),loadsize=self.opt.loadsize,rgb=True) | |
| feg_mask = impro.imread(os.path.join(video_dir,'mask','00001.png'),mod='gray',loadsize=self.opt.loadsize) | |
| self.mosaic_size,self.mod,self.rect_rat,self.feather = mosaic.get_random_parameter(feg_ori,feg_mask) | |
| self.startpos = [random.randint(0,self.mosaic_size),random.randint(0,self.mosaic_size)] | |
| self.loadsize = self.opt.loadsize | |
| #Init load pool | |
| for i in range(self.opt.S*self.opt.T): | |
| _ori_img = impro.imread(os.path.join(video_dir,'origin_image','%05d' % (i+1)+'.jpg'),loadsize=self.loadsize,rgb=True) | |
| _mask = impro.imread(os.path.join(video_dir,'mask','%05d' % (i+1)+'.png' ),mod='gray',loadsize=self.loadsize) | |
| _mosaic_img = mosaic.addmosaic_base(_ori_img, _mask, self.mosaic_size,0, self.mod,self.rect_rat,self.feather,self.startpos) | |
| _ori_img = data.random_transform_single_image(_ori_img,opt.finesize,self.transform_params) | |
| _mosaic_img = data.random_transform_single_image(_mosaic_img,opt.finesize,self.transform_params) | |
| self.ori_load_pool.append(self.normalize(_ori_img)) | |
| self.mosaic_load_pool.append(self.normalize(_mosaic_img)) | |
| self.ori_load_pool = np.array(self.ori_load_pool) | |
| self.mosaic_load_pool = np.array(self.mosaic_load_pool) | |
| #Init frist stream | |
| self.ori_stream = self.ori_load_pool [np.linspace(0, (self.opt.T-1)*self.opt.S,self.opt.T,dtype=np.int64)].copy() | |
| self.mosaic_stream = self.mosaic_load_pool[np.linspace(0, (self.opt.T-1)*self.opt.S,self.opt.T,dtype=np.int64)].copy() | |
| # stream B,T,H,W,C -> B,C,T,H,W | |
| self.ori_stream = self.ori_stream.reshape (1,self.opt.T,opt.finesize,opt.finesize,3).transpose((0,4,1,2,3)) | |
| self.mosaic_stream = self.mosaic_stream.reshape(1,self.opt.T,opt.finesize,opt.finesize,3).transpose((0,4,1,2,3)) | |
| #Init frist previous frame | |
| self.previous_pred = self.ori_load_pool[self.opt.S*self.opt.N-1].copy() | |
| # previous B,C,H,W | |
| self.previous_pred = self.previous_pred.reshape(1,opt.finesize,opt.finesize,3).transpose((0,3,1,2)) | |
| def normalize(self,data): | |
| ''' | |
| normalize to -1 ~ 1 | |
| ''' | |
| return (data.astype(np.float32)/255.0-0.5)/0.5 | |
| def anti_normalize(self,data): | |
| return np.clip((data*0.5+0.5)*255,0,255).astype(np.uint8) | |
| def next(self): | |
| # random | |
| if np.random.random()<0.05: | |
| self.startpos = [random.randint(0,self.mosaic_size),random.randint(0,self.mosaic_size)] | |
| if np.random.random()<0.02: | |
| self.transform_params['rate']['crop'] = [np.random.random(),np.random.random()] | |
| if np.random.random()<0.02: | |
| self.loadsize = np.random.randint(self.opt.finesize,self.opt.loadsize) | |
| if self.t != 0: | |
| self.previous_pred = None | |
| self.ori_load_pool [:self.opt.S*self.opt.T-1] = self.ori_load_pool [1:self.opt.S*self.opt.T] | |
| self.mosaic_load_pool[:self.opt.S*self.opt.T-1] = self.mosaic_load_pool[1:self.opt.S*self.opt.T] | |
| #print(os.path.join(self.video_dir,'origin_image','%05d' % (self.opt.S*self.opt.T+self.t)+'.jpg')) | |
| _ori_img = impro.imread(os.path.join(self.video_dir,'origin_image','%05d' % (self.opt.S*self.opt.T+self.t)+'.jpg'),loadsize=self.loadsize,rgb=True) | |
| _mask = impro.imread(os.path.join(self.video_dir,'mask','%05d' % (self.opt.S*self.opt.T+self.t)+'.png' ),mod='gray',loadsize=self.loadsize) | |
| _mosaic_img = mosaic.addmosaic_base(_ori_img, _mask, self.mosaic_size,0, self.mod,self.rect_rat,self.feather,self.startpos) | |
| _ori_img = data.random_transform_single_image(_ori_img,self.opt.finesize,self.transform_params) | |
| _mosaic_img = data.random_transform_single_image(_mosaic_img,self.opt.finesize,self.transform_params) | |
| _ori_img,_mosaic_img = self.normalize(_ori_img),self.normalize(_mosaic_img) | |
| self.ori_load_pool [self.opt.S*self.opt.T-1] = _ori_img | |
| self.mosaic_load_pool[self.opt.S*self.opt.T-1] = _mosaic_img | |
| self.ori_stream = self.ori_load_pool [np.linspace(0, (self.opt.T-1)*self.opt.S,self.opt.T,dtype=np.int64)].copy() | |
| self.mosaic_stream = self.mosaic_load_pool[np.linspace(0, (self.opt.T-1)*self.opt.S,self.opt.T,dtype=np.int64)].copy() | |
| # stream B,T,H,W,C -> B,C,T,H,W | |
| self.ori_stream = self.ori_stream.reshape (1,self.opt.T,self.opt.finesize,self.opt.finesize,3).transpose((0,4,1,2,3)) | |
| self.mosaic_stream = self.mosaic_stream.reshape(1,self.opt.T,self.opt.finesize,self.opt.finesize,3).transpose((0,4,1,2,3)) | |
| self.t += 1 | |
| class VideoDataLoader(object): | |
| """VideoDataLoader""" | |
| def __init__(self, opt, videolist, test_flag=False): | |
| super(VideoDataLoader, self).__init__() | |
| self.videolist = [] | |
| self.opt = opt | |
| self.test_flag = test_flag | |
| for i in range(self.opt.n_epoch): | |
| self.videolist += videolist.copy() | |
| random.shuffle(self.videolist) | |
| self.each_video_n_iter = self.opt.M -self.opt.S*(self.opt.T+1) | |
| self.n_iter = len(self.videolist)//self.opt.load_thread//self.opt.batchsize*self.each_video_n_iter*self.opt.load_thread | |
| self.queue = Queue(self.opt.load_thread) | |
| self.ori_stream = np.zeros((self.opt.batchsize,3,self.opt.T,self.opt.finesize,self.opt.finesize),dtype=np.float32)# B,C,T,H,W | |
| self.mosaic_stream = np.zeros((self.opt.batchsize,3,self.opt.T,self.opt.finesize,self.opt.finesize),dtype=np.float32)# B,C,T,H,W | |
| self.previous_pred = np.zeros((self.opt.batchsize,3,self.opt.finesize,self.opt.finesize),dtype=np.float32) | |
| self.load_init() | |
| def load(self,videolist): | |
| for load_video_iter in range(len(videolist)//self.opt.batchsize): | |
| iter_videolist = videolist[load_video_iter*self.opt.batchsize:(load_video_iter+1)*self.opt.batchsize] | |
| videoloaders = [VideoLoader(self.opt,os.path.join(self.opt.dataset,iter_videolist[i]),self.test_flag) for i in range(self.opt.batchsize)] | |
| for each_video_iter in range(self.each_video_n_iter): | |
| for i in range(self.opt.batchsize): | |
| self.ori_stream[i] = videoloaders[i].ori_stream | |
| self.mosaic_stream[i] = videoloaders[i].mosaic_stream | |
| if each_video_iter == 0: | |
| self.previous_pred[i] = videoloaders[i].previous_pred | |
| videoloaders[i].next() | |
| if each_video_iter == 0: | |
| self.queue.put([self.ori_stream.copy(),self.mosaic_stream.copy(),self.previous_pred]) | |
| else: | |
| self.queue.put([self.ori_stream.copy(),self.mosaic_stream.copy(),None]) | |
| def load_init(self): | |
| ptvn = len(self.videolist)//self.opt.load_thread #pre_thread_video_num | |
| for i in range(self.opt.load_thread): | |
| p = Process(target=self.load,args=(self.videolist[i*ptvn:(i+1)*ptvn],)) | |
| p.daemon = True | |
| p.start() | |
| def get_data(self): | |
| return self.queue.get() |