Hugging Face's logo

Spaces:
ICML2022
/
PointCloudC
Build error

App Files Community
1
PointCloudC
File size: 4,660 Bytes 
d7b89b7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
d015e2a
 
d7b89b7
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
3090213
 
 
 
 
 
 
 
 
 
 
 
d7b89b7
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
 
import gradio as gr
import mathutils
import math
import numpy as np
import matplotlib.pyplot as plt
import matplotlib
import matplotlib.cm as cmx
import os.path as osp
import h5py
import random
import torch
import torch.nn as nn

from GDANet_cls import GDANET
from DGCNN import DGCNN

with open('shape_names.txt') as f:
    CLASS_NAME = f.read().splitlines()

model_gda = GDANET()
model_gda = nn.DataParallel(model_gda)
model_gda.load_state_dict(torch.load('./GDANet_WOLFMix.t7', map_location=torch.device('cpu')))
model_gda.eval()

model_dgcnn = DGCNN()
model_dgcnn = nn.DataParallel(model_dgcnn)
model_dgcnn.load_state_dict(torch.load('./dgcnn.t7', map_location=torch.device('cpu')))
model_dgcnn.eval()

def pyplot_draw_point_cloud(points, corruption):
    rot1 = mathutils.Euler([-math.pi / 2, 0, 0]).to_matrix().to_3x3()
    rot2 = mathutils.Euler([0, 0, math.pi]).to_matrix().to_3x3()
    points = np.dot(points, rot1)
    points = np.dot(points, rot2)
    x, y, z = points[:, 0], points[:, 1], points[:, 2]
    colorsMap = 'winter'
    cs = y
    cm = plt.get_cmap(colorsMap)
    cNorm = matplotlib.colors.Normalize(vmin=-1, vmax=1)
    scalarMap = cmx.ScalarMappable(norm=cNorm, cmap=cm)
    fig = plt.figure(figsize=(5, 5))
    ax = fig.add_subplot(111, projection='3d')
    ax.scatter(x, y, z, c=scalarMap.to_rgba(cs))
    scalarMap.set_array(cs)
    ax.set_xlim(-1, 1)
    ax.set_ylim(-1, 1)
    ax.set_zlim(-1, 1)
    plt.axis('off')
    plt.title(corruption, fontsize=30)
    plt.tight_layout()
    plt.savefig('visualization.png', bbox_inches='tight', dpi=200)
    plt.close()



def load_dataset(corruption_idx, severity):
    corruptions = [
        'clean',
        'scale',
        'jitter',
        'rotate',
        'dropout_global',
        'dropout_local',
        'add_global',
        'add_local',
    ]
    corruption_type = corruptions[corruption_idx]
    if corruption_type == 'clean':
        f = h5py.File(osp.join('modelnet_c', corruption_type + '.h5'))
    else:
        f = h5py.File(osp.join('modelnet_c', corruption_type + '_{}'.format(severity-1) + '.h5'))
    data = f['data'][:].astype('float32')
    label = f['label'][:].astype('int64')
    f.close()
    return data, label

def recognize_pcd(model, pcd):
    pcd = torch.tensor(pcd).unsqueeze(0)
    pcd = pcd.permute(0, 2, 1)
    output = model(pcd)
    prediction = output.softmax(-1).flatten()
    _, top5_idx = torch.topk(prediction, 5)
    return {CLASS_NAME[i]: float(prediction[i]) for i in top5_idx.tolist()}

def run(seed, corruption_idx, severity):
    data, label = load_dataset(corruption_idx, severity)
    random.seed(seed)
    sample_indx = random.randint(0, data.shape[0])
    pcd, cls = data[sample_indx], label[sample_indx]
    pyplot_draw_point_cloud(pcd, CLASS_NAME[cls[0]])
    output = 'visualization.png'
    return output, recognize_pcd(model_dgcnn, pcd), recognize_pcd(model_gda, pcd)

if __name__ == '__main__':
    iface = gr.Interface(
        fn=run,
        inputs=[
            gr.components.Number(label='Sample Seed', precision=0),
            gr.components.Radio(
                ['Clean', 'Scale', 'Jitter', 'Rotate', 'Drop Global', 'Drop Local', 'Add Global', 'Add Local'],
                value='Clean', type="index", label='Corruption Type'),
            gr.components.Slider(1, 5, step=1, label='Corruption severity'),
        ],
        outputs=[
            gr.components.Image(type="file", label="Visualization"),
            gr.components.Label(num_top_classes=5, label="Baseline (DGCNN) Prediction"),
            gr.components.Label(num_top_classes=5, label="Ours (GDANet+WolfMix) Prediction")
        ],
        live=False,
        allow_flagging='never',
        title="ModelNet-C",
        description="""
        
                    Welcome to the demo of ModelNet-C! In this demo, you may: 
                    
                    - **Visualize** various types of corrupted point clouds in ModelNet-C,
                    
                    - **Compare** our proposed techniques to the baseline in terms of prediction robustness.
                    
                    For more details, checkout more our paper [Benchmarking and Analyzing Point Cloud Classification under Corruptions, ICML 2022](https://arxiv.org/abs/2202.03377)!
                    
                    """,
        examples=[
            [0, 'Jitter', 5],
            [999, 'Drop Local', 5],
        ],
        # css=".output-image, .image-preview {height: 500px !important}",
        article="<p style='text-align: center'><a href='https://github.com/jiawei-ren/ModelNet-C' target='_blank'>ModelNet-C @ GitHub</a></p> "
    )
    iface.launch()