From 06ea936242f6e0d1b8c5ff64cea98c0c5c8f8711 Mon Sep 17 00:00:00 2001 From: arbadacarba <63317640+arbadacarbaYK@users.noreply.github.com> Date: Thu, 18 Jul 2024 16:40:00 +0200 Subject: [PATCH] Update pixelateTG.py --- pixelateTG.py | 57 +++++++++++++++------------------------------------ 1 file changed, 16 insertions(+), 41 deletions(-) diff --git a/pixelateTG.py b/pixelateTG.py index ee7bbab..a3b506e 100644 --- a/pixelateTG.py +++ b/pixelateTG.py @@ -1,6 +1,6 @@ import os os.environ['CUDA_VISIBLE_DEVICES'] = '-1' -from dotenv import load_dotenv # Import the load_dotenv function from python-dotenv +from dotenv import load_dotenv import cv2 import random import imageio @@ -13,23 +13,21 @@ # Load environment variables from .env file load_dotenv() -TOKEN = os.getenv('TELEGRAM_BOT_TOKEN') # Get the Telegram bot token from the environment variable +TOKEN = os.getenv('TELEGRAM_BOT_TOKEN') MAX_THREADS = 15 PIXELATION_FACTOR = 0.04 -RESIZE_FACTOR = 1.5 # Common resize factor +RESIZE_FACTOR = 1.5 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!') + update.message.reply_text('Pixelbot helps you to anonymize pictures before share them with others. Start by sending me an image.') def detect_heads(image): detector = MTCNN() faces = detector.detect_faces(image) - # Sort faces by size (area) in descending order - faces.sort(key=lambda x: x['box'][2] * x['box'][3], reverse=True) + # Sort faces by size (area) in ascending order + faces.sort(key=lambda x: x['box'][2] * x['box'][3]) head_boxes = [(face['box'][0], face['box'][1], int(RESIZE_FACTOR * face['box'][2]), int(RESIZE_FACTOR * face['box'][3])) for face in faces] - # Sort faces based on y-coordinate (top to bottom) - head_boxes.sort(key=lambda box: box[1]) return head_boxes def overlay(photo_path, user_id, overlay_type, resize_factor, bot): @@ -44,19 +42,15 @@ def overlay(photo_path, user_id, overlay_type, resize_factor, bot): overlay_image = cv2.imread(random_overlay, cv2.IMREAD_UNCHANGED) original_aspect_ratio = overlay_image.shape[1] / overlay_image.shape[0] - # Calculate new dimensions for the overlay new_width = int(resize_factor * w) new_height = int(new_width / original_aspect_ratio) - # Ensure the overlay is centered on the face center_x = x + w // 2 center_y = y + h // 2 - # Overlay position adjusted for better centering overlay_x = center_x - (new_width // 2) overlay_y = center_y - (new_height // 2) - # Clamp values to ensure they are within the image boundaries overlay_x = max(0, overlay_x) overlay_y = max(0, overlay_y) if overlay_x + new_width > image.shape[1]: @@ -64,32 +58,26 @@ def overlay(photo_path, user_id, overlay_type, resize_factor, bot): if overlay_y + new_height > image.shape[0]: new_height = image.shape[0] - overlay_y - # Resize the overlay image overlay_image_resized = cv2.resize(overlay_image, (new_width, new_height), interpolation=cv2.INTER_AREA) - # Calculate the regions of interest (ROI) roi_start_x = overlay_x roi_start_y = overlay_y roi_end_x = min(image.shape[1], overlay_x + new_width) roi_end_y = min(image.shape[0], overlay_y + new_height) - # Blend the overlay onto the image - try: - overlay_part = overlay_image_resized[:roi_end_y - roi_start_y, :roi_end_x - roi_start_x] - alpha_mask = overlay_part[:, :, 3] / 255.0 - for c in range(3): - image[roi_start_y:roi_end_y, roi_start_x:roi_end_x, c] = ( - alpha_mask * overlay_part[:, :, c] + - (1 - alpha_mask) * image[roi_start_y:roi_end_y, roi_start_x:roi_end_x, c] - ) - except ValueError as e: - print(f"Error blending overlay: {e}") - continue + overlay_part = overlay_image_resized[:roi_end_y - roi_start_y, :roi_end_x - roi_start_x] + alpha_mask = overlay_part[:, :, 3] / 255.0 + for c in range(3): + image[roi_start_y:roi_end_y, roi_start_x:roi_end_x, c] = ( + alpha_mask * overlay_part[:, :, c] + + (1 - alpha_mask) * image[roi_start_y:roi_end_y, roi_start_x:roi_end_x, c] + ) processed_path = f"processed/{user_id}_{overlay_type}.jpg" cv2.imwrite(processed_path, image, [int(cv2.IMWRITE_JPEG_QUALITY), 95]) return processed_path + # Overlay functions def liotta_overlay(photo_path, user_id, bot): return overlay(photo_path, user_id, 'liotta', RESIZE_FACTOR, bot) @@ -143,7 +131,6 @@ def pixelate_faces(update: Update, context: CallbackContext) -> None: 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}')]) @@ -168,7 +155,6 @@ def pixelate_faces(update: Update, context: CallbackContext) -> None: 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()) @@ -197,6 +183,7 @@ def pixelate_command(update: Update, context: CallbackContext) -> None: [InlineKeyboardButton("⚔️ Pixel", callback_data=f'pixelate_{session_id}'), InlineKeyboardButton("CLOSE ME", callback_data=f'cancel_{session_id}')] ] + reply_markup = InlineKeyboardMarkup(keyboard) chat_data[session_id] = {'photo_path': photo_path, 'chat_id': update.message.chat.id} @@ -209,15 +196,10 @@ def process_image(photo_path, user_id, session_id, bot): 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] - - # Apply pixelation to the ROI - pixelation_size = max(1, int(PIXELATION_FACTOR * min(w, h))) # Ensure pixelation size is at least 1 + pixelation_size = max(1, int(PIXELATION_FACTOR * min(w, h))) 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) - - # Replace the original face region with the pixelated ROI image[y:y+h, x:x+w] = pixelated_roi processed_path = f"processed/{user_id}_{session_id}_pixelated.jpg" @@ -236,17 +218,12 @@ def button_callback(update: Update, context: CallbackContext) -> None: photo_path = data.get('photo_path') user_or_chat_id = data.get('user_id') or data.get('chat_id') - # Delete files in both 'downloads' and 'processed' directories if query.data.startswith('cancel'): - # Delete the corresponding processed file processed_file_path = f"processed/{user_or_chat_id}_{session_id}_pixelated.jpg" if os.path.exists(processed_file_path): os.remove(processed_file_path) - - # Delete the original photo/gif file from 'downloads' directory if os.path.exists(photo_path): os.remove(photo_path) - if session_id in user_data: del user_data[session_id] if session_id in chat_data: @@ -274,10 +251,8 @@ 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) - dispatcher = updater.dispatcher dispatcher.add_handler(CommandHandler("start", start))