app.py

import torch
from transformers import AutoModelForSequenceClassification, AutoTokenizer, AutoConfig
from torchvision import models, transforms
import torch.nn as nn
import os
import json
import cv2
from PIL import Image
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=2)
        
        # Visual model (ResNet50)
        self.visual_model = models.resnet50(weights=models.ResNet50_Weights.DEFAULT)
        visual_feature_dim = self.visual_model.fc.in_features
        self.visual_model.fc = nn.Identity()

        # Fusion and classification layer setup
        text_feature_dim = self.text_model.config.hidden_size
        self.fc1 = nn.Linear(text_feature_dim + visual_feature_dim, hidden_dim)
        self.dropout = nn.Dropout(dropout)
        self.fc2 = nn.Linear(hidden_dim, 1)

    def forward(self, encoding, frames):
        input_ids = encoding['input_ids'].squeeze(1).to(device)
        attention_mask = encoding['attention_mask'].squeeze(1).to(device)

        # Extract text and visual features
        text_features = self.text_model(input_ids=input_ids, attention_mask=attention_mask).logits
        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)
        visual_features = visual_features.view(batch_size, num_frames, -1).mean(dim=1)

        # Combine and classify
        combined_features = torch.cat((text_features, visual_features), dim=1)
        x = self.dropout(torch.relu(self.fc1(combined_features)))
        output = torch.sigmoid(self.fc2(x))

        return output

    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')

            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:
            hf_model = AutoModelForSequenceClassification.from_pretrained(load_directory, num_labels=2)
            model = cls(text_model_name=hf_model.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') # 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)
    frames = []
    frame_count = 0
    while frame_count < num_frames:  # Limit to a number of frames for efficiency
        success, frame = cap.read()
        if not success:
            break
        # Convert frame (NumPy array) to PIL image and apply transformations
        frame = Image.fromarray(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB))
        frame = transform(frame)
        frames.append(frame)
        frame_count += 1
    cap.release()

    # Stack frames and add batch dimension (1, num_frames, channels, height, width)
    frames = torch.stack(frames)
    frames = frames.unsqueeze(0)  # Add batch dimension
    return frames

def predict_video(model, video_path, text_input, tokenizer, transform):
    try:
        # Set model to evaluation mode
        model.eval()
        
        # Tokenize the text input
        encoding = tokenizer(
            text_input, padding='max_length', truncation=True, max_length=128, return_tensors='pt'
        )
        encoding = {key: val.to(device) for key, val in encoding.items()}
        
        # Load frames from the video
        frames = load_frames_from_video(video_path, transform)
        frames = frames.to(device)
        
        # Log input shapes and devices
        print(f"Encoding device: {next(iter(encoding.values())).device}, Frames shape: {frames.shape}")
        
        # Perform forward pass through the model
        with torch.no_grad():
            output = model(encoding, frames)
        
        # Apply sigmoid to get probability, then threshold to get prediction
        prediction = (output.squeeze(-1) > 0.5).float()
        
        return prediction.item()
    
    except Exception as e:
        print(f"Prediction error: {e}")
        return "Error during prediction"




transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])


# Define your video paths and captions
video_paths = [
    'https://drive.google.com/uc?export=download&id=1iWq1q1LM-jmf4iZxOqZTw4FaIBekJowM',
    'https://drive.google.com/uc?export=download&id=1_egBaC1HD2kIZgRRKsnCtsWG94vg1c7n',
    'https://drive.google.com/uc?export=download&id=12cGxBEkfU5Q1Ezg2jRk6zGyn2hoR3JLj'
]

video_captions = [
    "Everytime i start a diet  كل مرة أحاول أبدأ ريجيم 😓 #dietmemes #funnyvideos #animetiktok",
    "New sandwich from burger king 🍔👑 #mukbang #asmr #asmrmukbang #asmrsounds #eat #food #Foodie moe eats #yummy #cheese #chicken #burger #fries #burgerking @Burger King",
    "all workout guides l!nked in bi0 // honestly huge moment 😂 I’ve been so focused on growing my upper body that this feels like it finally shows! shorts from @KEEPTHATPUMP #upperbody #upperbodyworkout #glutegains #glutegrowth #gluteexercise #workout #strengthtraining #gym #trending #fyp"
]


def predict_risk(video_index):
    video_path = video_paths[video_index]
    text_input = video_captions[video_index]
    
    # Make prediction
    prediction = predict_video(model, video_path, text_input, tokenizer, transform)
    
    # Return the corresponding label
    if prediction == "Error during prediction":
        return "Error during prediction"
    return "Risky Health Behavior" if prediction == 1 else "Not Risky Health Behavior"

# Interface setup
with gr.Blocks() as interface:
    gr.Markdown("# Risk Behavior Prediction")
    gr.Markdown("Select a video to classify its behavior as risky or not.")
    
    # Input option selector
    video_selector = gr.Radio(["Video 1", "Video 2", "Video 3"], label="Choose a Video")

    # Use function to return URLs which are handled by the Gradio `gr.Video` component
    def show_selected_video(choice):
        idx = int(choice.split()[-1]) - 1
        return video_paths[idx], f"**Caption:** {video_captions[idx]}"

    video_player = gr.Video(width=320, height=240)
    caption_box = gr.Markdown()

    video_selector.change(
        fn=show_selected_video,
        inputs=video_selector,
        outputs=[video_player, caption_box]
    )
    
    # Prediction button and output
    predict_button = gr.Button("Predict Risk")
    output_text = gr.Textbox(label="Prediction")

    predict_button.click(
        fn=lambda idx: predict_risk(int(idx.split()[-1]) - 1),
        inputs=video_selector,
        outputs=output_text
    )

# Launch the app
interface.launch()