CoolFace
Apppublic

Abs6187/ISL_Sign_Language_Translation

sourceHugging Facemitupdated 1y agoView on Hugging Face
0likes
pose_models.py360 linesDownload Raw Back to root
1"""
2ISL Sign Language Translation - TechMatrix Solvers Initiative
3Model definitions for body pose and hand pose estimation
4Developed by: TechMatrix Solvers Team
5"""
6
7import torch
8from collections import OrderedDict
9import torch.nn as nn
10
11
12def construct_layers(layer_config, no_relu_layers, prelu_layers=[]):
13    """
14    Constructs neural network layers based on configuration
15    
16    Args:
17        layer_config: Dictionary defining layer parameters
18        no_relu_layers: List of layers that shouldn't have ReLU activation
19        prelu_layers: List of layers that should use PReLU instead of ReLU
20    """
21    layers = []
22    
23    for layer_name, params in layer_config.items():
24        if 'pool' in layer_name:
25            layer = nn.MaxPool2d(kernel_size=params[0], stride=params[1], padding=params[2])
26            layers.append((layer_name, layer))
27        else:
28            conv2d = nn.Conv2d(
29                in_channels=params[0], 
30                out_channels=params[1],
31                kernel_size=params[2], 
32                stride=params[3],
33                padding=params[4]
34            )
35            layers.append((layer_name, conv2d))
36            
37            if layer_name not in no_relu_layers:
38                if layer_name not in prelu_layers:
39                    layers.append(('relu_' + layer_name, nn.ReLU(inplace=True)))
40                else:
41                    layers.append(('prelu' + layer_name[4:], nn.PReLU(params[1])))
42
43    return nn.Sequential(OrderedDict(layers))
44
45
46def construct_multi_conv_layers(layer_config, no_relu_layers):
47    """
48    Constructs multiple convolution layers for complex architectures
49    """
50    modules = []
51    for layer_name, params in layer_config.items():
52        layers = []
53        if 'pool' in layer_name:
54            layer = nn.MaxPool2d(kernel_size=params[0], stride=params[1], padding=params[2])
55            layers.append((layer_name, layer))
56        else:
57            conv2d = nn.Conv2d(
58                in_channels=params[0], 
59                out_channels=params[1],
60                kernel_size=params[2], 
61                stride=params[3],
62                padding=params[4]
63            )
64            layers.append((layer_name, conv2d))
65            if layer_name not in no_relu_layers:
66                layers.append(('Mprelu' + layer_name[5:], nn.PReLU(params[1])))
67        modules.append(nn.Sequential(OrderedDict(layers)))
68    return nn.ModuleList(modules)
69
70
71class BodyPose25Model(nn.Module):
72    """
73    Body pose estimation model using 25-point skeleton
74    Developed by TechMatrix Solvers for ISL translation
75    """
76    
77    def __init__(self):
78        super(BodyPose25Model, self).__init__()
79        
80        # Define layers without ReLU activation
81        no_relu_layers = [
82            'Mconv7_stage0_L1', 'Mconv7_stage0_L2',
83            'Mconv7_stage1_L1', 'Mconv7_stage1_L2',
84            'Mconv7_stage2_L2', 'Mconv7_stage3_L2'
85        ]
86        prelu_layers = ['conv4_2', 'conv4_3_CPM', 'conv4_4_CPM']
87        
88        # Initial feature extraction layers
89        base_layers = OrderedDict([
90            ('conv1_1', [3, 64, 3, 1, 1]),
91            ('conv1_2', [64, 64, 3, 1, 1]),
92            ('pool1_stage1', [2, 2, 0]),
93            ('conv2_1', [64, 128, 3, 1, 1]),
94            ('conv2_2', [128, 128, 3, 1, 1]),
95            ('pool2_stage1', [2, 2, 0]),
96            ('conv3_1', [128, 256, 3, 1, 1]),
97            ('conv3_2', [256, 256, 3, 1, 1]),
98            ('conv3_3', [256, 256, 3, 1, 1]),
99            ('conv3_4', [256, 256, 3, 1, 1]),
100            ('pool3_stage1', [2, 2, 0]),
101            ('conv4_1', [256, 512, 3, 1, 1]),
102            ('conv4_2', [512, 512, 3, 1, 1]),
103            ('conv4_3_CPM', [512, 256, 3, 1, 1]),
104            ('conv4_4_CPM', [256, 128, 3, 1, 1])
105        ])
106        self.base_model = construct_layers(base_layers, no_relu_layers, prelu_layers)
107        
108        # Multi-stage refinement blocks
109        stage_blocks = {}
110        
111        # L2 branch - Stage 0
112        stage_blocks['Mconv1_stage0_L2'] = OrderedDict([
113            ('Mconv1_stage0_L2_0', [128, 96, 3, 1, 1]),
114            ('Mconv1_stage0_L2_1', [96, 96, 3, 1, 1]),
115            ('Mconv1_stage0_L2_2', [96, 96, 3, 1, 1])
116        ])
117        
118        for i in range(2, 6):
119            stage_blocks[f'Mconv{i}_stage0_L2'] = OrderedDict([
120                (f'Mconv{i}_stage0_L2_0', [288, 96, 3, 1, 1]),
121                (f'Mconv{i}_stage0_L2_1', [96, 96, 3, 1, 1]),
122                (f'Mconv{i}_stage0_L2_2', [96, 96, 3, 1, 1])
123            ])
124            
125        stage_blocks['Mconv6_7_stage0_L2'] = OrderedDict([
126            ('Mconv6_stage0_L2', [288, 256, 1, 1, 0]),
127            ('Mconv7_stage0_L2', [256, 52, 1, 1, 0])
128        ])
129        
130        # L2 branch - Stages 1-3
131        for stage in range(1, 4):
132            stage_blocks[f'Mconv1_stage{stage}_L2'] = OrderedDict([
133                (f'Mconv1_stage{stage}_L2_0', [180, 128, 3, 1, 1]),
134                (f'Mconv1_stage{stage}_L2_1', [128, 128, 3, 1, 1]),
135                (f'Mconv1_stage{stage}_L2_2', [128, 128, 3, 1, 1])
136            ])
137            for i in range(2, 6):
138                stage_blocks[f'Mconv{i}_stage{stage}_L2'] = OrderedDict([
139                    (f'Mconv{i}_stage{stage}_L2_0', [384, 128, 3, 1, 1]),
140                    (f'Mconv{i}_stage{stage}_L2_1', [128, 128, 3, 1, 1]),
141                    (f'Mconv{i}_stage{stage}_L2_2', [128, 128, 3, 1, 1])
142                ])
143            stage_blocks[f'Mconv6_7_stage{stage}_L2'] = OrderedDict([
144                (f'Mconv6_stage{stage}_L2', [384, 512, 1, 1, 0]),
145                (f'Mconv7_stage{stage}_L2', [512, 52, 1, 1, 0])
146            ])
147        
148        # L1 branch configurations
149        stage_blocks['Mconv1_stage0_L1'] = OrderedDict([
150            ('Mconv1_stage0_L1_0', [180, 96, 3, 1, 1]),
151            ('Mconv1_stage0_L1_1', [96, 96, 3, 1, 1]),
152            ('Mconv1_stage0_L1_2', [96, 96, 3, 1, 1])
153        ])
154        
155        for i in range(2, 6):
156            stage_blocks[f'Mconv{i}_stage0_L1'] = OrderedDict([
157                (f'Mconv{i}_stage0_L1_0', [288, 96, 3, 1, 1]),
158                (f'Mconv{i}_stage0_L1_1', [96, 96, 3, 1, 1]),
159                (f'Mconv{i}_stage0_L1_2', [96, 96, 3, 1, 1])
160            ])
161            
162        stage_blocks['Mconv6_7_stage0_L1'] = OrderedDict([
163            ('Mconv6_stage0_L1', [288, 256, 1, 1, 0]),
164            ('Mconv7_stage0_L1', [256, 26, 1, 1, 0])
165        ])
166        
167        stage_blocks['Mconv1_stage1_L1'] = OrderedDict([
168            ('Mconv1_stage1_L1_0', [206, 128, 3, 1, 1]),
169            ('Mconv1_stage1_L1_1', [128, 128, 3, 1, 1]),
170            ('Mconv1_stage1_L1_2', [128, 128, 3, 1, 1])
171        ])
172        
173        for i in range(2, 6):
174            stage_blocks[f'Mconv{i}_stage1_L1'] = OrderedDict([
175                (f'Mconv{i}_stage1_L1_0', [384, 128, 3, 1, 1]),
176                (f'Mconv{i}_stage1_L1_1', [128, 128, 3, 1, 1]),
177                (f'Mconv{i}_stage1_L1_2', [128, 128, 3, 1, 1])
178            ])
179            
180        stage_blocks['Mconv6_7_stage1_L1'] = OrderedDict([
181            ('Mconv6_stage1_L1', [384, 512, 1, 1, 0]),
182            ('Mconv7_stage1_L1', [512, 26, 1, 1, 0])
183        ])
184        
185        # Build multi-conv modules
186        for block_name in stage_blocks.keys():
187            stage_blocks[block_name] = construct_multi_conv_layers(stage_blocks[block_name], no_relu_layers)
188        
189        self.stage_models = nn.ModuleDict(stage_blocks)
190        
191        # Freeze parameters for efficiency
192        for param in self.parameters():
193            param.requires_grad = False
194            
195    def _multi_conv_forward(self, x, models):
196        """Forward pass through multi-convolution blocks"""
197        outputs = []
198        current_output = x
199        for model in models:
200            current_output = model(current_output)
201            outputs.append(current_output)
202        return torch.cat(outputs, 1)
203        
204    def forward(self, x):
205        """Forward pass through the body pose model"""
206        base_features = self.base_model(x)
207        
208        # L2 branch processing
209        current_features = base_features
210        for stage in range(4):
211            current_features = self._multi_conv_forward(
212                current_features, self.stage_models[f'Mconv1_stage{stage}_L2']
213            )
214            for layer in range(2, 6):
215                current_features = self._multi_conv_forward(
216                    current_features, self.stage_models[f'Mconv{layer}_stage{stage}_L2']
217                )
218            current_features = self.stage_models[f'Mconv6_7_stage{stage}_L2'][0](current_features)
219            current_features = self.stage_models[f'Mconv6_7_stage{stage}_L2'][1](current_features)
220            l2_output = current_features
221            current_features = torch.cat([base_features, current_features], 1)
222        
223        # L1 branch - Stage 0
224        current_features = self._multi_conv_forward(
225            current_features, self.stage_models['Mconv1_stage0_L1']
226        )
227        for layer in range(2, 6):
228            current_features = self._multi_conv_forward(
229                current_features, self.stage_models[f'Mconv{layer}_stage0_L1']
230            )
231        current_features = self.stage_models['Mconv6_7_stage0_L1'][0](current_features)
232        current_features = self.stage_models['Mconv6_7_stage0_L1'][1](current_features)
233        stage0_l1_output = current_features
234        current_features = torch.cat([base_features, stage0_l1_output, l2_output], 1)
235        
236        # L1 branch - Stage 1
237        current_features = self._multi_conv_forward(
238            current_features, self.stage_models['Mconv1_stage1_L1']
239        )
240        for layer in range(2, 6):
241            current_features = self._multi_conv_forward(
242                current_features, self.stage_models[f'Mconv{layer}_stage1_L1']
243            )
244        current_features = self.stage_models['Mconv6_7_stage1_L1'][0](current_features)
245        stage1_l1_output = self.stage_models['Mconv6_7_stage1_L1'][1](current_features)
246        
247        return l2_output, stage1_l1_output
248
249
250class HandPoseModel(nn.Module):
251    """
252    Hand pose estimation model using 21-point hand landmarks
253    Developed by TechMatrix Solvers for ISL translation
254    """
255    
256    def __init__(self):
257        super(HandPoseModel, self).__init__()
258        
259        # Layers without ReLU activation
260        no_relu_layers = [
261            'conv6_2_CPM', 'Mconv7_stage2', 'Mconv7_stage3',
262            'Mconv7_stage4', 'Mconv7_stage5', 'Mconv7_stage6'
263        ]
264        
265        # Stage 1 - Feature extraction
266        stage1_base = OrderedDict([
267            ('conv1_1', [3, 64, 3, 1, 1]),
268            ('conv1_2', [64, 64, 3, 1, 1]),
269            ('pool1_stage1', [2, 2, 0]),
270            ('conv2_1', [64, 128, 3, 1, 1]),
271            ('conv2_2', [128, 128, 3, 1, 1]),
272            ('pool2_stage1', [2, 2, 0]),
273            ('conv3_1', [128, 256, 3, 1, 1]),
274            ('conv3_2', [256, 256, 3, 1, 1]),
275            ('conv3_3', [256, 256, 3, 1, 1]),
276            ('conv3_4', [256, 256, 3, 1, 1]),
277            ('pool3_stage1', [2, 2, 0]),
278            ('conv4_1', [256, 512, 3, 1, 1]),
279            ('conv4_2', [512, 512, 3, 1, 1]),
280            ('conv4_3', [512, 512, 3, 1, 1]),
281            ('conv4_4', [512, 512, 3, 1, 1]),
282            ('conv5_1', [512, 512, 3, 1, 1]),
283            ('conv5_2', [512, 512, 3, 1, 1]),
284            ('conv5_3_CPM', [512, 128, 3, 1, 1])
285        ])
286
287        stage1_prediction = OrderedDict([
288            ('conv6_1_CPM', [128, 512, 1, 1, 0]),
289            ('conv6_2_CPM', [512, 22, 1, 1, 0])
290        ])
291
292        stage_blocks = {}
293        stage_blocks['stage1_base'] = stage1_base
294        stage_blocks['stage1_prediction'] = stage1_prediction
295
296        # Stages 2-6 refinement
297        for i in range(2, 7):
298            stage_blocks[f'stage{i}'] = OrderedDict([
299                (f'Mconv1_stage{i}', [150, 128, 7, 1, 3]),
300                (f'Mconv2_stage{i}', [128, 128, 7, 1, 3]),
301                (f'Mconv3_stage{i}', [128, 128, 7, 1, 3]),
302                (f'Mconv4_stage{i}', [128, 128, 7, 1, 3]),
303                (f'Mconv5_stage{i}', [128, 128, 7, 1, 3]),
304                (f'Mconv6_stage{i}', [128, 128, 1, 1, 0]),
305                (f'Mconv7_stage{i}', [128, 22, 1, 1, 0])
306            ])
307
308        # Build all stage models
309        for block_name in stage_blocks.keys():
310            stage_blocks[block_name] = construct_layers(stage_blocks[block_name], no_relu_layers)
311
312        self.stage1_base_model = stage_blocks['stage1_base']
313        self.stage1_prediction_model = stage_blocks['stage1_prediction']
314        self.stage2_model = stage_blocks['stage2']
315        self.stage3_model = stage_blocks['stage3']
316        self.stage4_model = stage_blocks['stage4']
317        self.stage5_model = stage_blocks['stage5']
318        self.stage6_model = stage_blocks['stage6']
319        
320        # Freeze parameters for efficiency
321        for param in self.parameters():
322            param.requires_grad = False
323
324    def forward(self, x):
325        """Forward pass through the hand pose model"""
326        base_features = self.stage1_base_model(x)
327        stage1_output = self.stage1_prediction_model(base_features)
328        
329        # Stage 2
330        stage2_input = torch.cat([stage1_output, base_features], 1)
331        stage2_output = self.stage2_model(stage2_input)
332        
333        # Stage 3
334        stage3_input = torch.cat([stage2_output, base_features], 1)
335        stage3_output = self.stage3_model(stage3_input)
336        
337        # Stage 4
338        stage4_input = torch.cat([stage3_output, base_features], 1)
339        stage4_output = self.stage4_model(stage4_input)
340        
341        # Stage 5
342        stage5_input = torch.cat([stage4_output, base_features], 1)
343        stage5_output = self.stage5_model(stage5_input)
344        
345        # Stage 6
346        stage6_input = torch.cat([stage5_output, base_features], 1)
347        stage6_output = self.stage6_model(stage6_input)
348        
349        return stage6_output
350
351
352# Factory functions for easy model instantiation
353def create_bodypose_model():
354    """Create and return body pose detection model"""
355    return BodyPose25Model()
356
357
358def create_handpose_model():
359    """Create and return hand pose detection model"""
360    return HandPoseModel()