-
Notifications
You must be signed in to change notification settings - Fork 0
/
main.py
146 lines (117 loc) · 6.55 KB
/
main.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
from data.database_handler import DatabaseHandler
from utils.audio_utils import AudioUtils
from utils.image_utils import ImageUtils
from dotenv import load_dotenv
import os
from telegram import Update
from telegram.ext import Application, ContextTypes, MessageHandler, filters
# .env
load_dotenv()
TELEGRAM_BOT_TOKEN=os.getenv('TELEGRAM_BOT_TOKEN')
IMAGE_DATA_PATH = './data/image_data/'
# =========================================================================================== BOT FUNCTION HANDLERS
async def error_handler(update: object, context: ContextTypes.DEFAULT_TYPE) -> None:
"""Log the error and send a telegram message to notify the developer."""
# Log the error, also we can implement a dev notification.
print("An exception was raised while handling an update.")
async def audio_message(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
"""Receive an audio message, changes the sampling rate to 16kHz and
saves it in .wav format.
Args:
update (Update): This object represents an incoming update.
context (ContextTypes.DEFAULT_TYPE): Determines the type of the context argument.
"""
# Edge case of the message been edited
if not update.message:
return None
# Obtain the user ID to know where to store its audio
user_id = update.effective_user.id
# Audio preprocessing and conversion
file_id = update.message.voice.file_id
audio_msg_file = await context.bot.get_file(file_id)
# NOTE: Here I ASSUME that the audio will be short enough to be stored as bytearray in RAM,
# allowing us to directly preprocess the audio without additional I/O latency
# of saving the file, reading it, and saving it again after preprocessing.
audio_data = await audio_msg_file.download_as_bytearray()
au = AudioUtils()
await au.processAudio(audio_data=audio_data, user_id=user_id)
await update.message.reply_text("Let's check this audio!")
async def image_message(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
"""Determines whether there is a face in the photos being sent or not,
saves only those where it is.
NOTE: According to the statement: "Determines whether there is a face in
the photos being sent or not, saves only those where it is" it's not clear
if we need to record the image relations with the User into a DB, neither
if the images should have some user_id, or if we need to rescale those images.
Here I'm assuming that we are making a face Dataset, so all the images
are named as [image_<0>.jpg, ..., image_<N>.jpg] with its original sizes.
Also no information about the userid will be included (it would be the same
code like in the audio part, if needed).
Args:
update (Update): This object represents an incoming update.
context (ContextTypes.DEFAULT_TYPE): Determines the type of the context argument.
"""
# Edge case of the message been edited
if not update.message:
return
iu = ImageUtils()
# for photo_obj in update.message.photo:
photo_obj = update.message.photo[-1]
new_file = await context.bot.get_file(photo_obj.file_id)
# NOTE: Here I ASSUME that the photo will be small enough to be stored as bytearray in RAM,
# allowing us to directly if there is a face without additional I/O latency
# of saving the file, reading it, and saving it again after preprocessing.
image_data = await new_file.download_as_bytearray()
has_face = await iu.processImage(img=image_data)
# Telegram converts all the images to .jpg, which is good for making the Dataset
# and we don't need to preprocess it further.
if has_face:
files_in_dir = iu.count_directory_files()
# In case of failure, we will simply skip this image.
if files_in_dir < 0:
return None
new_img_id = f'{IMAGE_DATA_PATH}image_{files_in_dir}.jpg'
# NOTE: Saving the original file from internet without rescaling could be
# a problem if someone starts sending very big images. But with
# "saves only those where it is" I understand that is original Image.
# At least Telegram downsizes big images, it's a plus in this task.
await new_file.download_to_drive(new_img_id)
if has_face:
await update.message.reply_text("What a beautiful face!")
else:
await update.message.reply_text("Wow, thanks!")
async def audioDBCount(update: Update, context: ContextTypes.DEFAULT_TYPE) -> None:
"""Development purposses only. Check the DB Audio Count for your ID.
Args:
update (Update): This object represents an incoming update.
context (ContextTypes.DEFAULT_TYPE): Determines the type of the context argument.
"""
db = DatabaseHandler()
res = f'We have {db.getUserAudioCount(update.effective_user.id)} audios from you'
await update.message.reply_text(res)
# =========================================================================================== MAIN APP
def main():
# --------------------------------------------------------------------------------------- Database initialization
db = DatabaseHandler()
try:
# Create tables: Users, Audios, <Images - in case of being needed>
db.create_tables()
except Exception as db_error:
print(f"Error initializing database: {db_error}")
return None
# ====================================================================================== Start the SrVladyslav Bot
try:
print("Starting the SrVladyslav Bot")
telegram_bot = Application.builder().token(TELEGRAM_BOT_TOKEN).build()
# ---------------------------------------------------------------------------------- Adding the bot handlers
telegram_bot.add_handler(MessageHandler(filters.VOICE, audio_message)) # Audio filtering task
telegram_bot.add_handler(MessageHandler(filters.PHOTO, image_message)) # Image filtering task
telegram_bot.add_error_handler(error_handler)
# telegram_bot.add_handler(CommandHandler("adb", audioDBCount)) # /adb returns the Audio Db count for your ID, (Dev. HELPER)
# ---------------------------------------------------------------------------------- Run the bot until Ctrl-C is pressed
print("The bot is running... Press Ctrl-C to stop.")
telegram_bot.run_polling(allowed_updates=Update.ALL_TYPES)
except Exception as e:
print(f"An error occurred while initializing the bot: {e}")
if __name__ == '__main__':
main()