Update app.py

app.py

@@ -1,135 +1,95 @@
-import os
-import gradio as gr
-from groq import Groq
-from huggingface_hub import InferenceClient
-from moviepy.editor import VideoFileClip
+import cv2
+import mediapipe as mp
 import numpy as np
-from io import BytesIO
+import gradio as gr
+import base64
+import time
 
-# Lưu khóa API
-if 'GROQ_API_KEY' not in os.environ:
-    os.environ['GROQ_API_KEY'] = input('Nhập khóa API Groq của bạn: ')
-if 'HF_TOKEN' not in os.environ:
-    os.environ['HF_TOKEN'] = input('Nhập Hugging Face API Token của bạn: ')
+# Initialize MediaPipe Selfie Segmentation
+mp_selfie_segmentation = mp.solutions.selfie_segmentation
+segmentation = mp_selfie_segmentation.SelfieSegmentation(model_selection=1)
 
-MAX_VIDEO_SIZE_MB = 35
+# Global settings
+settings = {
+    "seg_enabled": True,
+    "blur_bg": False,
+    "set_bg": False,
+    "set_color": False,
+    "bg_color": (0, 0, 0),  # BGR
+    "blur_intensity": 15
+}
+bg_image = None
 
-# Hàm gọi Groq API
-def call_groq_api(prompt, model_name="meta-llama/llama-4-scout-17b-16e-instruct", max_tokens=2048):
-    client = Groq(api_key=os.environ.get('GROQ_API_KEY'))
-    response = client.chat.completions.create(
-        model=model_name,
-        messages=prompt,
-        max_tokens=max_tokens
-    )
-    return response.choices[0].message.content
+def process_frame(frame, seg_enabled, blur_bg, set_bg, set_color, bg_color, blur_intensity, custom_image=None):
+    global bg_image
+    settings.update({
+        "seg_enabled": seg_enabled,
+        "blur_bg": blur_bg,
+        "set_bg": set_bg,
+        "set_color": set_color,
+        "bg_color": tuple(map(int, bg_color.split(","))) if set_color else (0, 0, 0),
+        "blur_intensity": blur_intensity
+    })
 
-# Hàm trích xuất audio từ video
-def extract_audio_from_video(video_path):
-    try:
-        video = VideoFileClip(video_path)
-        audio_path = "temp_audio.mp3"
-        video.audio.write_audiofile(audio_path)
-        video.close()
-        return audio_path
-    except Exception as e:
-        raise Exception(f"Lỗi khi trích xuất audio: {str(e)}")
+    if custom_image is not None and set_bg:
+        bg_image = custom_image
 
-# Hàm trích xuất thumbnails từ video
-def extract_thumbnails(video_path, num_frames=6):
-    try:
-        video = VideoFileClip(video_path)
-        duration = video.duration
-        step = duration / num_frames
-        thumbnails = []
-        for i in range(num_frames):
-            frame_time = i * step
-            frame = video.get_frame(frame_time)
-            thumbnails.append(frame)
-        video.close()
-        return thumbnails
-    except Exception as e:
-        raise Exception(f"Lỗi khi trích xuất thumbnails: {str(e)}")
+    process_start = time.time()
+    frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
+    result = segmentation.process(frame_rgb)
+    mask = result.segmentation_mask
 
-# Hàm gọi Hugging Face Inference API để chuyển audio thành văn bản
-def transcribe_audio(audio_path):
-    try:
-        client = InferenceClient(api_key=os.environ.get('HF_TOKEN'))
-        with open(audio_path, "rb") as audio_file:
-            transcription = client.automatic_speech_recognition(
-                audio_file,
-                model="openai/whisper-tiny"
-            )
-        return transcription['text']
-    except Exception as e:
-        raise Exception(f"Lỗi khi chuyển audio thành văn bản: {str(e)}")
+    # Create alpha mask
+    alpha = mask > 0.5
+    alpha = alpha.astype(np.uint8) * 255
+    alpha = cv2.merge([alpha, alpha, alpha])
 
-# Hàm tạo prompt cho Groq
-def create_prompt(social_media_type, transcription_text):
-    system_msg = (
-        "Bạn là chuyên gia trong việc tạo nội dung mạng xã hội và tạo bài đăng hiệu quả dựa trên nội dung người dùng. "
-        "Tuân thủ quy tắc và ràng buộc của nền tảng mạng xã hội."
-    )
-    user_msg = f"Nền tảng: {social_media_type}\nNội dung văn bản: {transcription_text}"
-    return [
-        {"role": "system", "content": system_msg},
-        {"role": "user", "content": user_msg}
-    ]
+    output_frame = frame.copy()
 
-# Hàm xử lý chính
-def process_and_generate_post(video_file, social_media_type, progress=gr.Progress()):
-    progress(0, desc="Đang khởi tạo...")
-    
-    if not video_file:
-        return "Vui lòng tải lên tệp video.", None, None, None
-    
-    try:
-        # Kiểm tra kích thước tệp video
-        video_size = os.path.getsize(video_file) / (1024 * 1024)  # Chuyển sang MB
-        if video_size > MAX_VIDEO_SIZE_MB:
-            return f"Tệp video lớn hơn {MAX_VIDEO_SIZE_MB} MB. Vui lòng tải lên tệp nhỏ hơn.", None, None, None
-        
-        progress(0.2, desc="Đang trích xuất audio")
-        audio_path = extract_audio_from_video(video_file)
-        
-        progress(0.4, desc="Đang trích xuất thumbnails")
-        thumbnails = extract_thumbnails(video_file)
-        
-        progress(0.6, desc="Đang chuyển audio thành văn bản")
-        transcription_text = transcribe_audio(audio_path)
-        
-        progress(0.8, desc="Đang tạo bài đăng mạng xã hội")
-        prompt = create_prompt(social_media_type, transcription_text)
-        social_media_post = call_groq_api(prompt)
-        
-        # Xóa tệp audio tạm
-        if os.path.exists(audio_path):
-            os.remove(audio_path)
-        
-        return social_media_post, thumbnails, transcription_text, audio_path
-    
-    except Exception as e:
-        return f"Đã xảy ra lỗi: {str(e)}", None, None, None
+    if settings["seg_enabled"]:
+        if settings["blur_bg"]:
+            bg = cv2.resize(frame, None, fx=0.1, fy=0.1, interpolation=cv2.INTER_LINEAR)
+            ksize = settings["blur_intensity"]
+            if ksize % 2 == 0:
+                ksize -= 1
+            bg = cv2.GaussianBlur(bg, (ksize, ksize), 0)
+            bg = cv2.resize(bg, (frame.shape[1], frame.shape[0]), interpolation=cv2.INTER_LINEAR)
+            output_frame = np.where(alpha == 255, frame, bg)
+        elif settings["set_bg"] and bg_image is not None:
+            if bg_image.shape[:2] != frame.shape[:2]:
+                bg_image = cv2.resize(bg_image, (frame.shape[1], frame.shape[0]))
+            output_frame = np.where(alpha == 255, frame, bg_image)
+        elif settings["set_color"]:
+            bg = np.full_like(frame, settings["bg_color"])
+            output_frame = np.where(alpha == 255, frame, bg)
+        else:
+            bg = np.zeros_like(frame)
+            output_frame = np.where(alpha == 255, frame, bg)
+
+    process_time = (time.time() - process_start) * 1000
+    return output_frame, f"{process_time:.2f} ms"
+
+# Gradio interface
+with gr.Blocks() as demo:
+    gr.Markdown("# AI Background Remover")
+    with gr.Row():
+        with gr.Column():
+            webcam = gr.Webcam(label="Live Video")
+            output_image = gr.Image(label="Processed Output")
+            seg_enabled = gr.Checkbox(label="Enable Background Removal", value=True)
+            blur_bg = gr.Checkbox(label="Blur Background")
+            set_bg = gr.Checkbox(label="Custom Image Background")
+            custom_image = gr.Image(label="Upload Custom Background")
+            set_color = gr.Checkbox(label="Solid Color Background")
+            bg_color = gr.Textbox(label="Background Color (R,G,B)", value="0,0,0")
+            blur_intensity = gr.Slider(label="Blur Intensity", minimum=5, maximum=25, value=15, step=2)
+            processing_time = gr.Textbox(label="Processing Time", value="0 ms")
+
+    webcam.stream(
+        fn=process_frame,
+        inputs=[webcam, seg_enabled, blur_bg, set_bg, set_color, bg_color, blur_intensity, custom_image],
+        outputs=[output_image, processing_time]
+    )
 
-# Giao diện Gradio
-def gradio_interface():
-    with gr.Blocks(theme=gr.themes.Base()) as demo:
-        gr.Markdown("### Công cụ Tạo Bài Đăng Mạng Xã Hội")
-        video_input = gr.File(label="Tải lên Video", file_types=[".mp4", ".avi", ".mov"])
-        social_media_type = gr.Radio(
-            choices=["X (Twitter)", "Facebook", "LinkedIn", "Instagram"],
-            value="X (Twitter)",
-            label="Nền tảng"
-        )
-        generate_btn = gr.Button("Tạo Bài Đăng")
-        output = gr.Textbox(label="Bài Đăng Đã Tạo")
-        generate_btn.click(
-            fn=process_and_generate_post,
-            inputs=[video_input, social_media_type],
-            outputs=[output]
-        )
-        return demo
-# Khởi chạy giao diện Gradio
 if __name__ == "__main__":
-    demo = gradio_interface()
-    demo.launch(share=True)
+    demo.launch()