You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
129 lines
4.8 KiB
129 lines
4.8 KiB
import argparse
|
|
from MeshRenderingPipeline import MeshRenderingPipeline
|
|
from watchdog.observers import Observer
|
|
from watchdog.events import FileSystemEventHandler
|
|
import threading
|
|
import time
|
|
import random
|
|
|
|
class ImageHandler(FileSystemEventHandler):
|
|
def __init__(self, mesh_pipeline):
|
|
self.mesh_pipeline = mesh_pipeline
|
|
|
|
def on_created(self, event):
|
|
if event.is_directory:
|
|
return
|
|
if event.src_path.endswith(('.png', '.jpg', '.webp')):
|
|
print(f"New image detected: {event.src_path}")
|
|
self.mesh_pipeline.process_image(event.src_path, 0, 1, None)
|
|
|
|
class App:
|
|
def __init__(self, args):
|
|
self.args = args
|
|
self.running = False
|
|
self.stop_event = threading.Event()
|
|
self.thread = None
|
|
self.input_thread = None
|
|
self.observer = None
|
|
self.mesh_pipeline = MeshRenderingPipeline(
|
|
args.config,
|
|
args.input_path,
|
|
args.output_path,
|
|
args.diffusion_steps,
|
|
args.seed,
|
|
args.scale,
|
|
args.distance,
|
|
args.view,
|
|
args.no_rembg,
|
|
args.export_texmap,
|
|
args.save_video,
|
|
args.gltf,
|
|
args.remove,
|
|
args.gltf_path,
|
|
args.local
|
|
)
|
|
|
|
def start_generation(self):
|
|
self.running = True
|
|
self.stop_event.clear()
|
|
self.thread = threading.Thread(target=self._generate_objects)
|
|
self.thread.start()
|
|
|
|
def stop_generation(self):
|
|
self.stop_event.set()
|
|
self.running = False
|
|
if self.thread:
|
|
self.thread.join()
|
|
self.stop_observer()
|
|
|
|
def monitor_new_images(self):
|
|
event_handler = ImageHandler(self.mesh_pipeline)
|
|
self.observer = Observer()
|
|
self.observer.schedule(event_handler, self.args.input_path, recursive=False)
|
|
self.observer.start()
|
|
print("\nWaiting for new image..")
|
|
while not self.stop_event.is_set():
|
|
time.sleep(1)
|
|
self.stop_observer() # Ensure the observer is stopped on interruption
|
|
|
|
|
|
def stop_observer(self):
|
|
if self.observer:
|
|
self.observer.stop()
|
|
self.observer.join()
|
|
self.observer = None
|
|
|
|
def _generate_objects(self):
|
|
while not self.stop_event.is_set():
|
|
self.run_pipeline()
|
|
time.sleep(1)
|
|
|
|
|
|
def run_pipeline(self):
|
|
self.mesh_pipeline.run_pipeline_sequence(self.stop_event)
|
|
self.monitor_new_images()
|
|
|
|
def run(self):
|
|
self.input_thread = threading.Thread(target=self._handle_input)
|
|
self.input_thread.start()
|
|
self.input_thread.join()
|
|
|
|
def _handle_input(self):
|
|
while True:
|
|
command = input("\nEnter a command, <start> <stop> <exit>: ")
|
|
if command.lower() == 'exit':
|
|
print("Exiting the program.")
|
|
self.stop_generation()
|
|
break
|
|
elif command.lower() == 'start':
|
|
if not self.running:
|
|
print("\nStarting continuous generation.")
|
|
self.start_generation()
|
|
else:
|
|
print("\nGeneration already running.")
|
|
elif command.lower() == 'stop':
|
|
print("\nStopping continuous generation.")
|
|
self.stop_generation()
|
|
|
|
|
|
if __name__ == '__main__':
|
|
parser = argparse.ArgumentParser()
|
|
parser.add_argument('config', type=str, help='Path to config file.')
|
|
parser.add_argument('input_path', type=str, help='Path to input image or directory.')
|
|
parser.add_argument('--output_path', type=str, default='outputs/', help='Output directory.')
|
|
parser.add_argument('--diffusion_steps', type=int, default=75, help='Denoising Sampling steps.')
|
|
parser.add_argument('--seed', type=int, default=42, help='Random seed for sampling.')
|
|
parser.add_argument('--scale', type=float, default=1.0, help='Scale of generated object.')
|
|
parser.add_argument('--distance', type=float, default=4.5, help='Render distance.')
|
|
parser.add_argument('--view', type=int, default=6, choices=[4, 6], help='Number of input views.')
|
|
parser.add_argument('--no_rembg', action='store_true', help='Do not remove input background.')
|
|
parser.add_argument('--export_texmap', action='store_true', help='Export a mesh with texture map.')
|
|
parser.add_argument('--save_video', action='store_true', help='Save a circular-view video.')
|
|
parser.add_argument('--gltf', action='store_true', help='Export a gtlf file.')
|
|
parser.add_argument('--remove', action='store_true', help='Removes obj, mtl, texmap, nv files.')
|
|
parser.add_argument('--gltf_path', type=str, default='C:/Users/caile/Desktop/InstantMesh/ex', help='Output directory.')
|
|
parser.add_argument('--local', action='store_true', help='Dont send over a network')
|
|
args = parser.parse_args()
|
|
|
|
app = App(args)
|
|
app.run()
|