Spaces:
Sleeping
Sleeping
File size: 4,274 Bytes
4d2f736 | 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 | import torch
from torch import nn
import torch.nn.functional as F
from torch.utils.data import DataLoader, TensorDataset
from torchaudio import transforms
from torchvision import models
class BlazeFace(nn.Module):
def __init__(self, input_channels=1, use_double_block=False, activation="relu", use_optional_block=True):
super(BlazeFace, self).__init__()
self.activation = activation
self.use_double_block = use_double_block
self.use_optional_block = use_optional_block
def conv_block(in_channels, out_channels, kernel_size, stride, padding):
return nn.Sequential(
nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size, stride=stride, padding=padding),
nn.BatchNorm2d(out_channels),
nn.ReLU() if activation == "relu" else nn.Sigmoid() # Apply ReLU activation (default) or Sigmoid
)
def depthwise_separable_block(in_channels, out_channels, stride):
return nn.Sequential(
nn.Conv2d(in_channels, in_channels, kernel_size=5, stride=stride, padding=2, groups=in_channels, bias=False),
nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0),
nn.BatchNorm2d(out_channels),
nn.ReLU() if activation == "relu" else nn.Sigmoid()
)
def double_block(in_channels, filters_1, filters_2, stride):
return nn.Sequential(
depthwise_separable_block(in_channels, filters_1, stride),
depthwise_separable_block(filters_1, filters_2, 1)
)
# Define layers (first part: conv layers)
self.conv1 = conv_block(input_channels, 24, kernel_size=5, stride=2, padding=2)
# Define single blocks (subsequent conv blocks)
self.single_blocks = nn.ModuleList([
depthwise_separable_block(24, 24, stride=1),
depthwise_separable_block(24, 24, stride=1),
depthwise_separable_block(24, 48, stride=2),
depthwise_separable_block(48, 48, stride=1),
depthwise_separable_block(48, 48, stride=1)
])
# Define double blocks if `use_double_block` is True
if self.use_double_block:
self.double_blocks = nn.ModuleList([
double_block(48, 24, 96, stride=2),
double_block(96, 24, 96, stride=1),
double_block(96, 24, 96, stride=2),
double_block(96, 24, 96, stride=1),
double_block(96, 24, 96, stride=2)
])
else:
self.double_blocks = nn.ModuleList([
depthwise_separable_block(48, 96, stride=2),
depthwise_separable_block(96, 96, stride=1),
depthwise_separable_block(96, 96, stride=2),
depthwise_separable_block(96, 96, stride=1),
depthwise_separable_block(96, 96, stride=2)
])
# Final convolutional head
self.conv_head = nn.Conv2d(96, 64, kernel_size=1, stride=1)
self.bn_head = nn.BatchNorm2d(64)
# Global Average Pooling
self.global_avg_pooling = nn.AdaptiveAvgPool2d(1)
def forward(self, x):
# First conv layer
x = self.conv1(x)
# Apply single blocks
for block in self.single_blocks:
x = block(x)
# Apply double blocks
for block in self.double_blocks:
x = block(x)
# Final head
x = self.conv_head(x)
x = self.bn_head(x)
x = F.relu(x)
# Global Average Pooling and Flatten
x = self.global_avg_pooling(x)
x = torch.flatten(x, 1)
return x
class BlazeFaceModel(nn.Module):
def __init__(self, input_channels, label_count, use_double_block=False, activation="relu", use_optional_block=True):
super(BlazeFaceModel, self).__init__()
self.blazeface_backbone = BlazeFace(input_channels=input_channels, use_double_block=use_double_block, activation=activation, use_optional_block=use_optional_block)
self.fc = nn.Linear(64, label_count)
def forward(self, x):
features = self.blazeface_backbone(x)
output = self.fc(features)
return output |