Use support library.

app.py

@@ -1,14 +1,12 @@
 import streamlit as st
-from huggingface_hub import HfFolder
+from huggingface_hub import HfFolder, snapshot_download
 HfFolder().save_token(st.secrets['etoken'])
+snapshot_download("OpenShape/openshape-demo-support", local_dir='.')
 
 
 import numpy
-import trimesh
-import objaverse
 import openshape
-import misc_utils
-import plotly.graph_objects as go
+from openshape.demo import misc_utils, classification
 
 
 @st.cache_resource
@@ -17,70 +15,21 @@ def load_openshape(name):
 
 
 f32 = numpy.float32
-model_b32 = openshape.load_pc_encoder('openshape-pointbert-vitb32-rgb')
-model_l14 = openshape.load_pc_encoder('openshape-pointbert-vitl14-rgb')
+# clip_model, clip_prep = load_openclip()
 model_g14 = openshape.load_pc_encoder('openshape-pointbert-vitg14-rgb')
 
 
 st.title("OpenShape Demo")
-objaid = st.text_input("Enter an Objaverse ID")
-model = st.file_uploader("Or upload a model (.glb/.obj/.ply)")
-npy = st.file_uploader("Or upload a point cloud numpy array (.npy of Nx3 XYZ or Nx6 XYZRGB)")
-swap_yz_axes = st.checkbox("Swap Y/Z axes of input (Y is up for OpenShape)")
+load_data = misc_utils.input_3d_shape()
 prog = st.progress(0.0, "Idle")
 
 
-def load_data():
-    # load the model
-    prog.progress(0.05, "Preparing Point Cloud")
-    if npy is not None:
-        pc: numpy.ndarray = numpy.load(npy)
-    elif model is not None:
-        pc = misc_utils.trimesh_to_pc(trimesh.load(model, model.name.split(".")[-1]))
-    elif objaid:
-        prog.progress(0.1, "Downloading Objaverse Object")
-        objamodel = objaverse.load_objects([objaid])[objaid]
-        prog.progress(0.2, "Preparing Point Cloud")
-        pc = misc_utils.trimesh_to_pc(trimesh.load(objamodel))
-    else:
-        raise ValueError("You have to supply 3D input!")
-    prog.progress(0.25, "Preprocessing Point Cloud")
-    assert pc.ndim == 2, "invalid pc shape: ndim = %d != 2" % pc.ndim
-    assert pc.shape[1] in [3, 6], "invalid pc shape: should have 3/6 channels, got %d" % pc.shape[1]
-    if swap_yz_axes:
-        pc[:, [1, 2]] = pc[:, [2, 1]]
-    pc[:, :3] = pc[:, :3] - numpy.mean(pc[:, :3], axis=0)
-    pc[:, :3] = pc[:, :3] / numpy.linalg.norm(pc[:, :3], axis=-1).max()
-    if pc.shape[1] == 3:
-        pc = numpy.concatenate([pc, numpy.ones_like(pc)], axis=-1)
-    prog.progress(0.3, "Preprocessed Point Cloud")
-    return pc.astype(f32)
-
-
-def render_pc(pc):
-    rand = numpy.random.permutation(len(pc))[:2048]
-    pc = pc[rand]
-    rgb = (pc[:, 3:] * 255).astype(numpy.uint8)
-    g = go.Scatter3d(
-        x=pc[:, 0], y=pc[:, 1], z=pc[:, 2],
-        mode='markers',
-        marker=dict(size=2, color=[f'rgb({rgb[i, 0]}, {rgb[i, 1]}, {rgb[i, 2]})' for i in range(len(pc))]),
-    )
-    fig = go.Figure(data=[g])
-    fig.update_layout(scene_camera=dict(up=dict(x=0, y=1, z=0)))
-    col1, col2 = st.columns(2)
-    with col1:
-        st.plotly_chart(fig, use_container_width=True)
-        # st.caption("Point Cloud Preview")
-    return col2
-
-
 try:
     if st.button("Run Classification on LVIS Categories"):
-        pc = load_data()
-        col2 = render_pc(pc)
+        pc = load_data(prog)
+        col2 = misc_utils.render_pc(pc)
         prog.progress(0.5, "Running Classification")
-        pred = openshape.pred_lvis_sims(model_g14, pc)
+        pred = classification.pred_lvis_sims(model_g14, pc)
         with col2:
             for i, (cat, sim) in zip(range(5), pred.items()):
                 st.text(cat)

misc_utils.py

@@ -1,90 +0,0 @@
-import numpy
-import trimesh
-import trimesh.sample
-import trimesh.visual
-import trimesh.proximity
-import streamlit as st
-import matplotlib.pyplot as plotlib
-
-
-def get_bytes(x: str):
-    import io, requests
-    return io.BytesIO(requests.get(x).content)
-
-
-def get_image(x: str):
-    try:
-        return plotlib.imread(get_bytes(x), 'auto')
-    except Exception:
-        raise ValueError("Invalid image", x)
-
-
-def model_to_pc(mesh: trimesh.Trimesh, n_sample_points=10000):
-    f32 = numpy.float32
-    rad = numpy.sqrt(mesh.area / (3 * n_sample_points))
-    for _ in range(24):
-        pcd, face_idx = trimesh.sample.sample_surface_even(mesh, n_sample_points, rad)
-        rad *= 0.85
-        if len(pcd) == n_sample_points:
-            break
-    else:
-        raise ValueError("Bad geometry, cannot finish sampling.", mesh.area)
-    if isinstance(mesh.visual, trimesh.visual.ColorVisuals):
-        rgba = mesh.visual.face_colors[face_idx]
-    elif isinstance(mesh.visual, trimesh.visual.TextureVisuals):
-        bc = trimesh.proximity.points_to_barycentric(mesh.triangles[face_idx], pcd)
-        if mesh.visual.uv is None or len(mesh.visual.uv) < mesh.faces[face_idx].max():
-            uv = numpy.zeros([len(bc), 2])
-            st.warning("Invalid UV, filling with zeroes")
-        else:
-            uv = numpy.einsum('ntc,nt->nc', mesh.visual.uv[mesh.faces[face_idx]], bc)
-        material = mesh.visual.material
-        if hasattr(material, 'materials'):
-            if len(material.materials) == 0:
-                rgba = numpy.ones_like(pcd) * 0.8
-                texture = None
-                st.warning("Empty MultiMaterial found, falling back to light grey")
-            else:
-                material = material.materials[0]
-        if hasattr(material, 'image'):
-            texture = material.image
-            if texture is None:
-                rgba = numpy.zeros([len(uv), len(material.main_color)]) + material.main_color
-        elif hasattr(material, 'baseColorTexture'):
-            texture = material.baseColorTexture
-            if texture is None:
-                rgba = numpy.zeros([len(uv), len(material.main_color)]) + material.main_color
-        else:
-            texture = None
-            rgba = numpy.ones_like(pcd) * 0.8
-            st.warning("Unknown material, falling back to light grey")
-        if texture is not None:
-            rgba = trimesh.visual.uv_to_interpolated_color(uv, texture)
-    if rgba.max() > 1:
-        if rgba.max() > 255:
-            rgba = rgba.astype(f32) / rgba.max()
-        else:
-            rgba = rgba.astype(f32) / 255.0
-    return numpy.concatenate([numpy.array(pcd, f32), numpy.array(rgba, f32)[:, :3]], axis=-1)
-
-
-def trimesh_to_pc(scene_or_mesh):
-    if isinstance(scene_or_mesh, trimesh.Scene):
-        meshes = []
-        for node_name in scene_or_mesh.graph.nodes_geometry:
-            # which geometry does this node refer to
-            transform, geometry_name = scene_or_mesh.graph[node_name]
-
-            # get the actual potential mesh instance
-            geometry = scene_or_mesh.geometry[geometry_name].copy()
-            if not hasattr(geometry, 'triangles'):
-                continue
-            geometry: trimesh.Trimesh
-            geometry = geometry.apply_transform(transform)
-            meshes.append(model_to_pc(geometry, 10000 // len(scene_or_mesh.geometry)))
-        if not len(meshes):
-            raise ValueError("Unsupported mesh object: no triangles found")
-        return numpy.concatenate(meshes)
-    else:
-        assert isinstance(scene_or_mesh, trimesh.Trimesh)
-        return model_to_pc(scene_or_mesh, 10000)

openshape/__init__.py

@@ -1,52 +0,0 @@
-import torch
-import torch.nn as nn
-from huggingface_hub import hf_hub_download
-from .ppat_rgb import Projected, PointPatchTransformer
-
-
-def module(state_dict: dict, name):
-    return {'.'.join(k.split('.')[1:]): v for k, v in state_dict.items() if k.startswith(name + '.')}
-
-
-def G14(s):
-    model = Projected(
-        PointPatchTransformer(512, 12, 8, 512*3, 256, 384, 0.2, 64, 6),
-        nn.Linear(512, 1280)
-    )
-    model.load_state_dict(module(s['state_dict'], 'module'))
-    return model
-
-
-def L14(s):
-    model = Projected(
-        PointPatchTransformer(512, 12, 8, 1024, 128, 64, 0.4, 256, 6),
-        nn.Linear(512, 768)
-    )
-    model.load_state_dict(module(s, 'pc_encoder'))
-    return model
-
-
-def B32(s):
-    model = PointPatchTransformer(512, 12, 8, 1024, 128, 64, 0.4, 256, 6)
-    model.load_state_dict(module(s, 'pc_encoder'))
-    return model
-
-
-model_list = {
-    "openshape-pointbert-vitb32-rgb": B32,
-    "openshape-pointbert-vitl14-rgb": L14,
-    "openshape-pointbert-vitg14-rgb": G14,
-}
-
-
-def load_pc_encoder(name):
-    s = torch.load(hf_hub_download("OpenShape/" + name, "model.pt", token=True), map_location='cpu')
-    model = model_list[name](s).eval()
-    if torch.cuda.is_available():
-        model.cuda()
-    return model
-
-
-# only import the functions in demo!
-# from .sd_pc2img import pc_to_image
-from .classification import pred_lvis_sims

openshape/classification.py

@@ -1,13 +0,0 @@
-import torch
-import torch.nn.functional as F
-from collections import OrderedDict
-from . import lvis
-
-
-@torch.no_grad()
-def pred_lvis_sims(pc_encoder: torch.nn.Module, pc):
-    ref_dev = next(pc_encoder.parameters()).device
-    enc = pc_encoder(torch.tensor(pc[:, [0, 2, 1, 3, 4, 5]].T[None], device=ref_dev)).cpu()
-    sim = torch.matmul(F.normalize(lvis.feats, dim=-1), F.normalize(enc, dim=-1).squeeze())
-    argsort = torch.argsort(sim, descending=True)
-    return OrderedDict((lvis.categories[i], sim[i]) for i in argsort if i < len(lvis.categories))

openshape/lvis.py

@@ -1,1162 +0,0 @@
-import os
-import torch
-
-
-feats = torch.load(os.path.join(os.path.dirname(__file__), 'lvis_cats.pt'))
-categories = [
-    'Band_Aid',
-    'Bible',
-    'CD_player',
-    'Christmas_tree',
-    'Dixie_cup',
-    'Ferris_wheel',
-    'Lego',
-    'Rollerblade',
-    'Sharpie',
-    'Tabasco_sauce',
-    'aerosol_can',
-    'air_conditioner',
-    'airplane',
-    'alarm_clock',
-    'alcohol',
-    'alligator',
-    'almond',
-    'ambulance',
-    'amplifier',
-    'anklet',
-    'antenna',
-    'apple',
-    'apricot',
-    'apron',
-    'aquarium',
-    'arctic_(type_of_shoe)',
-    'armband',
-    'armchair',
-    'armoire',
-    'armor',
-    'army_tank',
-    'artichoke',
-    'ashtray',
-    'asparagus',
-    'atomizer',
-    'automatic_washer',
-    'avocado',
-    'award',
-    'awning',
-    'ax',
-    'baboon',
-    'baby_buggy',
-    'backpack',
-    'bagel',
-    'baguet',
-    'bait',
-    'ball',
-    'ballet_skirt',
-    'balloon',
-    'bamboo',
-    'banana',
-    'bandage',
-    'bandanna',
-    'banjo',
-    'banner',
-    'barbell',
-    'barge',
-    'barrel',
-    'barrow',
-    'baseball',
-    'baseball_bat',
-    'baseball_cap',
-    'baseball_glove',
-    'basket',
-    'basketball',
-    'basketball_backboard',
-    'bass_horn',
-    'bat_(animal)',
-    'bath_mat',
-    'bath_towel',
-    'bathrobe',
-    'bathtub',
-    'battery',
-    'beachball',
-    'bead',
-    'beanbag',
-    'beanie',
-    'bear',
-    'bed',
-    'bedpan',
-    'bedspread',
-    'beef_(food)',
-    'beeper',
-    'beer_bottle',
-    'beer_can',
-    'beetle',
-    'bell',
-    'bell_pepper',
-    'belt',
-    'belt_buckle',
-    'bench',
-    'beret',
-    'bicycle',
-    'billboard',
-    'binder',
-    'binoculars',
-    'bird',
-    'birdbath',
-    'birdcage',
-    'birdfeeder',
-    'birdhouse',
-    'birthday_cake',
-    'birthday_card',
-    'blackberry',
-    'blackboard',
-    'blanket',
-    'blazer',
-    'blender',
-    'blimp',
-    'blouse',
-    'blueberry',
-    'boat',
-    'bob',
-    'bobbin',
-    'boiled_egg',
-    'bolo_tie',
-    'bolt',
-    'bonnet',
-    'book',
-    'bookcase',
-    'booklet',
-    'bookmark',
-    'boom_microphone',
-    'boot',
-    'bottle',
-    'bottle_cap',
-    'bottle_opener',
-    'bouquet',
-    'bow-tie',
-    'bow_(decorative_ribbons)',
-    'bow_(weapon)',
-    'bowl',
-    'bowler_hat',
-    'bowling_ball',
-    'box',
-    'boxing_glove',
-    'bracelet',
-    'brass_plaque',
-    'brassiere',
-    'bread',
-    'bread-bin',
-    'breechcloth',
-    'bridal_gown',
-    'briefcase',
-    'broach',
-    'broccoli',
-    'broom',
-    'brownie',
-    'brussels_sprouts',
-    'bubble_gum',
-    'bucket',
-    'bulldog',
-    'bulldozer',
-    'bullet_train',
-    'bulletin_board',
-    'bulletproof_vest',
-    'bullhorn',
-    'bun',
-    'bunk_bed',
-    'buoy',
-    'burrito',
-    'bus_(vehicle)',
-    'business_card',
-    'butter',
-    'butterfly',
-    'button',
-    'cab_(taxi)',
-    'cabana',
-    'cabin_car',
-    'cabinet',
-    'cake',
-    'calculator',
-    'calendar',
-    'calf',
-    'camcorder',
-    'camel',
-    'camera',
-    'camera_lens',
-    'camper_(vehicle)',
-    'can',
-    'can_opener',
-    'candle',
-    'candle_holder',
-    'candy_bar',
-    'candy_cane',
-    'canister',
-    'canoe',
-    'cantaloup',
-    'canteen',
-    'cap_(headwear)',
-    'cape',
-    'cappuccino',
-    'car_(automobile)',
-    'car_battery',
-    'card',
-    'cardigan',
-    'cargo_ship',
-    'carnation',
-    'carrot',
-    'cart',
-    'carton',
-    'cash_register',
-    'casserole',
-    'cassette',
-    'cast',
-    'cat',
-    'cauliflower',
-    'cayenne_(spice)',
-    'celery',
-    'cellular_telephone',
-    'chair',
-    'chaise_longue',
-    'chalice',
-    'chandelier',
-    'checkbook',
-    'checkerboard',
-    'cherry',
-    'chessboard',
-    'chicken_(animal)',
-    'chili_(vegetable)',
-    'chime',
-    'chinaware',
-    'chocolate_bar',
-    'chocolate_cake',
-    'chocolate_milk',
-    'chocolate_mousse',
-    'choker',
-    'chopping_board',
-    'chopstick',
-    'cider',
-    'cigar_box',
-    'cigarette',
-    'cigarette_case',
-    'cincture',
-    'cistern',
-    'clarinet',
-    'clasp',
-    'cleansing_agent',
-    'cleat_(for_securing_rope)',
-    'clementine',
-    'clip',
-    'clipboard',
-    'clippers_(for_plants)',
-    'cloak',
-    'clock',
-    'clock_tower',
-    'clothes_hamper',
-    'clothespin',
-    'clutch_bag',
-    'coaster',
-    'coat',
-    'coat_hanger',
-    'coatrack',
-    'cock',
-    'cockroach',
-    'cocoa_(beverage)',
-    'coconut',
-    'coffee_maker',
-    'coffee_table',
-    'coffeepot',
-    'coil',
-    'coin',
-    'colander',
-    'coloring_material',
-    'combination_lock',
-    'comic_book',
-    'compass',
-    'computer_keyboard',
-    'condiment',
-    'cone',
-    'control',
-    'convertible_(automobile)',
-    'cooker',
-    'cookie',
-    'cooking_utensil',
-    'cooler_(for_food)',
-    'cork_(bottle_plug)',
-    'corkboard',
-    'corkscrew',
-    'cornbread',
-    'cornet',
-    'cornice',
-    'cornmeal',
-    'corset',
-    'costume',
-    'cougar',
-    'cover',
-    'coverall',
-    'cow',
-    'cowbell',
-    'cowboy_hat',
-    'crab_(animal)',
-    'crabmeat',
-    'cracker',
-    'crape',
-    'crate',
-    'crawfish',
-    'crayon',
-    'cream_pitcher',
-    'crescent_roll',
-    'crib',
-    'crisp_(potato_chip)',
-    'crossbar',
-    'crouton',
-    'crow',
-    'crowbar',
-    'crown',
-    'crucifix',
-    'cruise_ship',
-    'crutch',
-    'cub_(animal)',
-    'cube',
-    'cucumber',
-    'cufflink',
-    'cup',
-    'cupboard',
-    'cupcake',
-    'curtain',
-    'cushion',
-    'cylinder',
-    'cymbal',
-    'dagger',
-    'dalmatian',
-    'dartboard',
-    'date_(fruit)',
-    'deadbolt',
-    'deck_chair',
-    'deer',
-    'desk',
-    'detergent',
-    'diaper',
-    'diary',
-    'die',
-    'dinghy',
-    'dining_table',
-    'dirt_bike',
-    'dish',
-    'dish_antenna',
-    'dishrag',
-    'dishtowel',
-    'dishwasher',
-    'dishwasher_detergent',
-    'dispenser',
-    'dog',
-    'dog_collar',
-    'doll',
-    'dollar',
-    'dollhouse',
-    'dolphin',
-    'domestic_ass',
-    'doorknob',
-    'doormat',
-    'doughnut',
-    'dove',
-    'dragonfly',
-    'drawer',
-    'dress',
-    'dress_hat',
-    'dress_suit',
-    'dresser',
-    'drill',
-    'drone',
-    'drum_(musical_instrument)',
-    'drumstick',
-    'duck',
-    'duckling',
-    'duct_tape',
-    'duffel_bag',
-    'dumbbell',
-    'dumpster',
-    'dustpan',
-    'eagle',
-    'earphone',
-    'earplug',
-    'earring',
-    'easel',
-    'eclair',
-    'edible_corn',
-    'eel',
-    'egg',
-    'egg_roll',
-    'egg_yolk',
-    'eggbeater',
-    'eggplant',
-    'elephant',
-    'elevator_car',
-    'elk',
-    'envelope',
-    'eraser',
-    'escargot',
-    'eyepatch',
-    'falcon',
-    'fan',
-    'faucet',
-    'fedora',
-    'ferret',
-    'ferry',
-    'fig_(fruit)',
-    'fighter_jet',
-    'figurine',
-    'file_(tool)',
-    'file_cabinet',
-    'fire_alarm',
-    'fire_engine',
-    'fire_extinguisher',
-    'fire_hose',
-    'fireplace',
-    'fireplug',
-    'first-aid_kit',
-    'fish',
-    'fish_(food)',
-    'fishbowl',
-    'fishing_rod',
-    'flag',
-    'flagpole',
-    'flamingo',
-    'flannel',
-    'flap',
-    'flash',
-    'flashlight',
-    'fleece',
-    'flip-flop_(sandal)',
-    'flipper_(footwear)',
-    'flower_arrangement',
-    'flowerpot',
-    'flute_glass',
-    'foal',
-    'folding_chair',
-    'food_processor',
-    'football_(American)',
-    'football_helmet',
-    'footstool',
-    'fork',
-    'forklift',
-    'freight_car',
-    'freshener',
-    'frisbee',
-    'frog',
-    'fruit_juice',
-    'frying_pan',
-    'fume_hood',
-    'funnel',
-    'futon',
-    'gameboard',
-    'garbage',
-    'garbage_truck',
-    'garden_hose',
-    'gargle',
-    'gargoyle',
-    'garlic',
-    'gasmask',
-    'gazelle',
-    'gelatin',
-    'gemstone',
-    'generator',
-    'giant_panda',
-    'gift_wrap',
-    'ginger',
-    'giraffe',
-    'glass_(drink_container)',
-    'globe',
-    'glove',
-    'goat',
-    'goggles',
-    'goldfish',
-    'golf_club',
-    'golfcart',
-    'gondola_(boat)',
-    'goose',
-    'gorilla',
-    'gourd',
-    'grape',
-    'grater',
-    'gravestone',
-    'gravy_boat',
-    'green_bean',
-    'green_onion',
-    'grill',
-    'grits',
-    'grizzly',
-    'grocery_bag',
-    'guitar',
-    'gull',
-    'gun',
-    'hair_dryer',
-    'hairbrush',
-    'hairnet',
-    'halter_top',
-    'ham',
-    'hamburger',
-    'hammer',
-    'hammock',
-    'hamper',
-    'hamster',
-    'hand_glass',
-    'hand_towel',
-    'handbag',
-    'handcart',
-    'handcuff',
-    'handkerchief',
-    'handle',
-    'handsaw',
-    'hardback_book',
-    'harmonium',
-    'hat',
-    'hatbox',
-    'headband',
-    'headboard',
-    'headlight',
-    'headscarf',
-    'headset',
-    'headstall_(for_horses)',
-    'heart',
-    'heater',
-    'helicopter',
-    'helmet',
-    'heron',
-    'highchair',
-    'hinge',
-    'hippopotamus',
-    'hockey_stick',
-    'hog',
-    'honey',
-    'hook',
-    'hookah',
-    'horned_cow',
-    'hornet',
-    'horse',
-    'horse_buggy',
-    'horse_carriage',
-    'hose',
-    'hot-air_balloon',
-    'hot_sauce',
-    'hotplate',
-    'hourglass',
-    'houseboat',
-    'hummingbird',
-    'iPod',
-    'ice_maker',
-    'ice_pack',
-    'ice_skate',
-    'icecream',
-    'identity_card',
-    'igniter',
-    'inhaler',
-    'inkpad',
-    'iron_(for_clothing)',
-    'ironing_board',
-    'jacket',
-    'jam',
-    'jar',
-    'jean',
-    'jeep',
-    'jersey',
-    'jet_plane',
-    'jewel',
-    'jewelry',
-    'joystick',
-    'jumpsuit',
-    'kayak',
-    'keg',
-    'kennel',
-    'kettle',
-    'key',
-    'keycard',
-    'kilt',
-    'kimono',
-    'kitchen_sink',
-    'kitchen_table',
-    'kite',
-    'kitten',
-    'kiwi_fruit',
-    'knee_pad',
-    'knife',
-    'knitting_needle',
-    'knob',
-    'knocker_(on_a_door)',
-    'koala',
-    'lab_coat',
-    'ladder',
-    'ladle',
-    'ladybug',
-    'lamb-chop',
-    'lamb_(animal)',
-    'lamp',
-    'lamppost',
-    'lampshade',
-    'lantern',
-    'laptop_computer',
-    'lasagna',
-    'latch',
-    'lawn_mower',
-    'leather',
-    'legging_(clothing)',
-    'legume',
-    'lemon',
-    'lemonade',
-    'lettuce',
-    'license_plate',
-    'life_buoy',
-    'life_jacket',
-    'lightbulb',
-    'lightning_rod',
-    'lime',
-    'limousine',
-    'lion',
-    'lip_balm',
-    'liquor',
-    'lizard',
-    'locker',
-    'log',
-    'lollipop',
-    'loveseat',
-    'machine_gun',
-    'magazine',
-    'magnet',
-    'mail_slot',
-    'mailbox_(at_home)',
-    'mallard',
-    'mallet',
-    'mammoth',
-    'manatee',
-    'mandarin_orange',
-    'manger',
-    'manhole',
-    'map',
-    'marker',
-    'martini',
-    'mascot',
-    'mashed_potato',
-    'mask',
-    'mast',
-    'mat_(gym_equipment)',
-    'matchbox',
-    'mattress',
-    'measuring_cup',
-    'measuring_stick',
-    'meatball',
-    'medicine',
-    'melon',
-    'microphone',
-    'microscope',
-    'microwave_oven',
-    'milestone',
-    'milk',
-    'milk_can',
-    'milkshake',
-    'minivan',
-    'mint_candy',
-    'mirror',
-    'mitten',
-    'mixer_(kitchen_tool)',
-    'money',
-    'monitor_(computer_equipment) computer_monitor',
-    'monkey',
-    'mop',
-    'motor',
-    'motor_scooter',
-    'motor_vehicle',
-    'motorcycle',
-    'mound_(baseball)',
-    'mouse_(computer_equipment)',
-    'mousepad',
-    'muffin',
-    'mug',
-    'mushroom',
-    'music_stool',
-    'musical_instrument',
-    'nailfile',
-    'napkin',
-    'neckerchief',
-    'necklace',
-    'necktie',
-    'needle',
-    'nest',
-    'newspaper',
-    'newsstand',
-    'nightshirt',
-    'notebook',
-    'notepad',
-    'nut',
-    'nutcracker',
-    'oar',
-    'octopus_(animal)',
-    'octopus_(food)',
-    'oil_lamp',
-    'olive_oil',
-    'omelet',
-    'onion',
-    'orange_(fruit)',
-    'orange_juice',
-    'ostrich',
-    'ottoman',
-    'oven',
-    'overalls_(clothing)',
-    'owl',
-    'pacifier',
-    'packet',
-    'paddle',
-    'padlock',
-    'paintbrush',
-    'painting',
-    'pajamas',
-    'palette',
-    'pan_(for_cooking)',
-    'pan_(metal_container)',
-    'pancake',
-    'papaya',
-    'paper_plate',
-    'paper_towel',
-    'paperback_book',
-    'paperweight',
-    'parachute',
-    'parakeet',
-    'parasail_(sports)',
-    'parasol',
-    'parchment',
-    'parka',
-    'parking_meter',
-    'parrot',
-    'passenger_car_(part_of_a_train)',
-    'passenger_ship',
-    'passport',
-    'pastry',
-    'patty_(food)',
-    'pea_(food)',
-    'peach',
-    'peanut_butter',
-    'pear',
-    'peeler_(tool_for_fruit_and_vegetables)',
-    'pegboard',
-    'pelican',
-    'pen',
-    'pencil',
-    'pencil_box',
-    'pencil_sharpener',
-    'pendulum',
-    'penguin',
-    'pennant',
-    'penny_(coin)',
-    'pepper',
-    'pepper_mill',
-    'perfume',
-    'persimmon',
-    'person',
-    'pet',
-    'pew_(church_bench)',
-    'phonebook',
-    'phonograph_record',
-    'piano',
-    'pickle',
-    'pickup_truck',
-    'pie',
-    'pigeon',
-    'piggy_bank',
-    'pillow',
-    'pineapple',
-    'pinecone',
-    'ping-pong_ball',
-    'pinwheel',
-    'pipe',
-    'pipe_bowl',
-    'pirate_flag',
-    'pistol',
-    'pita_(bread)',
-    'pitcher_(vessel_for_liquid)',
-    'pitchfork',
-    'pizza',
-    'place_mat',
-    'plastic_bag',
-    'plate',
-    'platter',
-    'playpen',
-    'pliers',
-    'plow_(farm_equipment)',
-    'plume',
-    'pocket_watch',
-    'pocketknife',
-    'poker_(fire_stirring_tool)',
-    'poker_chip',
-    'polar_bear',
-    'pole',
-    'police_cruiser',
-    'polo_shirt',
-    'poncho',
-    'pony',
-    'pool_table',
-    'pop_(soda)',
-    'popsicle',
-    'postbox_(public)',
-    'postcard',
-    'poster',
-    'pot',
-    'potato',
-    'potholder',
-    'pottery',
-    'pouch',
-    'power_shovel',
-    'prawn',
-    'pretzel',
-    'printer',
-    'projectile_(weapon)',
-    'projector',
-    'propeller',
-    'prune',
-    'pudding',
-    'puffer_(fish)',
-    'puffin',
-    'pug-dog',
-    'pumpkin',
-    'puncher',
-    'puppet',
-    'puppy',
-    'quesadilla',
-    'quiche',
-    'quilt',
-    'rabbit',
-    'race_car',
-    'racket',
-    'radar',
-    'radiator',
-    'radio_receiver',
-    'radish',
-    'raft',
-    'rag_doll',
-    'railcar_(part_of_a_train)',
-    'raincoat',
-    'ram_(animal)',
-    'raspberry',
-    'rat',
-    'reamer_(juicer)',
-    'rearview_mirror',
-    'receipt',
-    'recliner',
-    'record_player',
-    'reflector',
-    'refrigerator',
-    'remote_control',
-    'rhinoceros',
-    'rib_(food)',
-    'rifle',
-    'ring',
-    'river_boat',
-    'road_map',
-    'robe',
-    'rocking_chair',
-    'rodent',
-    'roller_skate',
-    'rolling_pin',
-    'root_beer',
-    'router_(computer_equipment)',
-    'rubber_band',
-    'runner_(carpet)',
-    'saddle_(on_an_animal)',
-    'saddle_blanket',
-    'saddlebag',
-    'safety_pin',
-    'sail',
-    'salad',
-    'salad_plate',
-    'salami',
-    'salmon_(fish)',
-    'salmon_(food)',
-    'salsa',
-    'saltshaker',
-    'sandal_(type_of_shoe)',
-    'sandwich',
-    'satchel',
-    'saucepan',
-    'saucer',
-    'sausage',
-    'sawhorse',
-    'saxophone',
-    'scale_(measuring_instrument)',
-    'scarecrow',
-    'scarf',
-    'school_bus',
-    'scissors',
-    'scoreboard',
-    'scraper',
-    'screwdriver',
-    'scrubbing_brush',
-    'sculpture',
-    'seabird',
-    'seahorse',
-    'seaplane',
-    'seashell',
-    'sewing_machine',
-    'shaker',
-    'shampoo',
-    'shark',
-    'sharpener',
-    'shaver_(electric)',
-    'shaving_cream',
-    'shawl',
-    'shears',
-    'sheep',
-    'shepherd_dog',
-    'sherbert',
-    'shield',
-    'shirt',
-    'shoe',
-    'shopping_bag',
-    'shopping_cart',
-    'short_pants',
-    'shot_glass',
-    'shoulder_bag',
-    'shovel',
-    'shower_cap',
-    'shower_curtain',
-    'shower_head',
-    'shredder_(for_paper)',
-    'signboard',
-    'silo',
-    'sink',
-    'skateboard',
-    'skewer',
-    'ski',
-    'ski_boot',
-    'ski_parka',
-    'ski_pole',
-    'skirt',
-    'skullcap',
-    'sled',
-    'sleeping_bag',
-    'slide',
-    'slipper_(footwear)',
-    'smoothie',
-    'snake',
-    'snowboard',
-    'snowman',
-    'snowmobile',
-    'soap',
-    'soccer_ball',
-    'sock',
-    'sofa',
-    'sofa_bed',
-    'softball',
-    'solar_array',
-    'sombrero',
-    'soup',
-    'soup_bowl',
-    'soupspoon',
-    'soya_milk',
-    'space_shuttle',
-    'sparkler_(fireworks)',
-    'spatula',
-    'speaker_(stero_equipment)',
-    'spear',
-    'spectacles',
-    'spice_rack',
-    'spider',
-    'sponge',
-    'spoon',
-    'sportswear',
-    'spotlight',
-    'squid_(food)',
-    'squirrel',
-    'stagecoach',
-    'stapler_(stapling_machine)',
-    'starfish',
-    'statue_(sculpture)',
-    'steak_(food)',
-    'steak_knife',
-    'steering_wheel',
-    'step_stool',
-    'stepladder',
-    'stereo_(sound_system)',
-    'stew',
-    'stirrer',
-    'stirrup',
-    'stool',
-    'stop_sign',
-    'stove',
-    'strainer',
-    'strap',
-    'straw_(for_drinking)',
-    'strawberry',
-    'street_sign',
-    'streetlight',
-    'string_cheese',
-    'stylus',
-    'subwoofer',
-    'sugar_bowl',
-    'sugarcane_(plant)',
-    'suit_(clothing)',
-    'suitcase',
-    'sunflower',
-    'sunglasses',
-    'sunhat',
-    'surfboard',
-    'sushi',
-    'suspenders',
-    'sweat_pants',
-    'sweatband',
-    'sweater',
-    'sweatshirt',
-    'sweet_potato',
-    'swimsuit',
-    'sword',
-    'syringe',
-    'table',
-    'table-tennis_table',
-    'table_lamp',
-    'tablecloth',
-    'tachometer',
-    'taco',
-    'tag',
-    'taillight',
-    'tambourine',
-    'tank_(storage_vessel)',
-    'tank_top_(clothing)',
-    'tape_(sticky_cloth_or_paper)',
-    'tape_measure',
-    'tapestry',
-    'tarp',
-    'tartan',
-    'tassel',
-    'teacup',
-    'teakettle',
-    'teapot',
-    'teddy_bear',
-    'telephone',
-    'telephone_booth',
-    'telephone_pole',
-    'telephoto_lens',
-    'television_camera',
-    'television_set',
-    'tennis_ball',
-    'tennis_racket',
-    'tequila',
-    'thermometer',
-    'thermos_bottle',
-    'thermostat',
-    'thimble',
-    'thread',
-    'thumbtack',
-    'tiara',
-    'tiger',
-    'tights_(clothing)',
-    'timer',
-    'tinfoil',
-    'tinsel',
-    'tissue_paper',
-    'toast_(food)',
-    'toaster',
-    'toaster_oven',
-    'tobacco_pipe',
-    'toilet',
-    'toilet_tissue',
-    'tomato',
-    'tongs',
-    'toolbox',
-    'toothbrush',
-    'toothpaste',
-    'toothpick',
-    'tortilla',
-    'tote_bag',
-    'tow_truck',
-    'towel',
-    'towel_rack',
-    'toy',
-    'tractor_(farm_equipment)',
-    'traffic_light',
-    'trailer_truck',
-    'train_(railroad_vehicle)',
-    'trampoline',
-    'trash_can',
-    'tray',
-    'trench_coat',
-    'triangle_(musical_instrument)',
-    'tricycle',
-    'tripod',
-    'trophy_cup',
-    'trousers',
-    'truck',
-    'truffle_(chocolate)',
-    'trunk',
-    'turban',
-    'turkey_(food)',
-    'turnip',
-    'turtle',
-    'turtleneck_(clothing)',
-    'tux',
-    'typewriter',
-    'umbrella',
-    'underdrawers',
-    'underwear',
-    'unicycle',
-    'urinal',
-    'urn',
-    'vacuum_cleaner',
-    'vase',
-    'veil',
-    'vending_machine',
-    'vent',
-    'vest',
-    'videotape',
-    'vinegar',
-    'violin',
-    'visor',
-    'vodka',
-    'volleyball',
-    'vulture',
-    'waffle',
-    'waffle_iron',
-    'wagon',
-    'walking_cane',
-    'walking_stick',
-    'wall_clock',
-    'wall_socket',
-    'wallet',
-    'walrus',
-    'wardrobe',
-    'washbasin',
-    'watch',
-    'water_bottle',
-    'water_cooler',
-    'water_faucet',
-    'water_gun',
-    'water_heater',
-    'water_jug',
-    'water_scooter',
-    'water_ski',
-    'water_tower',
-    'watering_can',
-    'watermelon',
-    'weathervane',
-    'webcam',
-    'wedding_cake',
-    'wedding_ring',
-    'wet_suit',
-    'wheel',
-    'wheelchair',
-    'whipped_cream',
-    'wig',
-    'wind_chime',
-    'windmill',
-    'window_box_(for_plants)',
-    'windsock',
-    'wine_bottle',
-    'wine_bucket',
-    'wineglass',
-    'wok',
-    'wolf',
-    'wooden_leg',
-    'wooden_spoon',
-    'wreath',
-    'wrench',
-    'wristband',
-    'wristlet',
-    'yacht',
-    'yogurt',
-    'zebra',
-    'zucchini'
-]

openshape/lvis_cats.pt

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:71baf2d3f89884a082f1db75d0e94ac9a3b8036553877a3fdd98861cd01c4aec
-size 5919467

openshape/pointnet_util.py

@@ -1,323 +0,0 @@
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from time import time
-import numpy as np
-import dgl.geometry
-
-def timeit(tag, t):
-    print("{}: {}s".format(tag, time() - t))
-    return time()
-
-def pc_normalize(pc):
-    l = pc.shape[0]
-    centroid = np.mean(pc, axis=0)
-    pc = pc - centroid
-    m = np.max(np.sqrt(np.sum(pc**2, axis=1)))
-    pc = pc / m
-    return pc
-
-def square_distance(src, dst):
-    """
-    Calculate Euclid distance between each two points.
-
-    src^T * dst = xn * xm + yn * ym + zn * zm；
-    sum(src^2, dim=-1) = xn*xn + yn*yn + zn*zn;
-    sum(dst^2, dim=-1) = xm*xm + ym*ym + zm*zm;
-    dist = (xn - xm)^2 + (yn - ym)^2 + (zn - zm)^2
-         = sum(src**2,dim=-1)+sum(dst**2,dim=-1)-2*src^T*dst
-
-    Input:
-        src: source points, [B, N, C]
-        dst: target points, [B, M, C]
-    Output:
-        dist: per-point square distance, [B, N, M]
-    """
-    B, N, _ = src.shape
-    _, M, _ = dst.shape
-    dist = -2 * torch.matmul(src, dst.permute(0, 2, 1))
-    dist += torch.sum(src ** 2, -1).view(B, N, 1)
-    dist += torch.sum(dst ** 2, -1).view(B, 1, M)
-    return dist
-
-
-def index_points(points, idx):
-    """
-
-    Input:
-        points: input points data, [B, N, C]
-        idx: sample index data, [B, S]
-    Return:
-        new_points:, indexed points data, [B, S, C]
-    """
-    device = points.device
-    B = points.shape[0]
-    view_shape = list(idx.shape)
-    view_shape[1:] = [1] * (len(view_shape) - 1)
-    repeat_shape = list(idx.shape)
-    repeat_shape[0] = 1
-    batch_indices = torch.arange(B, dtype=torch.long).to(device).view(view_shape).repeat(repeat_shape)
-    new_points = points[batch_indices, idx, :]
-    return new_points
-
-
-def farthest_point_sample(xyz, npoint):
-    """
-    Input:
-        xyz: pointcloud data, [B, N, 3]
-        npoint: number of samples
-    Return:
-        centroids: sampled pointcloud index, [B, npoint]
-    """
-    return dgl.geometry.farthest_point_sampler(xyz, npoint)
-    device = xyz.device
-    B, N, C = xyz.shape
-    centroids = torch.zeros(B, npoint, dtype=torch.long).to(device)
-    distance = torch.ones(B, N).to(device) * 1e10
-    farthest = torch.randint(0, N, (B,), dtype=torch.long).to(device)
-    batch_indices = torch.arange(B, dtype=torch.long).to(device)
-    for i in range(npoint):
-        centroids[:, i] = farthest
-        centroid = xyz[batch_indices, farthest, :].view(B, 1, 3)
-        dist = torch.sum((xyz - centroid) ** 2, -1)
-        mask = dist < distance
-        distance[mask] = dist[mask]
-        farthest = torch.max(distance, -1)[1]
-    return centroids
-
-
-def query_ball_point(radius, nsample, xyz, new_xyz):
-    """
-    Input:
-        radius: local region radius
-        nsample: max sample number in local region
-        xyz: all points, [B, N, 3]
-        new_xyz: query points, [B, S, 3]
-    Return:
-        group_idx: grouped points index, [B, S, nsample]
-    """
-    device = xyz.device
-    B, N, C = xyz.shape
-    _, S, _ = new_xyz.shape
-    group_idx = torch.arange(N, dtype=torch.long).to(device).view(1, 1, N).repeat([B, S, 1])
-    sqrdists = square_distance(new_xyz, xyz)
-    group_idx[sqrdists > radius ** 2] = N
-    group_idx = group_idx.sort(dim=-1)[0][:, :, :nsample]
-    group_first = group_idx[..., :1].repeat([1, 1, nsample])
-    mask = group_idx == N
-    group_idx[mask] = group_first[mask]
-    return group_idx
-
-
-def sample_and_group(npoint, radius, nsample, xyz, points, returnfps=False):
-    """
-    Input:
-        npoint:
-        radius:
-        nsample:
-        xyz: input points position data, [B, N, 3]
-        points: input points data, [B, N, D]
-    Return:
-        new_xyz: sampled points position data, [B, npoint, nsample, 3]
-        new_points: sampled points data, [B, npoint, nsample, 3+D]
-    """
-    B, N, C = xyz.shape
-    S = npoint
-    fps_idx = farthest_point_sample(xyz, npoint) # [B, npoint, C]
-    # torch.cuda.empty_cache()
-    new_xyz = index_points(xyz, fps_idx)
-    # torch.cuda.empty_cache()
-    idx = query_ball_point(radius, nsample, xyz, new_xyz)
-    # torch.cuda.empty_cache()
-    grouped_xyz = index_points(xyz, idx) # [B, npoint, nsample, C]
-    # torch.cuda.empty_cache()
-    grouped_xyz_norm = grouped_xyz - new_xyz.view(B, S, 1, C)
-    # torch.cuda.empty_cache()
-
-    if points is not None:
-        grouped_points = index_points(points, idx)
-        new_points = torch.cat([grouped_xyz_norm, grouped_points], dim=-1) # [B, npoint, nsample, C+D]
-    else:
-        new_points = grouped_xyz_norm
-    if returnfps:
-        return new_xyz, new_points, grouped_xyz, fps_idx
-    else:
-        return new_xyz, new_points
-
-
-def sample_and_group_all(xyz, points):
-    """
-    Input:
-        xyz: input points position data, [B, N, 3]
-        points: input points data, [B, N, D]
-    Return:
-        new_xyz: sampled points position data, [B, 1, 3]
-        new_points: sampled points data, [B, 1, N, 3+D]
-    """
-    device = xyz.device
-    B, N, C = xyz.shape
-    new_xyz = torch.zeros(B, 1, C).to(device)
-    grouped_xyz = xyz.view(B, 1, N, C)
-    if points is not None:
-        new_points = torch.cat([grouped_xyz, points.view(B, 1, N, -1)], dim=-1)
-    else:
-        new_points = grouped_xyz
-    return new_xyz, new_points
-
-
-class PointNetSetAbstraction(nn.Module):
-    def __init__(self, npoint, radius, nsample, in_channel, mlp, group_all):
-        super(PointNetSetAbstraction, self).__init__()
-        self.npoint = npoint
-        self.radius = radius
-        self.nsample = nsample
-        self.mlp_convs = nn.ModuleList()
-        self.mlp_bns = nn.ModuleList()
-        last_channel = in_channel
-        for out_channel in mlp:
-            self.mlp_convs.append(nn.Conv2d(last_channel, out_channel, 1))
-            self.mlp_bns.append(nn.BatchNorm2d(out_channel))
-            last_channel = out_channel
-        self.group_all = group_all
-
-    def forward(self, xyz, points):
-        """
-        Input:
-            xyz: input points position data, [B, C, N]
-            points: input points data, [B, D, N]
-        Return:
-            new_xyz: sampled points position data, [B, C, S]
-            new_points_concat: sample points feature data, [B, D', S]
-        """
-        xyz = xyz.permute(0, 2, 1)
-        if points is not None:
-            points = points.permute(0, 2, 1)
-
-        if self.group_all:
-            new_xyz, new_points = sample_and_group_all(xyz, points)
-        else:
-            new_xyz, new_points = sample_and_group(self.npoint, self.radius, self.nsample, xyz, points)
-        # new_xyz: sampled points position data, [B, npoint, C]
-        # new_points: sampled points data, [B, npoint, nsample, C+D]
-        new_points = new_points.permute(0, 3, 2, 1) # [B, C+D, nsample,npoint]
-        for i, conv in enumerate(self.mlp_convs):
-            bn = self.mlp_bns[i]
-            new_points =  F.relu(bn(conv(new_points)))
-
-        new_points = torch.max(new_points, 2)[0]
-        new_xyz = new_xyz.permute(0, 2, 1)
-        return new_xyz, new_points
-
-
-class PointNetSetAbstractionMsg(nn.Module):
-    def __init__(self, npoint, radius_list, nsample_list, in_channel, mlp_list):
-        super(PointNetSetAbstractionMsg, self).__init__()
-        self.npoint = npoint
-        self.radius_list = radius_list
-        self.nsample_list = nsample_list
-        self.conv_blocks = nn.ModuleList()
-        self.bn_blocks = nn.ModuleList()
-        for i in range(len(mlp_list)):
-            convs = nn.ModuleList()
-            bns = nn.ModuleList()
-            last_channel = in_channel + 3
-            for out_channel in mlp_list[i]:
-                convs.append(nn.Conv2d(last_channel, out_channel, 1))
-                bns.append(nn.BatchNorm2d(out_channel))
-                last_channel = out_channel
-            self.conv_blocks.append(convs)
-            self.bn_blocks.append(bns)
-
-    def forward(self, xyz, points):
-        """
-        Input:
-            xyz: input points position data, [B, C, N]
-            points: input points data, [B, D, N]
-        Return:
-            new_xyz: sampled points position data, [B, C, S]
-            new_points_concat: sample points feature data, [B, D', S]
-        """
-        xyz = xyz.permute(0, 2, 1)
-        if points is not None:
-            points = points.permute(0, 2, 1)
-
-        B, N, C = xyz.shape
-        S = self.npoint
-        new_xyz = index_points(xyz, farthest_point_sample(xyz, S))
-        new_points_list = []
-        for i, radius in enumerate(self.radius_list):
-            K = self.nsample_list[i]
-            group_idx = query_ball_point(radius, K, xyz, new_xyz)
-            grouped_xyz = index_points(xyz, group_idx)
-            grouped_xyz -= new_xyz.view(B, S, 1, C)
-            if points is not None:
-                grouped_points = index_points(points, group_idx)
-                grouped_points = torch.cat([grouped_points, grouped_xyz], dim=-1)
-            else:
-                grouped_points = grouped_xyz
-
-            grouped_points = grouped_points.permute(0, 3, 2, 1)  # [B, D, K, S]
-            for j in range(len(self.conv_blocks[i])):
-                conv = self.conv_blocks[i][j]
-                bn = self.bn_blocks[i][j]
-                grouped_points =  F.relu(bn(conv(grouped_points)))
-            new_points = torch.max(grouped_points, 2)[0]  # [B, D', S]
-            new_points_list.append(new_points)
-
-        new_xyz = new_xyz.permute(0, 2, 1)
-        new_points_concat = torch.cat(new_points_list, dim=1)
-        return new_xyz, new_points_concat
-
-
-class PointNetFeaturePropagation(nn.Module):
-    def __init__(self, in_channel, mlp):
-        super(PointNetFeaturePropagation, self).__init__()
-        self.mlp_convs = nn.ModuleList()
-        self.mlp_bns = nn.ModuleList()
-        last_channel = in_channel
-        for out_channel in mlp:
-            self.mlp_convs.append(nn.Conv1d(last_channel, out_channel, 1))
-            self.mlp_bns.append(nn.BatchNorm1d(out_channel))
-            last_channel = out_channel
-
-    def forward(self, xyz1, xyz2, points1, points2):
-        """
-        Input:
-            xyz1: input points position data, [B, C, N]
-            xyz2: sampled input points position data, [B, C, S]
-            points1: input points data, [B, D, N]
-            points2: input points data, [B, D, S]
-        Return:
-            new_points: upsampled points data, [B, D', N]
-        """
-        xyz1 = xyz1.permute(0, 2, 1)
-        xyz2 = xyz2.permute(0, 2, 1)
-
-        points2 = points2.permute(0, 2, 1)
-        B, N, C = xyz1.shape
-        _, S, _ = xyz2.shape
-
-        if S == 1:
-            interpolated_points = points2.repeat(1, N, 1)
-        else:
-            dists = square_distance(xyz1, xyz2)
-            dists, idx = dists.sort(dim=-1)
-            dists, idx = dists[:, :, :3], idx[:, :, :3]  # [B, N, 3]
-
-            dist_recip = 1.0 / (dists + 1e-8)
-            norm = torch.sum(dist_recip, dim=2, keepdim=True)
-            weight = dist_recip / norm
-            interpolated_points = torch.sum(index_points(points2, idx) * weight.view(B, N, 3, 1), dim=2)
-
-        if points1 is not None:
-            points1 = points1.permute(0, 2, 1)
-            new_points = torch.cat([points1, interpolated_points], dim=-1)
-        else:
-            new_points = interpolated_points
-
-        new_points = new_points.permute(0, 2, 1)
-        for i, conv in enumerate(self.mlp_convs):
-            bn = self.mlp_bns[i]
-            new_points = F.relu(bn(conv(new_points)))
-        return new_points
-

openshape/ppat_rgb.py

@@ -1,118 +0,0 @@
-import torch
-import torch.nn as nn
-import torch_redstone as rst
-from einops import rearrange
-from .pointnet_util import PointNetSetAbstraction
-
-
-class PreNorm(nn.Module):
-    def __init__(self, dim, fn):
-        super().__init__()
-        self.norm = nn.LayerNorm(dim)
-        self.fn = fn
-    def forward(self, x, *extra_args, **kwargs):
-        return self.fn(self.norm(x), *extra_args, **kwargs)
-
-class FeedForward(nn.Module):
-    def __init__(self, dim, hidden_dim, dropout = 0.):
-        super().__init__()
-        self.net = nn.Sequential(
-            nn.Linear(dim, hidden_dim),
-            nn.GELU(),
-            nn.Dropout(dropout),
-            nn.Linear(hidden_dim, dim),
-            nn.Dropout(dropout)
-        )
-    def forward(self, x):
-        return self.net(x)
-
-class Attention(nn.Module):
-    def __init__(self, dim, heads = 8, dim_head = 64, dropout = 0., rel_pe = False):
-        super().__init__()
-        inner_dim = dim_head *  heads
-        project_out = not (heads == 1 and dim_head == dim)
-
-        self.heads = heads
-        self.scale = dim_head ** -0.5
-
-        self.attend = nn.Softmax(dim = -1)
-        self.dropout = nn.Dropout(dropout)
-
-        self.to_qkv = nn.Linear(dim, inner_dim * 3, bias = False)
-
-        self.to_out = nn.Sequential(
-            nn.Linear(inner_dim, dim),
-            nn.Dropout(dropout)
-        ) if project_out else nn.Identity()
-
-        self.rel_pe = rel_pe
-        if rel_pe:
-            self.pe = nn.Sequential(nn.Conv2d(3, 64, 1), nn.ReLU(), nn.Conv2d(64, 1, 1))
-
-    def forward(self, x, centroid_delta):
-        qkv = self.to_qkv(x).chunk(3, dim = -1)
-        q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> b h n d', h = self.heads), qkv)
-
-        pe = self.pe(centroid_delta) if self.rel_pe else 0
-        dots = (torch.matmul(q, k.transpose(-1, -2)) + pe) * self.scale
-
-        attn = self.attend(dots)
-        attn = self.dropout(attn)
-
-        out = torch.matmul(attn, v)
-        out = rearrange(out, 'b h n d -> b n (h d)')
-        return self.to_out(out)
-
-
-class Transformer(nn.Module):
-    def __init__(self, dim, depth, heads, dim_head, mlp_dim, dropout = 0., rel_pe = False):
-        super().__init__()
-        self.layers = nn.ModuleList([])
-        for _ in range(depth):
-            self.layers.append(nn.ModuleList([
-                PreNorm(dim, Attention(dim, heads = heads, dim_head = dim_head, dropout = dropout, rel_pe = rel_pe)),
-                PreNorm(dim, FeedForward(dim, mlp_dim, dropout = dropout))
-            ]))
-    def forward(self, x, centroid_delta):
-        for attn, ff in self.layers:
-            x = attn(x, centroid_delta) + x
-            x = ff(x) + x
-        return x
-
-
-class PointPatchTransformer(nn.Module):
-    def __init__(self, dim, depth, heads, mlp_dim, sa_dim, patches, prad, nsamp, in_dim=3, dim_head=64, rel_pe=False, patch_dropout=0) -> None:
-        super().__init__()
-        self.patches = patches
-        self.patch_dropout = patch_dropout
-        self.sa = PointNetSetAbstraction(npoint=patches, radius=prad, nsample=nsamp, in_channel=in_dim + 3, mlp=[64, 64, sa_dim], group_all=False)
-        self.lift = nn.Sequential(nn.Conv1d(sa_dim + 3, dim, 1), rst.Lambda(lambda x: torch.permute(x, [0, 2, 1])), nn.LayerNorm([dim]))
-        self.cls_token = nn.Parameter(torch.randn(dim))
-        self.transformer = Transformer(dim, depth, heads, dim_head, mlp_dim, 0.0, rel_pe)
-
-    def forward(self, features):
-        self.sa.npoint = self.patches
-        if self.training:
-            self.sa.npoint -= self.patch_dropout
-        # print("input", features.shape)
-        centroids, feature = self.sa(features[:, :3], features)
-        # print("f", feature.shape, 'c', centroids.shape)
-        x = self.lift(torch.cat([centroids, feature], dim=1))
-
-        x = rst.supercat([self.cls_token, x], dim=-2)
-        centroids = rst.supercat([centroids.new_zeros(1), centroids], dim=-1)
-
-        centroid_delta = centroids.unsqueeze(-1) - centroids.unsqueeze(-2)
-        x = self.transformer(x, centroid_delta)
-
-        return x[:, 0]
-
-
-class Projected(nn.Module):
-    def __init__(self, ppat, proj) -> None:
-        super().__init__()
-        self.ppat = ppat
-        self.proj = proj
-
-    def forward(self, features: torch.Tensor):
-        return self.proj(self.ppat(features))

openshape/sd_pc2img.py

@@ -1,39 +0,0 @@
-import numpy
-import torch
-import torch_redstone as rst
-import transformers
-from diffusers import StableUnCLIPImg2ImgPipeline
-
-
-class Wrapper(transformers.modeling_utils.PreTrainedModel):
-    def __init__(self) -> None:
-        super().__init__(transformers.configuration_utils.PretrainedConfig())
-        self.param = torch.nn.Parameter(torch.tensor(0.))
-
-    def forward(self, x):
-        return rst.ObjectProxy(image_embeds=x)
-
-
-pipe = StableUnCLIPImg2ImgPipeline.from_pretrained(
-    "diffusers/stable-diffusion-2-1-unclip-i2i-l",
-    image_encoder = Wrapper()
-)
-if torch.cuda.is_available():
-    pipe = pipe.to('cuda:' + str(torch.cuda.current_device()))
-    pipe.enable_model_cpu_offload(torch.cuda.current_device())
-
-
-@torch.no_grad()
-def pc_to_image(pc_encoder: torch.nn.Module, pc, prompt, noise_level, width, height, cfg_scale, num_steps, callback):
-    ref_dev = next(pc_encoder.parameters()).device
-    enc = pc_encoder(torch.tensor(pc.T[None], device=ref_dev))
-    return pipe(
-        prompt="best quality, super high resolution, " + prompt,
-        negative_prompt="cropped, worst quality, low quality, normal quality, jpeg artifacts, signature, watermark, username, blurry",
-        image=torch.nn.functional.normalize(enc, dim=-1) * (768 ** 0.5) / 2,
-        width=width, height=height,
-        guidance_scale=cfg_scale,
-        noise_level=noise_level,
-        callback=callback,
-        num_inference_steps=num_steps
-    ).images[0]