Update pixelateTG.py

pixelateTG.py

@@ -4,13 +4,15 @@
 import imageio
 from telegram import Update, InlineKeyboardButton, InlineKeyboardMarkup
 from telegram.ext import Updater, CallbackContext, CommandHandler, CallbackQueryHandler, MessageHandler, Filters
+from concurrent.futures import ThreadPoolExecutor, wait
 from mtcnn.mtcnn import MTCNN
 from uuid import uuid4
 
 TOKEN = os.environ['TELEGRAM_BOT_TOKEN']
 MAX_THREADS = 15
 PIXELATION_FACTOR = 0.04
 RESIZE_FACTOR = 1.5  # Common resize factor
+executor = ThreadPoolExecutor(max_workers=MAX_THREADS)
 
 def start(update: Update, context: CallbackContext) -> None:
     update.message.reply_text('Send me a picture or a GIF, and I will pixelate faces in it!')
@@ -21,7 +23,7 @@ def detect_heads(image):
     head_boxes = [(face['box'][0], face['box'][1], int(RESIZE_FACTOR * face['box'][2]), int(RESIZE_FACTOR * face['box'][3])) for face in faces]
     return head_boxes
 
-def overlay(photo_path, user_id, overlay_type, resize_factor):
+def overlay(photo_path, user_id, overlay_type, resize_factor, bot):
     image = cv2.imread(photo_path)
     heads = detect_heads(image)
 
@@ -77,23 +79,23 @@ def overlay(photo_path, user_id, overlay_type, resize_factor):
 
 
 # Overlay functions
-def liotta_overlay(photo_path, user_id):
-    return overlay(photo_path, user_id, 'liotta', RESIZE_FACTOR)
+def liotta_overlay(photo_path, user_id, bot):
+    return overlay(photo_path, user_id, 'liotta', RESIZE_FACTOR, bot)
 
-def skull_overlay(photo_path, user_id):
-    return overlay(photo_path, user_id, 'skullofsatoshi', RESIZE_FACTOR)
+def skull_overlay(photo_path, user_id, bot):
+    return overlay(photo_path, user_id, 'skullofsatoshi', RESIZE_FACTOR, bot)
 
-def pepe_overlay(photo_path, user_id):
-    return overlay(photo_path, user_id, 'pepe', RESIZE_FACTOR)
+def pepe_overlay(photo_path, user_id, bot):
+    return overlay(photo_path, user_id, 'pepe', RESIZE_FACTOR, bot)
 
-def chad_overlay(photo_path, user_id):
-    return overlay(photo_path, user_id, 'chad', RESIZE_FACTOR)
+def chad_overlay(photo_path, user_id, bot):
+    return overlay(photo_path, user_id, 'chad', RESIZE_FACTOR, bot)
 
-def cats_overlay(photo_path, user_id):
-    return overlay(photo_path, user_id, 'cat', RESIZE_FACTOR)
+def cats_overlay(photo_path, user_id, bot):
+    return overlay(photo_path, user_id, 'cat', RESIZE_FACTOR, bot)
 
-def clowns_overlay(photo_path, user_id):
-    return overlay(photo_path, user_id, 'clown', RESIZE_FACTOR)
+def clowns_overlay(photo_path, user_id, bot):
+    return overlay(photo_path, user_id, 'clown', RESIZE_FACTOR, bot)
 
 def process_gif(gif_path, session_id, user_id, bot):
     frames = imageio.mimread(gif_path)
@@ -102,43 +104,59 @@ def process_gif(gif_path, session_id, user_id, bot):
     imageio.mimsave(processed_gif_path, processed_frames)
     return processed_gif_path
 
-def process_image(image, user_id, session_id):
-    faces = detect_heads(image)
-
-    for (x, y, w, h) in faces:
-        # Define the region of interest (ROI)
-        roi = image[y:y+h, x:x+w]
+def pixelate_faces(update: Update, context: CallbackContext) -> None:
+    session_id = str(uuid4())
+    user_data = context.user_data
 
-        # Apply pixelation to the ROI
-        pixelation_size = max(1, int(PIXELATION_FACTOR * min(w, h)))  # Ensure pixelation size is at least 1
-        pixelated_roi = cv2.resize(roi, (pixelation_size, pixelation_size), interpolation=cv2.INTER_NEAREST)
-        pixelated_roi = cv2.resize(pixelated_roi, (w, h), interpolation=cv2.INTER_NEAREST)
+    if update.message.photo:
+        file_id = update.message.photo[-1].file_id
+        file = context.bot.get_file(file_id)
+        file_name = file.file_path.split('/')[-1]
+        photo_path = f"downloads/{file_name}"
+        file.download(photo_path)
 
-        # Replace the original face region with the pixelated ROI
-        image[y:y+h, x:x+w] = pixelated_roi
+        image = cv2.imread(photo_path)
+        faces = detect_heads(image)
 
-    processed_path = f"processed/{user_id}_{session_id}_pixelated.jpg"
-    cv2.imwrite(processed_path, image, [int(cv2.IMWRITE_JPEG_QUALITY), 95])
-    return processed_path
+        if not faces:
+            update.message.reply_text('No faces detected in the image.')
+            return
 
+        keyboard = [
+            [InlineKeyboardButton("🤡 Clowns", callback_data=f'clowns_overlay_{session_id}'),
+             InlineKeyboardButton("😂 Liotta", callback_data=f'liotta_overlay_{session_id}'),
+             InlineKeyboardButton("☠️ Skull", callback_data=f'skull_overlay_{session_id}')],
+            [InlineKeyboardButton("🐈‍⬛ Cats", callback_data=f'cats_overlay_{session_id}'),
+             InlineKeyboardButton("🐸 Pepe", callback_data=f'pepe_overlay_{session_id}'),
+             InlineKeyboardButton("🏆 Chad", callback_data=f'chad_overlay_{session_id}')]
+        ]
+
+        # Check if it's a private chat, if yes, include the "⚔️ Pixel" button
+        if update.message.chat.type == 'private':
+            keyboard.append([InlineKeyboardButton("⚔️ Pixel", callback_data=f'pixelate_{session_id}')])
 
-def pixelate_faces(update: Update, context: CallbackContext) -> None:
-    session_id = str(uuid4())
-    user_data = context.user_data
+        keyboard.append([InlineKeyboardButton("CLOSE ME", callback_data=f'cancel_{session_id}')])
+
+        reply_markup = InlineKeyboardMarkup(keyboard)
+        user_data[session_id] = {'photo_path': photo_path, 'user_id': update.message.from_user.id}
 
-    if update.message.photo:
-        # Your existing code for handling photos...
+        update.message.reply_text('Press buttons until happy', reply_markup=reply_markup)
+        update.message.delete()
 
     elif update.message.document and update.message.document.mime_type == 'image/gif':
-        # Your existing code for handling GIFs...
-        processed_gif_path = process_gif(gif_path, session_id, update.message.from_user.id, context.bot)
+        file_id = update.message.document.file_id
+        file = context.bot.get_file(file_id)
+        file_name = file.file_path.split('/')[-1]
+        gif_path = f"downloads/{file_name}"
+        file.download(gif_path)
+
+        processed_gif_path = process_gif(gif_path, session_id, str(uuid4()), context.bot)
         context.bot.send_animation(chat_id=update.message.chat_id, animation=open(processed_gif_path, 'rb'))
 
     else:
         update.message.reply_text('Please send either a photo or a GIF.')
 
 
-
 def pixelate_command(update: Update, context: CallbackContext) -> None:
     if update.message.reply_to_message and update.message.reply_to_message.photo:
         session_id = str(uuid4())
@@ -174,12 +192,8 @@ def pixelate_command(update: Update, context: CallbackContext) -> None:
     else:
         update.message.reply_text('This only works as a reply to a picture.')
 
-def process_image(photo, user_id, session_id, bot):
-    if isinstance(photo, str):  # If photo_path is a file path
-        image = cv2.imread(photo)
-    else:  # If photo_path is an image array (frame of the GIF)
-        image = photo
-
+def process_image(photo_path, user_id, session_id, bot):
+    image = cv2.imread(photo_path)
     faces = detect_heads(image)
 
     for (x, y, w, h) in faces:
@@ -199,7 +213,6 @@ def process_image(photo, user_id, session_id, bot):
     return processed_path
 
 
-
 def button_callback(update: Update, context: CallbackContext) -> None:
     query = update.callback_query
     query.answer()
@@ -223,7 +236,7 @@ def button_callback(update: Update, context: CallbackContext) -> None:
         processed_path = None
 
         if query.data.startswith('pixelate'):
-            processed_path = process_image(photo_path, user_or_chat_id, session_id, context.bot)
+            processed_path = process_image(photo_path, user_or_chat_id, query.id, context.bot)
         elif query.data.startswith('liotta'):
             processed_path = liotta_overlay(photo_path, user_or_chat_id, context.bot)
         elif query.data.startswith('cats_overlay'):
@@ -240,7 +253,6 @@ def button_callback(update: Update, context: CallbackContext) -> None:
         if processed_path:
             context.bot.send_photo(chat_id=query.message.chat_id, photo=open(processed_path, 'rb'))
 
-
 def main() -> None:
     updater = Updater(TOKEN)