fossbk commited on
Commit
d3d18f9
·
verified ·
1 Parent(s): b94db6f

sử dụng mô hình phân loại video có sẵn từ Hugging Face

Browse files
Files changed (1) hide show
  1. app.py +5 -10
app.py CHANGED
@@ -1,15 +1,15 @@
1
  import gradio as gr
2
- import torch
3
  from transformers import pipeline
4
- from moviepy.editor import VideoFileClip
5
  from PIL import Image
6
- import os
7
 
8
  # Kiểm tra thiết bị sử dụng GPU hay CPU
9
  device = "cuda" if torch.cuda.is_available() else "cpu"
10
 
11
  # Tải các mô hình phân loại ảnh và video từ Hugging Face
12
  image_classifier = pipeline("image-classification", model="google/vit-base-patch16-224-in21k", device=0 if device == "cuda" else -1)
 
 
13
  video_classifier = pipeline("video-classification", model="google/vit-base-patch16-224-in21k", device=0 if device == "cuda" else -1)
14
 
15
  # Hàm phân loại ảnh
@@ -32,13 +32,8 @@ def classify_video(video, model_name):
32
  else:
33
  classifier = video_classifier # Chỉnh sửa ở đây nếu muốn hỗ trợ thêm các mô hình khác
34
 
35
- # Đọc video trích xuất các frame (ở đây đơn giản là lấy 1 frame đầu tiên)
36
- clip = VideoFileClip(video.name)
37
- frame = clip.get_frame(0)
38
- image = Image.fromarray(frame)
39
-
40
- # Phân loại frame đầu tiên của video
41
- result = classifier(image)
42
  return result[0]['label'], result[0]['score']
43
 
44
  # Giao diện Gradio
 
1
  import gradio as gr
 
2
  from transformers import pipeline
 
3
  from PIL import Image
4
+ import torch
5
 
6
  # Kiểm tra thiết bị sử dụng GPU hay CPU
7
  device = "cuda" if torch.cuda.is_available() else "cpu"
8
 
9
  # Tải các mô hình phân loại ảnh và video từ Hugging Face
10
  image_classifier = pipeline("image-classification", model="google/vit-base-patch16-224-in21k", device=0 if device == "cuda" else -1)
11
+
12
+ # Sử dụng mô hình phân loại video có sẵn trên Hugging Face
13
  video_classifier = pipeline("video-classification", model="google/vit-base-patch16-224-in21k", device=0 if device == "cuda" else -1)
14
 
15
  # Hàm phân loại ảnh
 
32
  else:
33
  classifier = video_classifier # Chỉnh sửa ở đây nếu muốn hỗ trợ thêm các mô hình khác
34
 
35
+ # Phân loại video trực tiếp không cần trích xuất frame
36
+ result = classifier(video)
 
 
 
 
 
37
  return result[0]['label'], result[0]['score']
38
 
39
  # Giao diện Gradio