Publish.

.gitignore

@@ -0,0 +1,2 @@
+__pycache__
+*.egg-info

openshape/__init__.py

@@ -0,0 +1,47 @@
+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

openshape/demo/caption.py

@@ -0,0 +1,163 @@
+from torch import nn
+import numpy as np
+import torch
+from typing import Tuple, List, Union, Optional
+from transformers import GPT2Tokenizer, GPT2LMHeadModel
+from huggingface_hub import hf_hub_download
+
+
+N = type(None)
+V = np.array
+ARRAY = np.ndarray
+ARRAYS = Union[Tuple[ARRAY, ...], List[ARRAY]]
+VS = Union[Tuple[V, ...], List[V]]
+VN = Union[V, N]
+VNS = Union[VS, N]
+T = torch.Tensor
+TS = Union[Tuple[T, ...], List[T]]
+TN = Optional[T]
+TNS = Union[Tuple[TN, ...], List[TN]]
+TSN = Optional[TS]
+TA = Union[T, ARRAY]
+
+
+D = torch.device
+
+
+class MLP(nn.Module):
+
+    def forward(self, x: T) -> T:
+        return self.model(x)
+
+    def __init__(self, sizes: Tuple[int, ...], bias=True, act=nn.Tanh):
+        super(MLP, self).__init__()
+        layers = []
+        for i in range(len(sizes) -1):
+            layers.append(nn.Linear(sizes[i], sizes[i + 1], bias=bias))
+            if i < len(sizes) - 2:
+                layers.append(act())
+        self.model = nn.Sequential(*layers)
+
+
+class ClipCaptionModel(nn.Module):
+
+    #@functools.lru_cache #FIXME
+    def get_dummy_token(self, batch_size: int, device: D) -> T:
+        return torch.zeros(batch_size, self.prefix_length, dtype=torch.int64, device=device)
+
+    def forward(self, tokens: T, prefix: T, mask: Optional[T] = None, labels: Optional[T] = None):
+        embedding_text = self.gpt.transformer.wte(tokens)
+        prefix_projections = self.clip_project(prefix).view(-1, self.prefix_length, self.gpt_embedding_size)
+        #print(embedding_text.size()) #torch.Size([5, 67, 768])
+        #print(prefix_projections.size()) #torch.Size([5, 1, 768])
+        embedding_cat = torch.cat((prefix_projections, embedding_text), dim=1)
+        if labels is not None:
+            dummy_token = self.get_dummy_token(tokens.shape[0], tokens.device)
+            labels = torch.cat((dummy_token, tokens), dim=1)
+        out = self.gpt(inputs_embeds=embedding_cat, labels=labels, attention_mask=mask)
+        return out
+
+    def __init__(self, prefix_length: int, prefix_size: int = 512):
+        super(ClipCaptionModel, self).__init__()
+        self.prefix_length = prefix_length
+        self.gpt = GPT2LMHeadModel.from_pretrained('gpt2')
+        self.gpt_embedding_size = self.gpt.transformer.wte.weight.shape[1]
+        if prefix_length > 10:  # not enough memory
+            self.clip_project = nn.Linear(prefix_size, self.gpt_embedding_size * prefix_length)
+        else:
+            self.clip_project = MLP((prefix_size, (self.gpt_embedding_size * prefix_length) // 2, self.gpt_embedding_size * prefix_length))
+
+
+class ClipCaptionPrefix(ClipCaptionModel):
+
+    def parameters(self, recurse: bool = True):
+        return self.clip_project.parameters()
+
+    def train(self, mode: bool = True):
+        super(ClipCaptionPrefix, self).train(mode)
+        self.gpt.eval()
+        return self
+
+
+def generate2(
+    model,
+    tokenizer,
+    tokens=None,
+    prompt=None,
+    embed=None,
+    entry_count=1,
+    entry_length=67,  # maximum number of words
+    top_p=0.8,
+    temperature=1.,
+    stop_token: str = '.',
+):
+    model.eval()
+    generated_num = 0
+    generated_list = []
+    stop_token_index = tokenizer.encode(stop_token)[0]
+    filter_value = -float("Inf")
+    device = next(model.parameters()).device
+    score_col = []
+    with torch.no_grad():
+
+        for entry_idx in range(entry_count):
+            if embed is not None:
+                generated = embed
+            else:
+                if tokens is None:
+                    tokens = torch.tensor(tokenizer.encode(prompt))
+                    tokens = tokens.unsqueeze(0).to(device)
+
+                generated = model.gpt.transformer.wte(tokens)
+
+            for i in range(entry_length):
+
+                outputs = model.gpt(inputs_embeds=generated)
+                logits = outputs.logits
+                logits = logits[:, -1, :] / (temperature if temperature > 0 else 1.0)
+                sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+                cumulative_probs = torch.cumsum(torch.softmax(sorted_logits, dim=-1), dim=-1)
+                sorted_indices_to_remove = cumulative_probs > top_p
+                sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[
+                                                    ..., :-1
+                                                    ].clone()
+                sorted_indices_to_remove[..., 0] = 0
+
+                indices_to_remove = sorted_indices[sorted_indices_to_remove]
+                logits[:, indices_to_remove] = filter_value
+                next_token = torch.argmax(torch.softmax(logits, dim=-1), -1).reshape(1, 1)
+                score = torch.softmax(logits, dim=-1).reshape(-1)[next_token.item()].item()
+                score_col.append(score)
+                next_token_embed = model.gpt.transformer.wte(next_token)
+                if tokens is None:
+                    tokens = next_token
+                else:
+                    tokens = torch.cat((tokens, next_token), dim=1)
+                generated = torch.cat((generated, next_token_embed), dim=1)
+                if stop_token_index == next_token.item():
+                    break
+
+            output_list = list(tokens.squeeze(0).cpu().numpy())
+            output_text = tokenizer.decode(output_list)
+            generated_list.append(output_text)
+    return generated_list[0]
+
+
+@torch.no_grad()
+def pc_caption(pc_encoder: torch.nn.Module, pc, cond_scale):
+    ref_dev = next(pc_encoder.parameters()).device
+    prefix = pc_encoder(torch.tensor(pc.T[None], device=ref_dev))
+    prefix = prefix.float() * cond_scale
+    prefix_embed = model.clip_project(prefix).reshape(1, prefix_length, -1)
+    text = generate2(model, tokenizer, embed=prefix_embed)
+    return text
+
+
+tokenizer = GPT2Tokenizer.from_pretrained("gpt2")
+prefix_length = 10
+model = ClipCaptionModel(prefix_length)
+# print(model.gpt_embedding_size)
+model.load_state_dict(torch.load(hf_hub_download('OpenShape/clipcap-cc', 'conceptual_weights.pt', token=True), map_location='cpu'))
+model.eval()
+if torch.cuda.is_available():
+    model = model.cuda()

openshape/demo/classification.py

@@ -0,0 +1,13 @@
+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/demo/lvis.py

@@ -0,0 +1,1162 @@
+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/demo/lvis_cats.pt

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

openshape/demo/misc_utils.py

@@ -0,0 +1,153 @@
+import numpy
+import trimesh
+import trimesh.sample
+import trimesh.visual
+import trimesh.proximity
+import objaverse
+import streamlit as st
+import plotly.graph_objects as go
+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(geometry)
+        total_area = sum(geometry.area for geometry in meshes)
+        if total_area < 1e-6:
+            raise ValueError("Bad geometry: total area too small (< 1e-6)")
+        pcs = []
+        for geometry in meshes:
+            pcs.append(model_to_pc(geometry, max(1, round(geometry.area / total_area * 10000))))
+        if not len(pcs):
+            raise ValueError("Unsupported mesh object: no triangles found")
+        return numpy.concatenate(pcs)
+    else:
+        assert isinstance(scene_or_mesh, trimesh.Trimesh)
+        return model_to_pc(scene_or_mesh, 10000)
+
+
+def input_3d_shape():
+    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)")
+    f32 = numpy.float32
+
+    def load_data(prog):
+        # 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 = 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 = 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)
+
+    return load_data
+
+
+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)))
+    fig.update_scenes(aspectmode="data")
+    col1, col2 = st.columns(2)
+    with col1:
+        st.plotly_chart(fig, use_container_width=True)
+        # st.caption("Point Cloud Preview")
+    return col2

openshape/demo/retrieval.py

@@ -0,0 +1,40 @@
+import json
+import torch
+import torch.nn.functional as F
+from huggingface_hub import hf_hub_download
+
+
+meta = json.load(
+    open(hf_hub_download("OpenShape/openshape-objaverse-embeddings", "objaverse_meta.json", token=True, repo_type='dataset'))
+)
+# {
+# "u": "94db219c315742909fee67deeeacae15",
+# "name": "knife", "like": 0, "view": 35,
+# "tags": ["game-ready", "damascus", "damascus_steel", "kabar-knife", "knife", "blender", "blender3d", "gameready"],
+# "cats": ["weapons-military"],
+# "img": "https://media.sketchfab.com/models/94db219c315742909fee67deeeacae15/thumbnails/c0bbbd475d264ff2a92972f5115564ee/0cd28a130ebd4d9c9ef73190f24d9a42.jpeg",
+# "desc": "", "faces": 1724, "size": 11955, "lic": "by",
+# "glb": "glbs/000-000/94db219c315742909fee67deeeacae15.glb"
+# }
+meta = {x['u']: x for x in meta['entries']}
+deser = torch.load(
+    hf_hub_download("OpenShape/openshape-objaverse-embeddings", "objaverse.pt", token=True, repo_type='dataset'), map_location='cpu'
+)
+us = deser['us']
+feats = deser['feats']
+
+
+def retrieve(embedding, top):
+    sims = []
+    embedding = F.normalize(embedding.detach().cpu(), dim=-1).squeeze()
+    for chunk in torch.split(feats, 10240):
+        sims.append(embedding @ F.normalize(chunk.float(), dim=-1).T)
+    sims = torch.cat(sims)
+    sims, idx = torch.topk(sims, top * 2)
+    results = []
+    for i, sim in zip(idx, sims):
+        if us[i] in meta:
+            results.append(dict(meta[us[i]], sim=sim))
+        if len(results) >= top:
+            break
+    return results

openshape/demo/sd_pc2img.py

@@ -0,0 +1,38 @@
+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]

openshape/pointnet_util.py

@@ -0,0 +1,323 @@
+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

@@ -0,0 +1,118 @@
+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))

setup.py

@@ -0,0 +1,21 @@
+import setuptools
+
+
+def packages():
+    return setuptools.find_packages()
+
+
+setuptools.setup(
+    name="openshape",
+    version="0.1",
+    author="flandre.info",
+    author_email="flandre@scarletx.cn",
+    description="Support library for OpenShape Demos.",
+    packages=packages(),
+    classifiers=[
+        "Programming Language :: Python :: 3 :: Only",
+        "License :: OSI Approved :: Apache Software License",
+        "Operating System :: OS Independent",
+    ],
+    python_requires='~=3.7',
+)