instant-mesh/application.py

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()