Update pixelateTG.py

pixelateTG.py

@@ -8,6 +8,10 @@
 from concurrent.futures import ThreadPoolExecutor, wait
 from mtcnn.mtcnn import MTCNN
 from uuid import uuid4
+import logging
+
+logging.basicConfig(level=logging.INFO)
+logger = logging.getLogger(__name__)
 
 # Load environment variables from .env file
 load_dotenv()
@@ -94,15 +98,17 @@ 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)
-    print(f"Number of frames in GIF: {len(frames)}")
-    processed_frames = [process_image(frame, user_id, session_id, bot) for frame in frames]
-    processed_gif_path = f"processed/{user_id}_{session_id}.gif"
-    imageio.mimsave(processed_gif_path, processed_frames)
-    print(f"Processed GIF saved at: {processed_gif_path}")
-    return processed_gif_path
-
-
+    try:
+        frames = imageio.mimread(gif_path)
+        logger.info(f"Number of frames in GIF: {len(frames)}")
+        processed_frames = [process_image(frame, user_id, session_id, bot) for frame in frames]
+        processed_gif_path = f"processed/{user_id}_{session_id}.gif"
+        imageio.mimsave(processed_gif_path, processed_frames)
+        logger.info(f"Processed GIF saved at: {processed_gif_path}")
+        return processed_gif_path
+    except Exception as e:
+        logger.error(f"Error processing GIF: {e}")
+        raise
 
 def pixelate_faces(update: Update, context: CallbackContext) -> None:
     session_id = str(uuid4())
@@ -116,6 +122,8 @@ def pixelate_faces(update: Update, context: CallbackContext) -> None:
         # Process the image
         photo_path = f"downloads/{file_name}"
         file.download(photo_path)
+
+        logger.info(f"Photo downloaded to {photo_path}")
 
         image = cv2.imread(photo_path)
         faces = detect_heads(image)
@@ -147,33 +155,43 @@ def pixelate_faces(update: Update, context: CallbackContext) -> None:
         # Download the GIF file
         gif_path = f"downloads/{file_name}"
         file.download(gif_path)
+
+        logger.info(f"GIF downloaded to {gif_path}")
 
         # Process the GIF
-        processed_gif_path = process_gif(gif_path, session_id, update.message.from_user.id, context.bot)
-        context.bot.send_animation(chat_id=update.message.from_user.id, animation=open(processed_gif_path, 'rb'))
+        try:
+            processed_gif_path = process_gif(gif_path, session_id, update.message.from_user.id, context.bot)
+            context.bot.send_animation(chat_id=update.message.from_user.id, animation=open(processed_gif_path, 'rb'))
+        except Exception as e:
+            logger.error(f"Error processing GIF: {e}")
 
         # Clean up temporary files
         os.remove(gif_path)
     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:
+    if update.message.reply_to_message and (update.message.reply_to_message.photo or update.message.reply_to_message.document.mime_type == 'image/gif'):
         session_id = str(uuid4())
         chat_data = context.chat_data
 
-        file_id = update.message.reply_to_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)
-
-        image = cv2.imread(photo_path)
-        faces = detect_heads(image)
-
-        if not faces:
+        if update.message.reply_to_message.photo:
+            file_id = update.message.reply_to_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)
+            image = cv2.imread(photo_path)
+            faces = detect_heads(image)
+        else: # GIF handling
+            file_id = update.message.reply_to_message.document.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)
+            # For GIFs, no need to detect faces at this step
+
+        if not faces and not update.message.reply_to_message.document.mime_type == 'image/gif':
             update.message.reply_text('No faces detected in the image.')
             return
 
@@ -192,7 +210,7 @@ def pixelate_command(update: Update, context: CallbackContext) -> None:
 
         update.message.reply_text('Press buttons until happy', reply_markup=reply_markup)
     else:
-        update.message.reply_text('This only works as a reply to a picture.')
+        update.message.reply_text('This only works as a reply to a picture or a GIF.')
 
 def process_image(photo_path, user_id, session_id, bot):
     image = cv2.imread(photo_path)
@@ -209,7 +227,6 @@ def process_image(photo_path, user_id, session_id, bot):
     cv2.imwrite(processed_path, image, [int(cv2.IMWRITE_JPEG_QUALITY), 95])
     return processed_path
 
-
 def button_callback(update: Update, context: CallbackContext) -> None:
     query = update.callback_query
     query.answer()