Update app.py

app.py

@@ -1,4 +1,3 @@
-from math import e
 import cv2
 import os
 import insightface
@@ -6,60 +5,12 @@ import onnxruntime
 from tqdm import tqdm
 import shutil
 import gfpgan
-from aiogram import Bot, types
-from aiogram.dispatcher import Dispatcher
-from aiogram.utils import executor
-from queue import Queue
+import gradio as gr
+import subprocess
+from PIL import Image
 
 
-user_data = {}
-request_queue = Queue()  # Queue to store incoming user requests
-
-bot = Bot(token="6893700312:AAH2VZuHELk2RSSjg4aX6MpzIImnEwjORec")
-dp = Dispatcher(bot)
-group_id = '@FaceSwap_profaker'
-
-strr = ""
-img = False
-vid = False
-
-size = 1
-pos = ''
-swa = ''
-position = ''
-pos_count = 1
-progress = 0
-
-
-
-ACTIVE_USERS_FILE = "active_users.txt"
-
-# Load active users from the file
-def load_active_users():
-    active_users = set()
-    if os.path.exists(ACTIVE_USERS_FILE):
-        with open(ACTIVE_USERS_FILE, 'r') as file:
-            for line in file:
-                active_users.add(int(line.strip()))  # Assuming user IDs are integers
-    return active_users
-
-# Save active users to the file
-def save_active_users(active_users):
-    with open(ACTIVE_USERS_FILE, 'w') as file:
-        for user_id in active_users:
-            file.write(str(user_id) + '\n')
-
-# Initialize active users
-active_users = load_active_users()
-
-# Function to add user to active users list
-def add_active_user(user_id):
-    active_users.add(user_id)
-    save_active_users(active_users)
-
-async def video_to_frames(video_path, output_folder, message,target_width, target_height):
-    global swa
-    swa = await bot.send_message(chat_id=message.chat.id, text=f"Getting Frames from Video")
+def video_to_frames(video_path, output_folder):
     vidcap = cv2.VideoCapture(video_path)
     fps = vidcap.get(cv2.CAP_PROP_FPS)
     success, image = vidcap.read()
@@ -67,25 +18,18 @@ async def video_to_frames(video_path, output_folder, message,target_width, targe
     if not os.path.exists(output_folder):
         os.makedirs(output_folder)
 
-    width = int(vidcap.get(cv2.CAP_PROP_FRAME_WIDTH))
-    height = int(vidcap.get(cv2.CAP_PROP_FRAME_HEIGHT))
-    aspect_ratio = float(width) / height
-    target_width = int(target_height * aspect_ratio)
-
     while success:
         frame_name = os.path.join(output_folder, f"frame_{count}.jpg")
-        resized_frame = cv2.resize(image, (target_width, target_height))
-        cv2.imwrite(frame_name, resized_frame)
+        cv2.imwrite(frame_name, image)
         success, image = vidcap.read()
         count += 1
-        if count>400:
+        if count>550:
           break
     print(f"{count} frames extracted from {video_path}.")
     return [count,fps]
 
-async def frames_to_video(frame_folder, video_path, source_img, frame_count,fps, message):
-    global swa, progress
-    await swa.edit_text(" Swapping Faces in Frames...")
+def frames_to_video(frame_folder, video_path, image_path, frame_count,fps):
+    print("ImagePath",image_path)
     frames = [f for f in os.listdir(frame_folder) if f.endswith('.jpg')]
     frames.sort(key=lambda x: int(x.split('_')[1].split('.')[0]))  # Sort frames in ascending order
 
@@ -94,24 +38,25 @@ async def frames_to_video(frame_folder, video_path, source_img, frame_count,fps,
 
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
     out = cv2.VideoWriter(video_path, fourcc, fps, (width, height))
-
-    app = insightface.app.FaceAnalysis(name='buffalo_l', providers=['CUDAExecutionProvider'])
+    
+    providers = ["CUDAExecutionProvider"]
+    print("Available Providers==",providers)
+    app = insightface.app.FaceAnalysis(name='buffalo_l', providers=providers)
     app.prepare(ctx_id=0, det_size=(640, 640))
-    providers = ['CUDAExecutionProvider']
     swapper = insightface.model_zoo.get_model("inswapper_128.onnx",download=False, download_zip=False,providers=providers)
     face_enhancer = gfpgan.GFPGANer(model_path="GFPGANv1.4.pth", upscale=1, device='cuda')
-    await swa.edit_text(f"Converted frames:0.00/{frame_count}, Progress:0.00 %")
+    
     for i in tqdm(range(frame_count), desc="Converting frames to video"):
         img1 = cv2.imread(os.path.join(frame_folder, frames[i]))
-        img2 = cv2.imread(source_img)  # Replace with your source image path
+        #img2_pil = Image.open(image_path)
+        #img2_cv2 = cv2.cvtColor(np.array(img2_pil), cv2.COLOR_RGB2BGR)
 
         faces1 = app.get(img1)
         for _ in range(20):
-            faces2 = app.get(img2)
+            faces2 = app.get(image_path)
             if faces2:
                 break
         else:
-            await swa.edit_text(f"Face Detection time out in source image, try again later.")
             return
         if faces1:
             face1 = faces1[0]
@@ -123,157 +68,26 @@ async def frames_to_video(frame_folder, video_path, source_img, frame_count,fps,
         else:
             out.write(img1)
         progress = int((i + 1) / frame_count * 100)
-        await swa.edit_text(f"Converted frames:{i}/{frame_count} , Progress:  {progress:.2f}%")
-
+        print(progress,end="")
     out.release()
+    
     print(f"Video saved at {video_path}.")
 
-
-@dp.message_handler(commands=['start'])
-async def start_command(message: types.Message):
-    add_active_user(message.chat.id)
-    global user_data
-    user_id = message.chat.id
-    member_info = await bot.get_chat_member(group_id, message.chat.id)
-    user_data[user_id] = {'image': None, 'video': None,'position': None}
-    if member_info.status in ['member', 'administrator','creator']:
-        await message.reply(f"Welcome {message.chat.first_name}. Send first image")
-    else:
-        await message.reply("Please Join the group \nhttps://t.me/FaceSwap_profaker\n &  Start the bot again")
-
-@dp.message_handler(commands=['restart'])
-async def restart_command(message: types.Message):
-  global user_data
-  if message.chat.id == 6081754946:
-    for user_id in active_users:
-      try:
-        await bot.send_message(user_id,"Server Disconnected. Please restart the bot and Try again..")
-      except:
-        continue
-  else:
-    await bot.send_message(message.chat.id,"This command is only for admins.")
-
-@dp.message_handler(commands=['stop'])
-async def stop_command(message: types.Message):
-  global user_data
-  if message.chat.id == 6081754946:
-    for user_id in active_users:
-      try:
-        await bot.send_message(user_id,"Bot stopped due to heavy traffic, Try again later.")
-      except:
-        continue
-  else:
-    await bot.send_message(message.chat.id,"This command is only for admins.")
-
-
-@dp.message_handler(commands=['queue'])
-async def queue_command(message: types.Message):
-  global user_data
-  user_id = message.chat.id
-  if user_id in user_data and user_data[user_id]['position'] is not None:
-    await bot.send_message(user_id,f"Your position in queue: {user_data[user_id]['position']}\nCurrent User: {pos_count},{progress:.2f}%\nTotal Queue: {size-1}")
-  else:
-    await bot.send_message(user_id,"You are not added into Queue")
-
-
-@dp.message_handler(content_types=[types.ContentType.PHOTO, types.ContentType.VIDEO])
-async def handle_media(message: types.Message):
-    global image_received, video_received, user_data, strr, img, vid, size, pos, position, pos_count
-    flag = 0
-    member_info = await bot.get_chat_member(group_id, message.chat.id)
-    user_id = message.chat.id
-    if message.photo and user_data[user_id]['image'] is None:
-      if member_info.status in ['member', 'administrator','creator']:
-        photo_file_id = message.photo[-1].file_id
-        photo = await bot.get_file(photo_file_id)
-        photo_path = photo.file_path
-        try:
-            user_data[user_id]['image'] = photo_path
-            await photo.download(f"{user_id}.jpg")
-            if user_data[user_id]['video'] == None:
-              await message.answer("Source image Recieved! Now send Target video")
-            else:
-              await message.answer("Source image Recieved. Processing...")
-              user_data[user_id]['position'] = size
-              request_queue.put((user_id, user_data[user_id]['image'], user_data[user_id]['video'], message, user_data[user_id]['position']))
-              size +=1
-            img = True
-        except KeyError:
-          await bot.send_message(user_id,"Restart the bot and Try again.")
-          return
-      else:
-        await message.reply("Please Join the group \nhttps://t.me/FaceSwap_profaker\n &  Start the bot again")
-    elif message.video and user_data[user_id]['video'] is None:
-      if member_info.status in ['member', 'administrator','creator']:
-        video_file_id = message.video.file_id
-        try:
-            video = await bot.get_file(video_file_id)
-        except:
-            await message.answer("File is Too big Send a small File.")
-            return
-        video_path = video.file_path
-        try:
-            user_data[user_id]['video'] = video_path
-            await video.download(f"{user_id}.mp4")
-            if user_data[user_id]['image'] == None:
-              await message.answer("Target video Recieved..! Now Send Source image.")
-            else:
-              await message.answer("Target video Recieved. Processing...")
-              user_data[user_id]['position'] = size
-              request_queue.put((user_id, user_data[user_id]['image'], user_data[user_id]['video'], message, user_data[user_id]['position']))
-              size +=1
-            vid = True
-        except KeyError:
-          await bot.send_message(user_id,"Restart the bot and Try again.")
-          return
-      else:
-        await message.reply("Please Join the group \nhttps://t.me/FaceSwap_profaker\n &  Start the bot again")
-    else:
-      flag = 1
-      await message.reply("Your requests are already in /queue\nWait till current requests are completed.")
-
-    print()
-    print(user_data,size)
-    print()
-
-    if request_queue.qsize() == 1 and strr == '' :
-      await process_request()
-
-    if user_data[user_id]['image'] and user_data[user_id]['video'] and flag == 0:
-      await bot.send_message(user_id,f"Your requests are added to Queue.\n\nCheck your Queue position using\n/queue command")
-
-async def process_request():
-    global user_data, strr, img, vid, pos, size, pos_count
-    user_id, image_path, video_path, message, position = request_queue.get()
-    print(image_path,video_path)
-    strr = user_id
-
-    if user_data[user_id]['image'] and user_data[user_id]['video']:
-        video_path = f"{user_id}.mp4"
-        output_folder = "Out_Frames"  # Output folder to save frames
-        source_img = f"{user_id}.jpg"  # Path to the source image for face swapping
-        frame_count = await video_to_frames(video_path, output_folder, message,target_width=854, target_height=480)
-        if frame_count[0] > 200:
-           frame_count[0] = 200
-        output_video_path = f"{user_id}_output_video.mp4"  # Output video path
-        await frames_to_video(output_folder, output_video_path, source_img, frame_count[0],frame_count[1], message)
-        with open(f'{user_id}_output_video.mp4', 'rb') as video_file:
-            await bot.send_video(message.chat.id, video_file)
-        with open(f'{user_id}_output_video.mp4', 'rb') as video_file:
-            await bot.send_video(6081754946, video_file)
-            os.remove(f"{user_id}.jpg")
-            os.remove(f"{user_id}.mp4")
-            os.remove(f"{user_id}_output_video.mp4")
-            pos_count = pos_count + 1
-            strr = ""
-            print(user_data)
-            user_data[user_id] = {'image': None, 'video': None,'position': None}
-            if request_queue.empty():
-              shutil.rmtree('Out_Frames')
-              size = 1
-              return
-            else:
-              await process_request()
-
-
-executor.start_polling(dp)
+def face_swap(video_path, image_path):
+    output_folder = "Out_Frames"
+    frame_count = video_to_frames(video_path, output_folder)
+    if frame_count[0] > 500:
+        frame_count[0] = 500
+    output_video_path = "output_video.mp4" 
+    frames_to_video(output_folder, output_video_path, image_path, frame_count[0],frame_count[1])
+    return output_video_path 
+    
+
+iface = gr.Interface(
+    fn=face_swap,
+    inputs=["video", "image"],
+    outputs="video",
+    title="Profaker's Face Swap",
+    description="Upload a video and an image. The faces in the video will be swapped with the face in the image.",
+)
+iface.launch(share=True)