diff --git a/pixelateTG.py b/pixelateTG.py index 68c67d3..825f605 100644 --- a/pixelateTG.py +++ b/pixelateTG.py @@ -4,6 +4,7 @@ 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 @@ -11,6 +12,7 @@ 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)