Ali Mohsin
feat: Add virtual try-on system components including DensePose, SMPL, and pix2pixHD models, rendering, and utilities.
5db43ff
| import cv2 | |
| import numpy as np | |
| from threading import Thread | |
| import queue | |
| import time | |
| import imageio | |
| import os | |
| import gc | |
| class MultithreadVideoWriter: | |
| def __init__(self, outvid, fps=30,quality=5): | |
| self.outvid = outvid | |
| self.quality=quality | |
| self.fps = fps | |
| self.vid = None | |
| self.buff_size = 160 | |
| self.stop = False | |
| self.finished = False | |
| self.quit = False | |
| self.frame_queue = queue.Queue(maxsize=self.buff_size) | |
| self.t = Thread(target=self.stream2video, args=()) | |
| self.t.daemon = True | |
| self.t.start() | |
| def append(self, image, isRGB=False): | |
| if not isRGB: | |
| image = image[:, :, [2, 1, 0]] | |
| self.frame_queue.put(image) | |
| def stream2video(self): | |
| while True: | |
| if self.stop and self.frame_queue.empty(): | |
| self.finished = True | |
| break | |
| try: | |
| frame = self.frame_queue.get(timeout=1) | |
| if self.vid is None: | |
| self.vid = imageio.get_writer(self.outvid, fps=self.fps, codec='libx264',quality=self.quality) | |
| self.vid.append_data(frame) | |
| except queue.Empty: | |
| continue | |
| while True: | |
| time.sleep(0.1) | |
| if self.quit: | |
| break | |
| def close(self): | |
| self.quit = True | |
| self.t.join() | |
| if self.vid is not None: | |
| self.vid.close() | |
| def make_video(self): | |
| self.stop = True | |
| while not self.finished: | |
| time.sleep(0.1) | |
| if self.vid is not None: | |
| self.vid.close() | |
| # Force garbage collection | |
| gc.collect() | |