321 lines
12 KiB
Python
321 lines
12 KiB
Python
import logging
|
|
import os
|
|
import subprocess
|
|
import tempfile
|
|
from PIL import Image
|
|
import imagehash
|
|
from sqlalchemy.orm import Session
|
|
from .samba_manager import SambaManager
|
|
from .config import Settings
|
|
from . import models
|
|
import itertools
|
|
import logging
|
|
import time
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class DuplicatesManager:
|
|
def __init__(self, settings: Settings, db: Session):
|
|
self.settings = settings
|
|
self.db = db
|
|
self.videos_root = "/videos"
|
|
|
|
def _get_video_duration(self, filepath):
|
|
logger.debug(f"Running ffprobe for duration of {filepath}")
|
|
try:
|
|
command = [
|
|
"ffprobe",
|
|
"-v",
|
|
"error",
|
|
"-show_entries",
|
|
"format=duration",
|
|
"-of",
|
|
"default=noprint_wrappers=1:nokey=1",
|
|
filepath,
|
|
]
|
|
logger.debug(f"ffprobe command: {' '.join(command)}")
|
|
result = subprocess.run(
|
|
command,
|
|
capture_output=True,
|
|
text=True,
|
|
check=True,
|
|
)
|
|
logger.debug(f"ffprobe stdout: {result.stdout.strip()}")
|
|
logger.debug(f"ffprobe stderr: {result.stderr.strip()}")
|
|
return float(result.stdout)
|
|
except (subprocess.CalledProcessError, FileNotFoundError) as e:
|
|
logger.error(f"ffprobe failed for {filepath}: {e}")
|
|
return None
|
|
|
|
def _get_frame_hash(self, filepath):
|
|
logger.debug(f"Running ffmpeg for frame hash of {filepath}")
|
|
tmp_frame_path = ""
|
|
try:
|
|
with tempfile.NamedTemporaryFile(suffix=".jpg", delete=False) as tmp_frame:
|
|
tmp_frame_path = tmp_frame.name
|
|
command = ["ffmpeg"]
|
|
if self.settings.gpu_enabled:
|
|
command.extend(["-hwaccel", "cuda"])
|
|
|
|
command.extend([
|
|
"-i",
|
|
filepath,
|
|
"-ss",
|
|
"00:00:10",
|
|
"-vframes",
|
|
"1",
|
|
"-y",
|
|
tmp_frame_path,
|
|
])
|
|
|
|
logger.debug(f"ffmpeg command: {' '.join(command)}")
|
|
result = subprocess.run(
|
|
command,
|
|
capture_output=True,
|
|
check=True,
|
|
)
|
|
logger.debug(f"ffmpeg stdout: {result.stdout.strip()}")
|
|
logger.debug(f"ffmpeg stderr: {result.stderr.strip()}")
|
|
|
|
if os.path.exists(tmp_frame_path):
|
|
logger.debug(f"Temporary frame file exists: {tmp_frame_path}, size: {os.path.getsize(tmp_frame_path)} bytes")
|
|
phash = imagehash.phash(Image.open(tmp_frame_path))
|
|
return str(phash)
|
|
else:
|
|
logger.warning(f"Temporary frame file was not created: {tmp_frame_path}")
|
|
return None
|
|
except (subprocess.CalledProcessError, FileNotFoundError) as e:
|
|
logger.error(f"ffmpeg failed for {filepath}: {e}")
|
|
return None
|
|
except Image.UnidentifiedImageError as e:
|
|
logger.error(f"PIL.UnidentifiedImageError for {filepath} with temp file {tmp_frame_path}: {e}")
|
|
return None
|
|
finally:
|
|
if os.path.exists(tmp_frame_path):
|
|
os.remove(tmp_frame_path)
|
|
|
|
def _is_video_file(self, filename):
|
|
video_extensions = ['.mp4', '.mkv', '.avi', '.mov', '.wmv', '.flv', '.webm']
|
|
return any(filename.lower().endswith(ext) for ext in video_extensions)
|
|
|
|
def _process_video_file(self, samba_manager: SambaManager, filepath, filename, size):
|
|
logger.info(f"Processing video: {filepath} ({filename})")
|
|
existing_video = self.db.query(models.VideoFile).filter_by(filepath=filepath).first()
|
|
if existing_video:
|
|
if existing_video.size == size:
|
|
logger.info(f"Skipping already processed and unaltered video: {filepath}")
|
|
return
|
|
else:
|
|
logger.info(f"File {filepath} has altered size ({existing_video.size} -> {size}). Re-processing.")
|
|
self.db.delete(existing_video)
|
|
self.db.commit()
|
|
|
|
with tempfile.NamedTemporaryFile(delete=False) as tmp_video:
|
|
try:
|
|
samba_manager.download_file(filepath, tmp_video)
|
|
except Exception as e:
|
|
logger.error(f"Failed to download {filepath}: {e}")
|
|
return
|
|
|
|
duration = self._get_video_duration(tmp_video.name)
|
|
phash = self._get_frame_hash(tmp_video.name)
|
|
|
|
os.remove(tmp_video.name)
|
|
|
|
if duration is not None and phash is not None:
|
|
video_file = models.VideoFile(
|
|
filepath=filepath,
|
|
filename=filename,
|
|
size=size,
|
|
duration=duration,
|
|
phash=phash,
|
|
)
|
|
self.db.add(video_file)
|
|
self.db.commit()
|
|
logger.info(f"Processed video: {filepath}")
|
|
else:
|
|
logger.warning(f"Could not get duration or hash for {filepath}")
|
|
|
|
def scan_videos(self, samba_manager: SambaManager):
|
|
state_file = "logs/duplicates_scan.state"
|
|
progress_file = "logs/duplicates_scan.progress"
|
|
pause_file = "logs/duplicates_scan.pause"
|
|
|
|
dirs_to_scan = []
|
|
|
|
try:
|
|
if os.path.exists(progress_file):
|
|
with open(progress_file, 'r') as f:
|
|
last_processed_dir = f.read().strip()
|
|
logger.info(f"Resuming scan from last in-progress directory: {last_processed_dir}")
|
|
dirs_to_scan.append(last_processed_dir)
|
|
|
|
if os.path.exists(state_file):
|
|
with open(state_file, 'r') as f:
|
|
dirs_to_scan.extend([line.strip() for line in f if line.strip()])
|
|
logger.info(f"Loaded {len(dirs_to_scan)} directories from state file.")
|
|
|
|
if not dirs_to_scan:
|
|
dirs_to_scan = [self.videos_root]
|
|
|
|
logger.info(f"Starting scan for videos in {self.videos_root} on share 'isolation'")
|
|
|
|
while dirs_to_scan:
|
|
while os.path.exists(pause_file):
|
|
logger.info("Scan is paused. Waiting for resume signal...")
|
|
time.sleep(5)
|
|
|
|
current_path = dirs_to_scan.pop(0)
|
|
|
|
with open(progress_file, 'w') as f:
|
|
f.write(current_path)
|
|
|
|
logger.info(f"Scanning directory: {current_path}")
|
|
files_and_dirs = samba_manager.list_path(current_path)
|
|
|
|
if "error" in files_and_dirs:
|
|
logger.error(f"Failed to list path {current_path}: {files_and_dirs['error']}")
|
|
continue
|
|
|
|
subdirs = []
|
|
for item in files_and_dirs:
|
|
while os.path.exists(pause_file):
|
|
logger.info("Scan is paused. Waiting for resume signal...")
|
|
time.sleep(5)
|
|
|
|
if item["is_directory"]:
|
|
subdirs.append(item["path"])
|
|
elif self._is_video_file(item["name"]):
|
|
self._process_video_file(samba_manager, item["path"], item["name"], item["size"])
|
|
|
|
dirs_to_scan = subdirs + dirs_to_scan
|
|
|
|
if os.path.exists(progress_file):
|
|
os.remove(progress_file)
|
|
|
|
with open(state_file, 'w') as f:
|
|
for d in dirs_to_scan:
|
|
f.write(d + '\n')
|
|
|
|
if os.path.exists(state_file):
|
|
os.remove(state_file)
|
|
|
|
logger.info("Video scan complete.")
|
|
return {"status": "Scan complete"}
|
|
|
|
except Exception as e:
|
|
logger.error(f"An error occurred during video scan: {e}", exc_info=True)
|
|
return {"status": "Scan failed", "error": str(e)}
|
|
|
|
def _hamming_distance(self, s1, s2):
|
|
return sum(c1 != c2 for c1, c2 in zip(s1, s2))
|
|
|
|
def _calculate_similarity(self, file1: models.VideoFile, file2: models.VideoFile):
|
|
size_similarity = 1 - (abs(file1.size - file2.size) / max(file1.size, file2.size))
|
|
duration_similarity = 1 - (abs(file1.duration - file2.duration) / max(file1.duration, file2.duration))
|
|
hash_similarity = 1 - (self._hamming_distance(file1.phash, file2.phash) / len(file1.phash))
|
|
return (size_similarity * 0.2) + (duration_similarity * 0.3) + (hash_similarity * 0.5)
|
|
|
|
def find_duplicates(self, threshold=0.95):
|
|
report = models.DuplicateReport(status="running")
|
|
self.db.add(report)
|
|
self.db.commit()
|
|
|
|
videos = self.db.query(models.VideoFile).all()
|
|
groups = []
|
|
processed_videos = set()
|
|
|
|
for video1, video2 in itertools.combinations(videos, 2):
|
|
if video1.id in processed_videos or video2.id in processed_videos:
|
|
continue
|
|
|
|
score = self._calculate_similarity(video1, video2)
|
|
if score >= threshold:
|
|
existing_group = None
|
|
for group in groups:
|
|
if video1.id in group["video_ids"] or video2.id in group["video_ids"]:
|
|
existing_group = group
|
|
break
|
|
|
|
if existing_group:
|
|
existing_group["video_ids"].add(video1.id)
|
|
existing_group["video_ids"].add(video2.id)
|
|
existing_group["scores"].append(score)
|
|
else:
|
|
groups.append({"video_ids": {video1.id, video2.id}, "scores": [score]})
|
|
|
|
processed_videos.add(video1.id)
|
|
processed_videos.add(video2.id)
|
|
|
|
for group_data in groups:
|
|
avg_score = sum(group_data["scores"]) / len(group_data["scores"])
|
|
db_group = models.DuplicateFileGroup(report_id=report.id, score=avg_score)
|
|
self.db.add(db_group)
|
|
self.db.commit()
|
|
|
|
for video_id in group_data["video_ids"]:
|
|
db_file = models.DuplicateFile(group_id=db_group.id, video_file_id=video_id)
|
|
self.db.add(db_file)
|
|
|
|
report.status = "completed"
|
|
self.db.commit()
|
|
|
|
return {"report_id": report.id, "status": "completed"}
|
|
|
|
def get_reports(self):
|
|
return self.db.query(models.DuplicateReport).all()
|
|
|
|
def get_duplicate_report(self, report_id: int):
|
|
report = self.db.query(models.DuplicateReport).filter(models.DuplicateReport.id == report_id).first()
|
|
if not report:
|
|
return {"error": "Report not found"}
|
|
|
|
groups = []
|
|
for group in report.groups:
|
|
files = []
|
|
for duplicate_file in group.files:
|
|
files.append(duplicate_file.video_file)
|
|
groups.append({
|
|
"group_id": group.id,
|
|
"score": group.score,
|
|
"files": files,
|
|
})
|
|
|
|
return {
|
|
"report_id": report.id,
|
|
"created_at": report.created_at,
|
|
"status": report.status,
|
|
"groups": groups,
|
|
}
|
|
|
|
def delete_files(self, filepaths: list[str]):
|
|
samba_manager = SambaManager(
|
|
self.settings.samba_server_ip,
|
|
"isolation", # Assuming all duplicates are in the isolation share
|
|
self.settings.samba_username,
|
|
self.settings.samba_password,
|
|
)
|
|
try:
|
|
for filepath in filepaths:
|
|
# Delete from Samba
|
|
samba_manager.delete_file(filepath)
|
|
|
|
# Delete from database
|
|
video_file = self.db.query(models.VideoFile).filter_by(filepath=filepath).first()
|
|
if video_file:
|
|
# Delete associations in DuplicateFile
|
|
self.db.query(models.DuplicateFile).filter_by(video_file_id=video_file.id).delete()
|
|
self.db.delete(video_file)
|
|
|
|
self.db.commit()
|
|
return {"status": "success"}
|
|
except Exception as e:
|
|
self.db.rollback()
|
|
logger.error(f"Error deleting files: {e}", exc_info=True)
|
|
return {"status": "error", "message": str(e)}
|
|
finally:
|
|
samba_manager.close()
|
|
|
|
|