Update pixelateTG.py

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,52 +42,42 @@ 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]:
             new_width = image.shape[1] - overlay_x
         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))