initial commit

app.py

@@ -1,6 +1,6 @@
 import torch
 from transformers import AutoModelForSequenceClassification, AutoTokenizer, AutoConfig
-from torchvision import transforms, models
+from torchvision import models, transforms
 import torch.nn as nn
 import os
 import json
@@ -11,22 +11,18 @@ import gradio as gr
 class MultimodalRiskBehaviorModel(nn.Module):
     def __init__(self, text_model_name="bert-base-uncased", hidden_dim=512, dropout=0.3):
         super(MultimodalRiskBehaviorModel, self).__init__()
-        
+
         # Text model using AutoModelForSequenceClassification
         self.text_model_name = text_model_name
         self.text_model = AutoModelForSequenceClassification.from_pretrained(text_model_name, num_labels=1)
         
         # Visual model (ResNet50)
         self.visual_model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
-        
-        # Save the original `fc.in_features` before replacing it
         visual_feature_dim = self.visual_model.fc.in_features
-        self.visual_model.fc = nn.Identity()  # Replace with identity layer
+        self.visual_model.fc = nn.Identity()
 
-        # Get the hidden dimension of the text model
+        # Fusion and classification layer setup
         text_feature_dim = self.text_model.config.hidden_size
-        
-        # Fusion and classification layers
         self.fc1 = nn.Linear(text_feature_dim + visual_feature_dim, hidden_dim)
         self.dropout = nn.Dropout(dropout)
         self.fc2 = nn.Linear(hidden_dim, 1)
@@ -34,36 +30,66 @@ class MultimodalRiskBehaviorModel(nn.Module):
     def forward(self, encoding, frames):
         input_ids = encoding['input_ids'].squeeze(1).to(device)
         attention_mask = encoding['attention_mask'].squeeze(1).to(device)
-        
-        # Text embeddings from BERT
+
+        # Extract text and visual features
         text_features = self.text_model(input_ids=input_ids, attention_mask=attention_mask).logits
-        
-        # Visual features from ResNet50
         frames = frames.to(device)
         
         batch_size, num_frames, channels, height, width = frames.size()
         frames = frames.view(batch_size * num_frames, channels, height, width)
         visual_features = self.visual_model(frames)
-        
-        # Reshape back to (batch_size, num_frames, visual_feature_dim)
         visual_features = visual_features.view(batch_size, num_frames, -1).mean(dim=1)
-        
-        # Combine text and visual features
+
+        # Combine and classify
         combined_features = torch.cat((text_features, visual_features), dim=1)
-        
-        # Pass through the classifier
         x = self.dropout(torch.relu(self.fc1(combined_features)))
         output = torch.sigmoid(self.fc2(x))
-        
+
         return output
 
-# Loading the tokenizer and model
+    def save_pretrained(self, save_directory):
+        os.makedirs(save_directory, exist_ok=True)
+        torch.save(self.state_dict(), os.path.join(save_directory, 'pytorch_model.bin'))
+        config = {
+            "text_model_name": self.text_model_name,
+            "hidden_dim": self.fc1.out_features
+        }
+        with open(os.path.join(save_directory, 'config.json'), 'w') as f:
+            json.dump(config, f)
+
+    @classmethod
+    def from_pretrained(cls, load_directory, map_location=None):
+        if os.path.exists(load_directory):
+            config_path = os.path.join(load_directory, 'config.json')
+            state_dict_path = os.path.join(load_directory, 'pytorch_model.bin')
+
+            if not os.path.exists(config_path):
+                raise FileNotFoundError(f"{config_path} not found.")
+            if not os.path.exists(state_dict_path):
+                raise FileNotFoundError(f"{state_dict_path} not found.")
+
+            with open(config_path, 'r') as f:
+                config_dict = json.load(f)
+            model = cls(text_model_name=config_dict["text_model_name"], hidden_dim=config_dict["hidden_dim"])
+            state_dict = torch.load(state_dict_path, map_location=map_location)
+            model.load_state_dict(state_dict)
+        
+        else:
+            config = AutoConfig.from_pretrained(load_directory)
+            hf_model = AutoModelForSequenceClassification.from_pretrained(load_directory, config=config)
+            model = cls(text_model_name=config.name_or_path, hidden_dim=hf_model.config.hidden_size)
+            model.text_model = hf_model
+
+        return model
+
 tokenizer = AutoTokenizer.from_pretrained('Souha-BH/BERT_Resnet50')
-model = MultimodalRiskBehaviorModel.from_pretrained('Souha-BH/BERT_Resnet50', map_location='cpu')
+model = MultimodalRiskBehaviorModel.from_pretrained('Souha-BH/BERT_Resnet50') # if cpu add arg map_location='cpu'
+
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 model.to(device)
 
 
+
 # Function to load frames from a video
 def load_frames_from_video(video_path, transform, num_frames=10):
     cap = cv2.VideoCapture(video_path)