import argparse
import threading
import time
import random
import open3d as o3d
import open3d . visualization . gui as gui
import open3d . visualization . rendering as rendering
from ShapeGenerator import ShapeGenerator
from TextGenerator import TextGenerator
# command example
# python app.py --output_dir /mnt/c/Users/caile/Desktop/output
class GCApp :
def __init__ ( self , output_dir , batch_size , step_size , guidance_scale , vis ) :
self . output_dir = output_dir
self . obj_gen = ShapeGenerator ( self . output_dir , batch_size , step_size , guidance_scale )
self . running = False
self . stop_event = threading . Event ( )
self . thread = None
self . vis = vis
self . waste_items = [
" Plastic bottle " , " Aluminum can " , " Glass bottle " , " Food wrapper " ,
" Cardboard box " , " Paper bag " , " Plastic bag " , " Electronics " ,
" Old smartphone " , " Broken TV " , " Computer parts " , " Batteries " ,
" Light bulbs " , " Old furniture " , " Styrofoam cup " , " Food container " ,
" Takeout box " , " Cigarette butts " , " Plastic utensils " , " Straws " ,
" Bottle caps " , " Rubber tires " , " Broken toys " , " Old clothes " ,
" Shoes " , " Wooden pallets " , " Paint cans " , " Cleaning products " ,
" Old appliances " , " Wires " , " Cables " , " Extension cords " ,
" Old magazines " , " Newspapers " , " Scrap metal " , " Construction debris " ,
" Yard waste " , " Grass clippings " , " Leaves " , " Old mattresses " ,
" Carpeting " , " Food scraps " , " Pet waste " , " Diapers " ,
" Sanitary products " , " Receipts " , " Plastic wrap " , " Packing peanuts " ,
" Ice cream containers " , " Fast food containers " , " Takeaway cups " ,
" Clamshell packaging " , " Plastic film " , " Broken glass " ,
" Old books " , " VCR tapes " , " CDs " , " DVDs " ,
" Game consoles " , " Remote controls " , " Ink cartridges " ,
" Toner cartridges " , " Old tools " , " Gardening tools " ,
" Bike parts " , " Fishing gear " , " Beach toys " , " Pool floats " ,
" Old bicycles " , " Skateboards " , " Surfboards " , " Helmets " ,
" Used batteries " , " Old jewelry " , " Keyboards " , " Mice (computer) " ,
" Speakers " , " Old cameras " , " Projectors " , " Printers " ,
" Scanners " , " Shredded paper " , " Bubble wrap " , " Plastic sheeting " ,
" Tarps " , " Old car parts " , " Motor oil containers " ,
" Propane tanks " , " Oil filters " , " Windshield wipers " ,
" Car batteries " , " Antifreeze containers " , " Used tires " ,
" Old propane tanks " , " Scrap wood " , " Broken furniture " ,
" Old carpets " , " Leather scraps " , " Textile waste " ,
" Compostable waste "
]
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 ( )
def get_random_item_prompt ( self ) :
return random . choice ( self . waste_items )
def add_to_scene ( self , mesh , prompt ) :
scene = self . vis . get_scene ( )
self . vis . update_label ( prompt )
scene . scene . clear_geometry ( )
scene . scene . add_geometry ( " name " , mesh , rendering . MaterialRecord ( ) )
def _generate_objects ( self ) :
while not self . stop_event . is_set ( ) :
mesh , prompt = self . obj_gen . generate_object ( self . get_random_item_prompt ( ) )
self . add_to_scene ( mesh , prompt )
time . sleep ( 1 )
def run ( self ) :
self . obj_gen . run ( )
while True :
command = input ( " Enter a command, <start> <stop> <generate (prompt)>: " )
if command . lower ( ) == ' exit ' :
print ( " Exiting the program. " )
self . stop_generation ( )
break
elif command . lower ( ) == ' start ' :
if not self . running :
print ( " Starting continuous generation. " )
self . start_generation ( )
else :
print ( " Generation already running. " )
elif command . lower ( ) == ' stop ' :
print ( " Stopping continuous generation. " )
self . stop_generation ( )
elif command . startswith ( ' generate ' ) :
shape = command [ len ( ' generate ' ) : ]
mesh , prompt = self . obj_gen . generate_object ( shape )
self . add_to_scene ( mesh , prompt )
else :
print ( " Unknown command. " )
class VisApp :
def __init__ ( self ) :
self . _id = 0
self . window = gui . Application . instance . create_window (
" garbage collector " , 1024 , 768 )
self . obj_label = None
self . layout = None
self . scene = None
self . create_scene ( )
self . create_layout ( )
def create_layout ( self ) :
self . layout = gui . Vert ( 0 , gui . Margins ( 10 , 10 , 10 , 10 ) )
self . obj_label = gui . Label ( " None " )
self . layout . add_child ( self . obj_label )
self . window . set_on_layout ( self . _on_layout )
self . window . add_child ( self . layout )
def create_scene ( self ) :
self . scene = gui . SceneWidget ( )
self . scene . scene = rendering . Open3DScene ( self . window . renderer )
self . scene . scene . set_background ( [ 1 , 1 , 1 , 1 ] )
self . scene . scene . scene . set_sun_light (
[ - 1 , - 1 , - 1 ] , # direction
[ 1 , 1 , 1 ] , # color
100000 ) # intensity
self . scene . scene . scene . enable_sun_light ( True )
bbox = o3d . geometry . AxisAlignedBoundingBox ( [ - 5 , - 5 , - 5 ] ,
[ 5 , 5 , 5 ] )
self . scene . setup_camera ( 20 , bbox , [ 0 , 0 , 0 ] )
self . window . add_child ( self . scene )
def _on_layout ( self , layout_context ) :
# The on_layout callback should set the frame (position + size) of every
# child correctly. After the callback is done the window will layout
# the grandchildren.
r = self . window . content_rect
self . scene . frame = r
width = 17 * layout_context . theme . font_size
height = min (
r . height ,
self . layout . calc_preferred_size (
layout_context , gui . Widget . Constraints ( ) ) . height )
self . layout . frame = gui . Rect ( r . get_right ( ) - width , r . y , width ,
height )
def get_scene ( self ) :
return self . scene
def update_label ( self , name ) :
self . obj_label . text = " Object: " + name
def main ( output_dir , batch_size , step_size , guidance_scale ) :
gui . Application . instance . initialize ( )
vis = VisApp ( )
app = GCApp ( output_dir , batch_size , step_size , guidance_scale , vis )
threading . Thread ( target = app . run , daemon = True ) . start ( )
gui . Application . instance . run ( )
if __name__ == " __main__ " :
parser = argparse . ArgumentParser ( description = " Generate shapes with the ShapeGenerator. " )
parser . add_argument ( " --output_dir " , type = str , required = True , help = " The directory to save generated shapes. " )
parser . add_argument ( " --batch_size " , type = int , default = 4 , help = " The number of batches for shap-e. the higher the batch size the longer it will take to process but will output a more refined mesh. " )
parser . add_argument ( " --step_size " , type = int , default = 64 , help = " The number of steps/iterations for shap-e. the higher the step size the longer it will take to process but will output a more refined mesh. " )
parser . add_argument ( " --guidance_scale " , type = int , default = 30 , help = " The guidance scale in context to the text prompt. The higher this value, the model will generate something closer to the text description (CLIP). " )
args = parser . parse_args ( )
main ( args . output_dir , args . batch_size , args . step_size , args . guidance_scale )