Untitled
unknown
plain_text
3 months ago
20 kB
4
Indexable
import os
import sys
import json
from pathlib import Path
import numpy as np
from tqdm import tqdm
import pandas as pd
import glob
import random
from sklearn.metrics import roc_auc_score, roc_curve, auc
# plotting
import matplotlib.pyplot as plt
import plotly.graph_objects as go
from plotly.subplots import make_subplots
# DeepLearning
import torch
import torch.nn as nn
import torch.nn.functional as F
from torch.autograd import Variable
import torch.optim as optim
print("Imports OK")
print(np.__version__)
print(torch.__version__)
# DATASET SPLITS
DATA_FOLDER = 'Shapenetcore_benchmark/'
class_name_id_map = {'Airplane': 0, 'Bag': 1, 'Cap': 2, 'Car': 3, 'Chair': 4,
'Earphone': 5, 'Guitar': 6, 'Knife': 7, 'Lamp': 8, 'Laptop': 9,
'Motorbike': 10, 'Mug': 11, 'Pistol': 12, 'Rocket': 13,
'Skateboard': 14, 'Table': 15}
#class_name_id_map = { 'Car': 3,'Motorbike': 10,}
class_id_name_map = {v:k for k,v in class_name_id_map.items()}
PCD_SCENE=dict(xaxis=dict(visible=False), yaxis=dict(visible=False), zaxis=dict(visible=False), aspectmode='data')
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
train_split_data = json.load(open('Shapenetcore_benchmark/train_split.json', 'r'))
train_class_count = np.array([x[0] for x in train_split_data])
# plot classwise count in train set
train_dist_plots = [go.Bar(x=list(class_name_id_map.keys()), y= np.bincount(train_class_count))]
layout = dict(title="Train dataset", title_x=0.5)
#fig = go.Figure(data=train_dist_plots, layout=layout)
#fig.show()
points_list = glob.glob("Shapenetcore_benchmark/04379243/points/*.npy")
print(len(points_list))
# LOAD A POINT CLOUD
# Random point cloud
idx = random.randint(0,len(points_list))
# load random point cloud
points = np.load(points_list[idx])
print(f"points shape = {points.shape}, min xyz = {np.min(points, axis=0)}, max xyz = {np.max(points, axis=0)}")
# load segmentation lables
seg_file_path = points_list[idx].replace('points', 'points_label').replace('.npy', '.seg')
seg_labels = np.loadtxt(seg_file_path).astype(np.int8)
print(f"seg_labels shape = {seg_labels.shape}, unique labels = {np.unique(seg_labels)}")
import plotly.graph_objects as go
# There are max of 16 parts in an object in Shapenet core dataset
# Creating random colors according to part label
NUM_PARTS = 16
PART_COLORS = np.random.choice(range(255), size=(NUM_PARTS, 3))
# Create a color array for eacoint cloud datah point based on its segment label
colors = PART_COLORS[seg_labels - 1] / 255.0 # Normalize to [0, 1] range
# Create a 3D scatter plot using plotly
fig = go.Figure(data=[go.Scatter3d(
x=points[:, 0],
y=points[:, 1],
z=points[:, 2],
mode='markers',
marker=dict(
size=2,
color=colors, # Set color to the colors array
opacity=0.8
)
)])
# Set layout for the plot
fig.update_layout(
template="plotly_dark",
title="Raw Point Cloud",
scene=dict(
xaxis_title='X',
yaxis_title='Y',
zaxis_title='Z'
),
title_x=0.5
)
# Show the plot
#fig.show()
class ShapeNetDataset(torch.utils.data.Dataset):
def __init__(self, root_dir, split_type, num_samples=2500):
self.root_dir = root_dir
self.split_type = split_type
self.num_samples = num_samples
with open(os.path.join(root_dir, f'{self.split_type}_split.json'), 'r') as f:
self.split_data = json.load(f)
def __getitem__(self, index):
# read point cloud data
class_id, class_name, point_cloud_path, seg_label_path = self.split_data[index]
# point cloud data
point_cloud_path = os.path.join(self.root_dir, point_cloud_path)
pc_data = np.load(point_cloud_path)
# segmentation labels
# -1 is to change part values from [1-16] to [0-15]
# which helps when running segmentation
pc_seg_labels = np.loadtxt(os.path.join(self.root_dir, seg_label_path)).astype(np.int8) - 1
#pc_seg_labels = pc_seg_labels.reshape(pc_seg_labels.size,1)
# Sample fixed number of points
num_points = pc_data.shape[0]
if num_points < self.num_samples:
# Duplicate random points if the number of points is less than max_num_points
additional_indices = np.random.choice(num_points, self.num_samples - num_points, replace=True)
pc_data = np.concatenate((pc_data, pc_data[additional_indices]), axis=0)
pc_seg_labels = np.concatenate((pc_seg_labels, pc_seg_labels[additional_indices]), axis=0)
else:
# Randomly sample max_num_points from the available points
random_indices = np.random.choice(num_points, self.num_samples)
pc_data = pc_data[random_indices]
pc_seg_labels = pc_seg_labels[random_indices]
# return variable
data_dict= {}
data_dict['class_id'] = class_id
data_dict['class_name'] = class_name
data_dict['points'] = pc_data
data_dict['seg_labels'] = pc_seg_labels
return data_dict
def __len__(self):
return len(self.split_data)
train_set = ShapeNetDataset(root_dir = DATA_FOLDER, split_type='train')
val_set = ShapeNetDataset(root_dir = DATA_FOLDER, split_type='val')
test_set = ShapeNetDataset(root_dir = DATA_FOLDER, split_type='test')
def showDataLayout():
plt.bar(["Train","Valid","Test"],
[len(train_set), len(val_set), len(test_set)], align='center',color=[ 'black','red', 'blue'])
plt.legend()
plt.ylabel('Number of images')
plt.title('Data distribution')
plt.savefig('data2018.png')
plt.show()
#print(len(train_set))
#print(len(val_set))
#print(len(test_set))
#showDataLayout()
data_dict= train_set[18]
print(f"class in dataset sample = {list(data_dict.keys())}")
points = data_dict['points']
seg_labels = data_dict['seg_labels']
print(f"class_id = {data_dict['class_id']}, class_name = {data_dict['class_name']}")
# Create a 3D scatter plot using plotly
pc_plots = go.Scatter3d(
x=points[:, 0],
y=points[:, 1],
z=points[:, 2],
mode='markers',
marker=dict(
size=2,
color=colors, # Set color to the colors array
opacity=0.8
)
)
# Set layout for the plot
layout = dict(
template="plotly_dark",
title=f"{data_dict['class_name']}, class id = {data_dict['class_id']}, Shapenet Dataset",
scene=PCD_SCENE,
title_x=0.5
)
# Create the figure and show the plot
fig = go.Figure(data=[pc_plots], layout=layout)
#fig.show()
# processing a list of batch items, each containing information about a 3D point cloud
def collate_fn(batch_list):
ret = {}
ret['class_id'] = torch.from_numpy(np.array([x['class_id'] for x in batch_list])).long()
ret['class_name'] = np.array([x['class_name'] for x in batch_list])
ret['points'] = torch.from_numpy(np.stack([x['points'] for x in batch_list], axis=0)).float()
ret['seg_labels'] = torch.from_numpy(np.stack([x['seg_labels'] for x in batch_list], axis=0)).long()
return ret
sample_loader = torch.utils.data.DataLoader(train_set, batch_size=16, num_workers=2, shuffle=True, collate_fn=collate_fn)
dataloader_iter = iter(sample_loader)
batch_dict = next(dataloader_iter)
print(batch_dict.keys())
for key in ['points','seg_labels', 'class_id']:
print(f"batch_dict[{key}].shape = {batch_dict[key].shape}")
batchSize= 64
workers = 2
train_loader = torch.utils.data.DataLoader(train_set, batch_size=batchSize, shuffle=True, num_workers=workers, collate_fn=collate_fn)
val_loader = torch.utils.data.DataLoader(val_set, batch_size=batchSize, shuffle=True, num_workers=workers, collate_fn=collate_fn)
test_loader = torch.utils.data.DataLoader(test_set, batch_size=batchSize,shuffle=True, num_workers=workers, collate_fn=collate_fn)
# T-Net
class STN3d(nn.Module):
def __init__(self, num_points = 2500):
super(STN3d, self).__init__()
self.num_points = num_points
self.conv1 = torch.nn.Conv1d(3, 64, 1)
self.conv2 = torch.nn.Conv1d(64, 128, 1)
self.conv3 = torch.nn.Conv1d(128, 1024, 1)
self.mp1 = torch.nn.MaxPool1d(num_points)
# FC layers
self.fc1 = nn.Linear(1024, 512)
self.fc2 = nn.Linear(512, 256)
self.fc3 = nn.Linear(256, 9)
self.relu = nn.ReLU()
self.bn1 = nn.BatchNorm1d(64)
self.bn2 = nn.BatchNorm1d(128)
self.bn3 = nn.BatchNorm1d(1024)
self.bn4 = nn.BatchNorm1d(512)
self.bn5 = nn.BatchNorm1d(256)
def forward(self, x):
batchsize = x.size()[0]
# Expected input shape = (bs, 3, num_points)
x = F.relu(self.bn1(self.conv1(x)))
x = F.relu(self.bn2(self.conv2(x)))
x = F.relu(self.bn3(self.conv3(x)))
x = self.mp1(x)
x = x.view(-1, 1024)
x = F.relu(self.bn4(self.fc1(x)))
x = F.relu(self.bn5(self.fc2(x)))
x = self.fc3(x)
iden = Variable(torch.from_numpy(np.array([1,0,0,0,1,0,0,0,1]).astype(np.float32))).view(1,9).repeat(batchsize,1)
if x.is_cuda:
iden = iden.cuda()
x = x + iden
x = x.view(-1, 3, 3)
return x
test_model = STN3d().to(device)
sim_data = Variable(torch.rand(32,3,2500)).to(device)
out = test_model(sim_data)
print('stn', out.size())
# OpenShape
class OpenShape(nn.Module):
def __init__(self, num_points = 2500):
super(OpenShape, self).__init__()
self.conv1 = torch.nn.Conv1d(3, 64, 1)
self.conv2 = torch.nn.Conv1d(64, 128, 1)
self.conv3 = torch.nn.Conv1d(128, 1024, 1)
self.mp1 = torch.nn.MaxPool1d(num_points)
def forward(self, x):
x = F.elu(self.conv1(x))
x = F.elu(self.conv2(x))
x = F.elu(self.conv3(x))
x = self.mp1(x)
return x
test_model = OpenShape().to(device)
sim_data = Variable(torch.rand(32,3,2500)).to(device)
out = test_model(sim_data)
print('stn', out.size())
#FeatureNet
class PointNetfeat(nn.Module):
def __init__(self, num_points = 2500, global_feat = True):
super(PointNetfeat, self).__init__()
self.stn = STN3d(num_points = num_points)
self.conv1 = torch.nn.Conv1d(3, 64, 1)
self.conv2 = torch.nn.Conv1d(64, 128, 1)
self.conv3 = torch.nn.Conv1d(128, 1024, 1)
self.bn1 = nn.BatchNorm1d(64)
self.bn2 = nn.BatchNorm1d(128)
self.bn3 = nn.BatchNorm1d(1024)
self.mp1 = torch.nn.MaxPool1d(num_points)
self.num_points = num_points
self.global_feat = global_feat
self.OpenShape = OpenShape(num_points = num_points)
def forward(self, x):
batchsize = x.size()[0]
trans = self.stn(x)
x = x.transpose(2,1)
x = torch.bmm(x, trans)
x = x.transpose(2,1)
x = F.elu(self.bn1(self.conv1(x)))
pointfeat = x
x = F.relu(self.bn2(self.conv2(x)))
x = self.bn3(self.conv3(x))
x = self.mp1(x)
x = x.view(-1, 1024)
if self.global_feat:
return x, trans
else:
x = x.view(-1, 1024, 1).repeat(1, 1, self.num_points)
return torch.cat([x, pointfeat], 1), trans
pointfeat = PointNetfeat(global_feat=True).to(device)
out, _ = pointfeat(sim_data)
print('global feat', out.size())
pointfeat = PointNetfeat(global_feat=False).to(device)
out, _ = pointfeat(sim_data)
print('point feat', out.size())
#Train
def dice_score(pred, target):
smooth = 1e-8
intersection = torch.sum(pred * target)
union = torch.sum(pred) + torch.sum(target)
dice = (2. * intersection + smooth) / (union + smooth)
return dice.item()
def train_model(model, num_epochs, criterion, optimizer, dataloader_train,
label_str='class_id', lr_scheduler=None, output_name='pointnet.pth'):
# move model to device
model.to(device)
for epoch in range(num_epochs):
print(f"Starting {epoch + 1} epoch ...")
# Training
model.train()
train_loss = 0.0
correct_predictions = 0
total_samples = 0
dice_scores = 0.0
y_true = []
y_scores = []
for batch_dict in tqdm(dataloader_train, total=len(dataloader_train)):
# Forward pass
x = batch_dict['points'].transpose(1, 2).to(device)
labels = batch_dict[label_str].to(device)
pred, _,_ = model(x)
loss = criterion(pred, labels)
train_loss += loss.item()
# Calculate accuracy
_, predicted_labels = torch.max(pred, 1)
correct_predictions += (predicted_labels == labels).sum().item()
total_samples += labels.size(0)
# Calculate Dice Score
pred_binary = F.softmax(pred, dim=1)[:, 1] > 0.5 # Assuming binary classification
dice_scores += dice_score(pred_binary, labels)
# Backward pass
loss.backward()
optimizer.step()
optimizer.zero_grad()
# adjust learning rate
if lr_scheduler is not None:
lr_scheduler.step()
# compute per epoch losses, accuracy, Dice Score, and AUC
train_loss = train_loss / len(dataloader_train)
accuracy = correct_predictions / total_samples
avg_dice_score = dice_scores / len(dataloader_train)
print(f'Epoch: {epoch + 1}, Train Loss: {train_loss:6.5f}, Dice Score: {avg_dice_score:.4f}')
torch.save(model.state_dict(), output_name)
class PointNetDenseSeg(nn.Module):
def __init__(self, num_points=2500, num_classes=16, num_instances=10):
super(PointNetDenseSeg, self).__init__()
self.num_points = num_points
self.num_classes = num_classes
self.num_instances = num_instances
# Feature extraction using PointNetfeat
self.feat = PointNetfeat(num_points, global_feat=False)
# Semantic segmentation layers
self.conv1_sem = torch.nn.Conv1d(1088, 512, 1)
self.conv2_sem = torch.nn.Conv1d(512, 256, 1)
self.conv3_sem = torch.nn.Conv1d(256, 128, 1)
self.conv4_sem = torch.nn.Conv1d(128, num_classes, 1)
self.bn1_sem = nn.BatchNorm1d(512)
self.bn2_sem = nn.BatchNorm1d(256)
self.bn3_sem = nn.BatchNorm1d(128)
# Instance segmentation layers
self.conv1_inst = torch.nn.Conv1d(1088, 512, 1)
self.conv2_inst = torch.nn.Conv1d(512, 256, 1)
self.conv3_inst = torch.nn.Conv1d(256, 128, 1)
self.conv4_inst = torch.nn.Conv1d(128, num_instances, 1)
self.bn1_inst = nn.BatchNorm1d(512)
self.bn2_inst = nn.BatchNorm1d(256)
self.bn3_inst = nn.BatchNorm1d(128)
def forward(self, x):
batchsize = x.size()[0]
x, trans = self.feat(x)
# Semantic segmentation branch
x_sem = F.relu(self.bn1_sem(self.conv1_sem(x)))
x_sem = F.relu(self.bn2_sem(self.conv2_sem(x_sem)))
x_sem = F.relu(self.bn3_sem(self.conv3_sem(x_sem)))
x_sem = self.conv4_sem(x_sem)
# Instance segmentation branch
x_inst = F.relu(self.bn1_inst(self.conv1_inst(x)))
x_inst = F.relu(self.bn2_inst(self.conv2_inst(x_inst)))
x_inst = F.relu(self.bn3_inst(self.conv3_inst(x_inst)))
x_inst = self.conv4_inst(x_inst)
return x_sem, x_inst, trans
sseg = PointNetDenseSeg().to(device)
out, _,_ = sseg(sim_data)
print('sseg', out.size())
epoch = 10
weight_decay = 1e-4
max_lr = 1e-3
num_points = 2500
num_classes = 16
criterion = nn.CrossEntropyLoss()
# create model, optimizer, lr_scheduler and pass to training function
num_classes = len(class_id_name_map.items())
classifier = PointNetDenseSeg(num_points = num_points)
# DEFINE OPTIMIZERS
optimizer = torch.optim.AdamW(classifier.parameters(), lr=max_lr, weight_decay=weight_decay)
if torch.cuda.is_available():
classifier.cuda()
#Train
# history = train_model(classifier, epoch, criterion, optimizer, train_loader,
# label_str='seg_labels', output_name='pointnet_se_seg.pth')
# Model Test
classifier.load_state_dict(torch.load('pointnet_se_seg.pth'))
classifier.eval()
total_loss = 0.0
correct_predictions = 0
total_samples = 0
with torch.no_grad():
for batch_dict in tqdm(test_loader, total=len(test_loader)):
x = batch_dict['points'].transpose(1, 2).to(device)
labels = batch_dict['seg_labels'].to(device)
pred, _,_ = classifier(x)
# calculate loss
loss = criterion(pred, labels)
total_loss += loss.item()
# calculate accuracy
_, predicted = torch.max(pred, 1)
correct_predictions += (predicted == labels).sum().item()
total_samples += labels.size(0)
# Print ground truth and predicted labels for each sample in the batch
for i in range(labels.size(0)): # iterate over each sample in the batch
gt_label = labels[i].item() # ground truth label
pred_label = predicted[i].item() # predicted label
print(f"Ground truth label: {gt_label}, Predicted label: {pred_label}")
evaluation_loss = total_loss / len(test_loader)
accuracy = correct_predictions / total_samples
print(evaluation_loss)
print(f"Evaluation Loss: {evaluation_loss:.4f}")
# # Now Classifie a point cloud
# point_cloud_path = 'Shapenetcore_benchmark/02691156/points/1d269dbde96f067966cf1b4a8fc3914e.npy'
# point_cloud = np.load(point_cloud_path)
# # Ensure the point cloud is of size [num_points, 3]
# if point_cloud.shape[0] < 2500:
# # If fewer points, replicate the points to reach 2500
# additional_points = np.random.choice(point_cloud.shape[0], 2500 - point_cloud.shape[0], replace=True)
# point_cloud = np.concatenate((point_cloud, point_cloud[additional_points]), axis=0)
# else:
# # If more points, randomly sample 2500 points
# point_cloud = point_cloud[np.random.choice(point_cloud.shape[0], 2500, replace=False)]
# # Reshape the point cloud to match the input shape [batch_size, 3, num_points]
# point_cloud = point_cloud.T # Shape (3, num_points)
# point_cloud = point_cloud[np.newaxis, :, :] # Add batch dimension [1, 3, num_points]
# # Convert point cloud to torch tensor
# point_cloud = torch.tensor(point_cloud, dtype=torch.float32).to(device)
# # Perform inference
# with torch.no_grad():
# pred_sem, pred_inst, _ = classifier(point_cloud)
# # You can get the semantic segmentation results (for each point) by selecting the class with the highest score
# pred_class = torch.argmax(pred_sem, dim=1) # Get class predictions for each point (semantic segmentation)
# # Now, to classify the entire point cloud (global classification)
# # You can use the class with the highest average probability across all points
# class_predictions = torch.argmax(pred_sem, dim=2) # (1, num_classes) (global class)
# _, predicted_class = torch.max(class_predictions, 1)
# # Get the predicted class as a scalar (if needed) for the first item in the batch
# predicted_class = predicted_class.item() # This gives a scalar for the first sample in the batch
# # Convert the predicted class id to the actual class name
# predicted_class_name = class_id_name_map[predicted_class]
# print(f"Predicted class: {predicted_class_name} (ID: {predicted_class})")Editor is loading...
Leave a Comment