Ali Mohsin
feat: Add virtual try-on system components including DensePose, SMPL, and pix2pixHD models, rendering, and utilities.
5db43ff
| import time | |
| import cv2 | |
| import numpy as np | |
| import torch | |
| from threading import Thread | |
| import queue | |
| class MultithreadVideoLoader: | |
| def __init__(self, path,max_height=None): | |
| self.stop = False | |
| if isinstance(path,list): | |
| self.path_list = path | |
| else: | |
| self.path_list=[path] | |
| self.cap_list=[cv2.VideoCapture(p) for p in self.path_list] | |
| self.fps_list = [cap.get(cv2.CAP_PROP_FPS) for cap in self.cap_list] | |
| self.nframe_list=[int(cap.get(cv2.CAP_PROP_FRAME_COUNT)) for cap in self.cap_list] | |
| self.nframe=self.__len__() | |
| self.max_height=max_height | |
| self.buff_size=64 | |
| self.frame_queue=queue.Queue(maxsize=self.buff_size) | |
| self.t=Thread(target=self.load_buff,args=()) | |
| self.t.daemon=True | |
| self.t.start() | |
| def get_fps(self): | |
| return self.fps_list[0] | |
| def load_buff(self): | |
| i=0 | |
| while True: | |
| if self.stop: | |
| break | |
| if i>self.nframe-1: | |
| break | |
| ret, frame = self.get_item(i) | |
| if not ret: | |
| break | |
| self.frame_queue.put(frame) | |
| i+=1 | |
| while True: | |
| time.sleep(1) | |
| if self.stop: | |
| break | |
| def global2local(self,i): | |
| n_cap=0 | |
| assert i<=self.__len__()-1 | |
| for n in self.nframe_list: | |
| if i<=n-1: | |
| return n_cap, i | |
| else: | |
| i=i-n | |
| n_cap+=1 | |
| def cap(self): | |
| try: | |
| result = self.frame_queue.get(timeout=1) | |
| except queue.Empty: | |
| result = None | |
| return result | |
| def get_item(self, idx): | |
| n_cap,n_frame=self.global2local(idx) | |
| #self.cap_list[n_cap].set(cv2.CAP_PROP_POS_FRAMES, n_frame - 1) | |
| #sequential loading is faster! | |
| res, frame = self.cap_list[n_cap].read() | |
| if not res: | |
| return res, frame | |
| h=frame.shape[0] | |
| w=frame.shape[1] | |
| need_resize=False | |
| if self.max_height is not None: | |
| if h>self.max_height: | |
| need_resize=True | |
| if need_resize: | |
| tmp = h | |
| h = self.max_height | |
| w = int(w* h / tmp) | |
| frame=cv2.resize(frame,(w,h)) | |
| return res, frame | |
| def __len__(self): | |
| nframes=0 | |
| for n in self.nframe_list: | |
| nframes+=n | |
| return nframes | |
| def close(self): | |
| self.stop=True | |
| while not self.frame_queue.empty(): | |
| self.frame_queue.get() | |
| self.t.join() | |
| for cap in self.cap_list: | |
| cap.release() | |