Download handler.py from winno/video-analysis: direct link, hf CLI and curl.
- Browser
- Download file 14.8 kB
-
https://huggingface.co/winno/video-analysis/resolve/main/handler.py
- Command line
-
hf download hf://winno/video-analysis/handler.py
-
curl -L -o handler.py https://huggingface.co/winno/video-analysis/resolve/main/handler.py
14.8 kB
| # File: handler.py | |
| import base64 | |
| import io | |
| import os | |
| import tempfile | |
| from typing import Dict, Any | |
| import cv2 | |
| import librosa | |
| import numpy as np | |
| import pandas as pd | |
| import torch | |
| import whisper_timestamped as whisper | |
| from fer import FER | |
| from moviepy.editor import VideoFileClip, AudioFileClip | |
| from torch.nn.functional import softmax | |
| from transformers import AutoModelForAudioClassification, pipeline | |
| from translate import Translator | |
| class EndpointHandler: | |
| def __init__(self, path=""): | |
| """ | |
| Loads all models onto the device. This is called once when the endpoint starts. | |
| """ | |
| print("Loading models...") | |
| # Use GPU if available | |
| self.device = "cuda" if torch.cuda.is_available() else "cpu" | |
| print(f"Using device: {self.device}") | |
| # 1. Audio Emotion Model | |
| self.audio_model = AutoModelForAudioClassification.from_pretrained( | |
| "3loi/SER-Odyssey-Baseline-WavLM-Categorical-Attributes", trust_remote_code=True | |
| ).to(self.device) | |
| self.audio_mean = self.audio_model.config.mean | |
| self.audio_std = self.audio_model.config.std | |
| # 2. Facial Emotion Model | |
| self.face_detector = FER(mtcnn=True) | |
| # 3. Text Emotion Model | |
| self.text_classifier = pipeline( | |
| "text-classification", model="j-hartmann/emotion-english-distilroberta-base", top_k=None, device=self.device | |
| ) | |
| # 4. Transcription Model | |
| self.transcription_model = whisper.load_model("medium", device=self.device) | |
| # 5. Translator | |
| self.translator = Translator(from_lang='ko', to_lang='en') | |
| print("All models loaded successfully.") | |
| def __call__(self, data: Dict[str, Any]) -> Dict[str, Any]: | |
| """ | |
| Handles an inference request. | |
| Args: | |
| data (Dict[str, Any]): Dictionary containing request parameters. Expected keys: | |
| 'video': Base64 encoded video string. | |
| 'analysis_type': One of "audio", "facial", or "text". | |
| """ | |
| print("Received inference request.") | |
| # --- 1. Parameter Extraction --- | |
| if 'inputs' in data and isinstance(data['inputs'], dict): | |
| params = data['inputs'] | |
| else: | |
| params = data | |
| b64_video = params.get("video") | |
| if not b64_video: | |
| raise ValueError("Missing 'video' parameter (base64 encoded string)") | |
| analysis_type = params.get("analysis_type") | |
| if analysis_type not in ["audio", "facial", "text"]: | |
| raise ValueError("Missing or invalid 'analysis_type'. Must be 'audio', 'facial', or 'text'.") | |
| # --- 2. Video Decoding --- | |
| video_bytes = base64.b64decode(b64_video) | |
| # Use a temporary file to store the video, as the original functions expect a path | |
| with tempfile.NamedTemporaryFile(suffix=".mp4", delete=True) as temp_video_file: | |
| temp_video_file.write(video_bytes) | |
| temp_video_file.flush() # Ensure all data is written | |
| video_path = temp_video_file.name | |
| print(f"Video saved to temporary file: {video_path}") | |
| # --- 3. Dispatch to correct analysis function --- | |
| try: | |
| if analysis_type == "audio": | |
| result = self._analyze_audio_emotions(video_path) | |
| elif analysis_type == "facial": | |
| result = self._detect_faces_and_emotions(video_path) | |
| elif analysis_type == "text": | |
| result = self._process_video_text(video_path) | |
| print("Analysis completed successfully.") | |
| return {"status": "success", **result} | |
| except Exception as e: | |
| print(f"Error during {analysis_type} analysis: {e}") | |
| # It's good practice to return a structured error | |
| return {"status": "error", "message": str(e)} | |
| # =================================================================== | |
| # REFACTORED ANALYSIS FUNCTIONS | |
| # =================================================================== | |
| def _analyze_audio_emotions(self, video_path: str) -> Dict: | |
| temp_audio_path = None | |
| try: | |
| # Extract audio | |
| with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_audio_file: | |
| temp_audio_path = temp_audio_file.name | |
| VideoFileClip(video_path).audio.write_audiofile(temp_audio_path, codec="pcm_s16le", logger=None) | |
| raw_wav, _ = librosa.load(temp_audio_path, sr=self.audio_model.config.sampling_rate) | |
| norm_wav = (raw_wav - self.audio_mean) / (self.audio_std + 1e-6) | |
| times, emotions_dfs = [], [] | |
| for start_time in range(0, len(norm_wav), self.audio_model.config.sampling_rate): | |
| audio_segment = norm_wav[start_time:start_time + self.audio_model.config.sampling_rate] | |
| # Process segment | |
| audio_np = np.array(audio_segment) | |
| mask = torch.ones(1, len(audio_np)).to(self.device) | |
| wavs = torch.tensor(audio_np).unsqueeze(0).to(self.device) | |
| with torch.no_grad(): | |
| pred = self.audio_model(wavs, mask) | |
| logits = pred.logits if hasattr(pred, 'logits') else pred[0] | |
| labels = {0: 'Angry', 1: 'Sad', 2: 'Happy', 3: 'Surprise', 4: 'Fear', 5: 'Disgust', 7: 'Neutral'} | |
| probabilities = softmax(logits, dim=-1).squeeze(0)[[0, 1, 2, 3, 4, 5, 7]] | |
| probabilities = probabilities / probabilities.sum() | |
| df = pd.DataFrame([probabilities.cpu().numpy()], columns=labels.values()) | |
| times.append(start_time / self.audio_model.config.sampling_rate) | |
| emotions_dfs.append(df) | |
| emotions_df = pd.concat(emotions_dfs, ignore_index=True) | |
| emotions_df.insert(0, "Time(s)", times) | |
| emotion_rename_map = {'Angry': 'anger', 'Sad': 'sadness', 'Happy': 'happy', 'Surprise': 'surprise', | |
| 'Fear': 'fear', 'Disgust': 'disgust', 'Neutral': 'neutral'} | |
| emotions_df.rename(columns=emotion_rename_map, inplace=True) | |
| # Return DataFrame as JSON | |
| return {"emotions_data": emotions_df.to_json(orient='split')} | |
| finally: | |
| if temp_audio_path and os.path.exists(temp_audio_path): | |
| os.remove(temp_audio_path) | |
| def _detect_faces_and_emotions(self, video_path: str) -> Dict: | |
| emotions_data = [] | |
| output_video_path = None | |
| # =================================================================== | |
| # NEW: Confidence threshold for filtering false positives. | |
| # Only faces where at least one emotion has a score > 0.35 will be kept. | |
| # You can adjust this value. Higher = stricter filtering. (e.g., 0.40) | |
| # =================================================================== | |
| CONFIDENCE_THRESHOLD = 0.6 | |
| try: | |
| # Create a temporary file for the output video | |
| with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as temp_out_video: | |
| output_video_path = temp_out_video.name | |
| original_video = VideoFileClip(video_path) | |
| cap = cv2.VideoCapture(video_path) | |
| fps = int(cap.get(cv2.CAP_PROP_FPS)) | |
| if fps == 0: # Handle potential issue with video metadata | |
| fps = 30 | |
| w = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH)) | |
| h = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT)) | |
| # Use a temporary path for the video writer intermediate file | |
| with tempfile.NamedTemporaryFile(suffix=".mp4", delete=False) as temp_video_writer_file: | |
| temp_video_writer_path = temp_video_writer_file.name | |
| out = cv2.VideoWriter(temp_video_writer_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (w, h)) | |
| frame_number = 0 | |
| while cap.isOpened(): | |
| ret, frame = cap.read() | |
| if not ret: break | |
| time_seconds = round(frame_number / fps) | |
| result = self.face_detector.detect_emotions(frame) | |
| # Process each face found in the frame | |
| for face in result: | |
| # =================================================================== | |
| # NEW: Filtering logic starts here | |
| # =================================================================== | |
| emotions = face["emotions"] | |
| max_emotion_score = max(emotions.values()) | |
| # If the highest emotion score is below our threshold, skip this face | |
| if max_emotion_score < CONFIDENCE_THRESHOLD: | |
| continue # Ignore this low-confidence detection | |
| # =================================================================== | |
| # End of new filtering logic. The rest of the loop proceeds as before. | |
| # =================================================================== | |
| box = face["box"] | |
| emotions["Time(s)"] = time_seconds | |
| emotions_data.append(emotions) | |
| cv2.rectangle(frame, (box[0], box[1]), (box[0] + box[2], box[1] + box[3]), (0, 155, 255), 2) | |
| # Find the dominant emotion to display on the video | |
| dominant_emotion = max(emotions, key=lambda k: emotions[k] if k != 'Time(s)' else -1) | |
| text_to_display = f"{dominant_emotion}: {emotions[dominant_emotion]:.2f}" | |
| cv2.putText(frame, text_to_display, (box[0], box[1] - 10), | |
| cv2.FONT_HERSHEY_SIMPLEX, 0.6, (0, 255, 0), 2, cv2.LINE_AA) | |
| out.write(frame) | |
| frame_number += 1 | |
| cap.release() | |
| out.release() | |
| # Combine processed video frames with original audio | |
| processed_video_clip = VideoFileClip(temp_video_writer_path) | |
| final_clip = processed_video_clip.set_audio(original_video.audio) | |
| final_clip.write_videofile(output_video_path, codec='libx264', logger=None) | |
| os.remove(temp_video_writer_path) # Clean up intermediate video | |
| # Read the final video bytes and encode to base64 | |
| with open(output_video_path, "rb") as f: | |
| processed_video_b64 = base64.b64encode(f.read()).decode("utf-8") | |
| # Process DataFrame | |
| emotions_df = pd.DataFrame(emotions_data) | |
| if not emotions_df.empty: | |
| emotions_df['Time(s)'] = emotions_df['Time(s)'].round().astype(int) | |
| max_time = emotions_df['Time(s)'].max() | |
| all_times = pd.DataFrame({'Time(s)': range(max_time + 1)}) | |
| avg_scores = emotions_df.groupby("Time(s)").mean().reset_index() | |
| df_merged = pd.merge(all_times, avg_scores, on='Time(s)', how='left').fillna(0) | |
| df_merged['Time(s)'] = df_merged['Time(s)'].astype(str) + " sec" | |
| else: | |
| df_merged = pd.DataFrame() | |
| return { | |
| "emotions_data": df_merged.to_json(orient='split'), | |
| "processed_video": processed_video_b64 | |
| } | |
| finally: | |
| if output_video_path and os.path.exists(output_video_path): | |
| os.remove(output_video_path) | |
| def _process_video_text(self, video_path: str) -> Dict: | |
| temp_audio_path = None | |
| try: | |
| video_clip = VideoFileClip(video_path) | |
| with tempfile.NamedTemporaryFile(suffix=".wav", delete=False) as temp_audio_file: | |
| temp_audio_path = temp_audio_file.name | |
| video_clip.audio.write_audiofile(temp_audio_path, logger=None) | |
| audio = whisper.load_audio(temp_audio_path) | |
| result = whisper.transcribe(self.transcription_model, audio) | |
| segments_data = [{'text': seg['text'], 'start': seg['start'], 'end': seg['end']} for seg in | |
| result['segments']] | |
| segments_df = pd.DataFrame(segments_data) | |
| if segments_df.empty: | |
| return {"words_data": pd.DataFrame().to_json(orient='split'), | |
| "segments_data": pd.DataFrame().to_json(orient='split')} | |
| segments_df['Translated_Text'] = segments_df['text'].apply(lambda x: self.translator.translate(x)) | |
| segments_df['Sentiment_Scores'] = segments_df['Translated_Text'].apply( | |
| lambda x: {entry['label']: entry['score'] for entry in self.text_classifier(x)[0]}) | |
| sentiment_df = segments_df['Sentiment_Scores'].apply(pd.Series) | |
| final_segments_df = pd.concat([segments_df.drop(columns=['Sentiment_Scores']), sentiment_df], axis=1) | |
| # Process words data | |
| word_texts, word_starts, word_ends = [], [], [] | |
| for segment in result['segments']: | |
| for word in segment['words']: | |
| word_texts.append(word['text']) | |
| word_starts.append(word['start']) | |
| word_ends.append(word['end']) | |
| words_df = pd.DataFrame({'text': word_texts, 'start': word_starts, 'end': word_ends}) | |
| words_df['second'] = words_df['start'].apply(lambda x: int(np.ceil(x))) | |
| words_grouped = words_df.groupby('second').agg( | |
| {'text': lambda x: ' '.join(x), 'start': 'min', 'end': 'max'}).reset_index() | |
| max_second = int(video_clip.duration) | |
| all_seconds = pd.DataFrame({'second': np.arange(0, max_second + 1)}) | |
| words_grouped = all_seconds.merge(words_grouped, on='second', how='left').fillna( | |
| {'text': '', 'start': 0, 'end': 0}) | |
| emotion_columns = final_segments_df.columns.difference(['text', 'start', 'end', 'Translated_Text']) | |
| for col in emotion_columns: | |
| words_grouped[col] = np.nan | |
| for i, row in words_grouped.iterrows(): | |
| matching_segment = final_segments_df[ | |
| (final_segments_df['start'] <= row['start']) & (final_segments_df['end'] >= row['end'])] | |
| if not matching_segment.empty: | |
| for emotion in emotion_columns: | |
| words_grouped.at[i, emotion] = matching_segment.iloc[0][emotion] | |
| words_grouped[emotion_columns] = words_grouped[emotion_columns].fillna(0) | |
| return { | |
| "words_data": words_grouped.to_json(orient='split'), | |
| "segments_data": final_segments_df.to_json(orient='split') | |
| } | |
| finally: | |
| if temp_audio_path and os.path.exists(temp_audio_path): | |
| os.remove(temp_audio_path) |