Spaces:
Running
Running
| # ===================================================== | |
| # Apckeyl_RealESRGAN | |
| # Version 2.0 | |
| # batch_processor.py | |
| # ===================================================== | |
| from image_upscaler import upscale_image | |
| from priority_queue import PriorityQueue | |
| class BatchProcessor: | |
| """ | |
| Batch Processor. | |
| Version 2.0 | |
| Использует Priority Queue. | |
| """ | |
| def __init__(self): | |
| self.results = [] | |
| self.queue = PriorityQueue() | |
| # ------------------------------------------------- | |
| def clear(self): | |
| self.results.clear() | |
| self.queue.clear() | |
| # ------------------------------------------------- | |
| def add_task( | |
| self, | |
| image, | |
| mode, | |
| quality, | |
| denoise, | |
| owner=False, | |
| ): | |
| self.queue.add( | |
| { | |
| "image": image, | |
| "mode": mode, | |
| "quality": quality, | |
| "denoise": denoise, | |
| }, | |
| owner=owner, | |
| ) | |
| # ------------------------------------------------- | |
| def process_queue(self): | |
| self.results.clear() | |
| while self.queue.size() > 0: | |
| task = self.queue.get_next() | |
| result = upscale_image( | |
| image=task["image"], | |
| mode=task["mode"], | |
| quality=task["quality"], | |
| denoise=task["denoise"], | |
| ) | |
| self.results.append( | |
| result | |
| ) | |
| return self.results | |
| # ------------------------------------------------- | |
| def process( | |
| self, | |
| images, | |
| mode, | |
| quality, | |
| denoise, | |
| owner=False, | |
| ): | |
| self.clear() | |
| for image in images: | |
| self.add_task( | |
| image=image, | |
| mode=mode, | |
| quality=quality, | |
| denoise=denoise, | |
| owner=owner, | |
| ) | |
| return self.process_queue() | |
| # ------------------------------------------------- | |
| def process_with_limit( | |
| self, | |
| images, | |
| mode, | |
| quality, | |
| denoise, | |
| owner=False, | |
| max_images=10, | |
| ): | |
| if images is None: | |
| return [] | |
| images = images[:max_images] | |
| return self.process( | |
| images=images, | |
| mode=mode, | |
| quality=quality, | |
| denoise=denoise, | |
| owner=owner, | |
| ) | |
| # ------------------------------------------------- | |
| def count(self): | |
| return len( | |
| self.results | |
| ) | |
| # ------------------------------------------------- | |
| def queue_size(self): | |
| return self.queue.size() |