pillow

app.py

@@ -2,19 +2,18 @@ import gradio as gr
 from transformers import pipeline
 from PIL import Image
 import torch
+import tempfile
+import os
+from moviepy.editor import VideoFileClip
 
 # Kiểm tra thiết bị sử dụng GPU hay CPU
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# Tải các mô hình phân loại ảnh và video từ Hugging Face
+# Tải mô hình phân loại ảnh từ Hugging Face (sử dụng mô hình ảnh cho video)
 image_classifier = pipeline("image-classification", model="google/vit-base-patch16-224-in21k", device=0 if device == "cuda" else -1)
 
-# Sử dụng mô hình phân loại video có sẵn trên Hugging Face
-video_classifier = pipeline("video-classification", model="google/vit-base-patch16-224-in21k", device=0 if device == "cuda" else -1)
-
 # Hàm phân loại ảnh
 def classify_image(image, model_name):
-    # Tùy chọn chọn model ảnh khác nếu người dùng yêu cầu
     if model_name == "ViT":
         classifier = image_classifier
     else:
@@ -24,16 +23,20 @@ def classify_image(image, model_name):
     result = classifier(image)
     return result[0]['label'], result[0]['score']
 
-# Hàm phân loại video
+# Hàm phân loại video (trích xuất frame đầu tiên của video)
 def classify_video(video, model_name):
-    # Tùy chọn chọn model video khác nếu người dùng yêu cầu
+    # Trích xuất frame đầu tiên của video
+    video_clip = VideoFileClip(video.name)
+    frame = video_clip.get_frame(0)  # Lấy frame đầu tiên
+    image = Image.fromarray(frame)
+
+    # Phân loại frame đầu tiên của video
     if model_name == "ViT":
-        classifier = video_classifier
+        classifier = image_classifier
     else:
-        classifier = video_classifier  # Chỉnh sửa ở đây nếu muốn hỗ trợ thêm các mô hình khác
+        classifier = image_classifier  # Chỉnh sửa ở đây nếu muốn hỗ trợ thêm các mô hình khác
 
-    # Phân loại video trực tiếp mà không cần trích xuất frame
-    result = classifier(video)
+    result = classifier(image)
     return result[0]['label'], result[0]['score']
 
 # Giao diện Gradio