114 lines
4.2 KiB
Python
114 lines
4.2 KiB
Python
import io
|
|
import os
|
|
import tempfile
|
|
import time
|
|
from bson import ObjectId
|
|
from minio import Minio
|
|
from pymongo import MongoClient
|
|
import whisper
|
|
|
|
class MinioHandler:
|
|
def __init__(self, endpoint, access_key, secret_key):
|
|
self.client = Minio(
|
|
endpoint,
|
|
access_key=access_key,
|
|
secret_key=secret_key,
|
|
secure=True
|
|
)
|
|
|
|
def get_video_data(self, bucket, object_name):
|
|
response = self.client.get_object(bucket, object_name)
|
|
data = response.read()
|
|
print(f"Read {len(data)} bytes from Minio for {object_name}")
|
|
return data
|
|
|
|
class DatabaseHandler:
|
|
def __init__(self, mongo_uri, database_name, collection_name):
|
|
self.client = MongoClient(mongo_uri)
|
|
self.db = self.client[database_name]
|
|
self.collection = self.db[collection_name]
|
|
|
|
def get_unprocessed_videos(self):
|
|
return self.collection.find({
|
|
'bucket_name': 'raw',
|
|
'object_name': {'$regex': 'douyin/'},
|
|
'whisper_transcription': {'$exists': False}
|
|
})
|
|
|
|
def update_transcription(self, video_id, transcription):
|
|
self.collection.update_one(
|
|
{'_id': video_id},
|
|
{'$set': {'whisper_transcription': transcription}}
|
|
)
|
|
|
|
class WhisperProcessor:
|
|
def __init__(self, model_name="large-v3"):
|
|
self.model = whisper.load_model(model_name)
|
|
|
|
def transcribe_audio(self, video_data):
|
|
with tempfile.NamedTemporaryFile(delete=False, suffix=".mp4") as temp_video:
|
|
temp_video.write(video_data)
|
|
temp_video_path = temp_video.name
|
|
|
|
try:
|
|
result = self.model.transcribe(temp_video_path)
|
|
return result["text"]
|
|
finally:
|
|
os.unlink(temp_video_path)
|
|
|
|
class WhisperTranscriptionSystem:
|
|
def __init__(self, minio_endpoint, minio_access_key, minio_secret_key,
|
|
mongo_uri, db_name, collection_name):
|
|
self.minio_handler = MinioHandler(minio_endpoint, minio_access_key, minio_secret_key)
|
|
self.db_handler = DatabaseHandler(mongo_uri, db_name, collection_name)
|
|
self.whisper_processor = WhisperProcessor()
|
|
|
|
def process_video(self, video_doc):
|
|
start_time = time.time()
|
|
try:
|
|
video_data = self.minio_handler.get_video_data(video_doc['bucket_name'], video_doc['object_name'])
|
|
transcription = self.whisper_processor.transcribe_audio(video_data)
|
|
|
|
self.db_handler.update_transcription(video_doc['_id'], transcription)
|
|
|
|
end_time = time.time()
|
|
processing_time = end_time - start_time
|
|
|
|
print(f"Processed video: {video_doc['object_name']}")
|
|
print(f"Processing time: {processing_time:.2f} seconds")
|
|
except Exception as e:
|
|
end_time = time.time()
|
|
processing_time = end_time - start_time
|
|
|
|
print(f"Error processing video {video_doc['object_name']}: {str(e)}")
|
|
print(f"Processing time (including error): {processing_time:.2f} seconds")
|
|
import traceback
|
|
traceback.print_exc()
|
|
|
|
def run(self):
|
|
while True:
|
|
unprocessed_videos = list(self.db_handler.get_unprocessed_videos())
|
|
|
|
if not unprocessed_videos:
|
|
print("No new videos to process. Waiting for 60 seconds before checking again...")
|
|
time.sleep(60)
|
|
continue
|
|
|
|
print(f"Found {len(unprocessed_videos)} videos to process.")
|
|
for video_doc in unprocessed_videos:
|
|
self.process_video(video_doc)
|
|
|
|
print("Finished processing current batch of videos. Checking for more...")
|
|
|
|
if __name__ == "__main__":
|
|
minio_endpoint = "api.obscura.work"
|
|
minio_access_key = "MnHTAG2NOLyXXIZrwDLp"
|
|
minio_secret_key = "WVlmMgww0aRIU43pCJ1XCjubXQO6YsbHysxX2hBf"
|
|
|
|
mongo_uri = "mongodb://minio_mongo:[email protected]:27017/minio_mongo"
|
|
db_name = "minio_mongo"
|
|
collection_name = "douyin_results"
|
|
|
|
system = WhisperTranscriptionSystem(minio_endpoint, minio_access_key, minio_secret_key,
|
|
mongo_uri, db_name, collection_name)
|
|
system.run() |