488 lines
20 KiB
Python
488 lines
20 KiB
Python
import logging
|
|
import os
|
|
import subprocess
|
|
import tempfile
|
|
import json
|
|
import time
|
|
import itertools
|
|
from PIL import Image
|
|
import imagehash
|
|
from sqlalchemy.orm import Session
|
|
from .samba_manager import SambaManager
|
|
from .config import Settings
|
|
from .stash_service import StashService
|
|
from . import models
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
class CancellationException(Exception):
|
|
pass
|
|
|
|
class CancellableWriter:
|
|
def __init__(self, file_obj, check_cancel_func):
|
|
self.file_obj = file_obj
|
|
self.check_cancel_func = check_cancel_func
|
|
|
|
def write(self, data):
|
|
if self.check_cancel_func():
|
|
raise CancellationException("Scan canceled by user")
|
|
return self.file_obj.write(data)
|
|
|
|
def close(self):
|
|
return self.file_obj.close()
|
|
|
|
def flush(self):
|
|
return self.file_obj.flush()
|
|
|
|
def tell(self):
|
|
return self.file_obj.tell()
|
|
|
|
def seek(self, offset, whence=0):
|
|
return self.file_obj.seek(offset, whence)
|
|
|
|
class DuplicatesManager:
|
|
def __init__(self, settings: Settings, db: Session):
|
|
self.settings = settings
|
|
self.db = db
|
|
self.stash_service = StashService(settings)
|
|
self.exclusions_file = "resources/config/exclusions.json"
|
|
self._load_exclusions()
|
|
|
|
def _load_exclusions(self):
|
|
if os.path.exists(self.exclusions_file):
|
|
try:
|
|
with open(self.exclusions_file, 'r') as f:
|
|
self.exclusions = json.load(f)
|
|
except:
|
|
self.exclusions = []
|
|
else:
|
|
self.exclusions = []
|
|
|
|
def _save_exclusions(self):
|
|
os.makedirs(os.path.dirname(self.exclusions_file), exist_ok=True)
|
|
with open(self.exclusions_file, 'w') as f:
|
|
json.dump(self.exclusions, f)
|
|
|
|
def add_exclusion(self, path):
|
|
if path not in self.exclusions:
|
|
self.exclusions.append(path)
|
|
self._save_exclusions()
|
|
|
|
def remove_exclusion(self, path):
|
|
if path in self.exclusions:
|
|
self.exclusions.remove(path)
|
|
self._save_exclusions()
|
|
|
|
def _is_excluded(self, path):
|
|
for excl in self.exclusions:
|
|
if path.startswith(excl):
|
|
return True
|
|
return False
|
|
|
|
def _get_video_duration(self, filepath):
|
|
try:
|
|
command = [
|
|
"ffprobe", "-v", "error", "-show_entries", "format=duration",
|
|
"-of", "default=noprint_wrappers=1:nokey=1", filepath
|
|
]
|
|
result = subprocess.run(command, capture_output=True, text=True, check=True)
|
|
return float(result.stdout)
|
|
except Exception as e:
|
|
logger.error(f"ffprobe failed for {filepath}: {e}")
|
|
return None
|
|
|
|
def _get_frame_hash(self, filepath, algorithm='phash'):
|
|
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"])
|
|
|
|
# Extract frame at 10s or 10%? Fixed 10s for now.
|
|
command.extend(["-i", filepath, "-ss", "00:00:10", "-vframes", "1", "-y", tmp_frame_path])
|
|
|
|
subprocess.run(command, capture_output=True, check=True)
|
|
|
|
if os.path.exists(tmp_frame_path) and os.path.getsize(tmp_frame_path) > 0:
|
|
img = Image.open(tmp_frame_path)
|
|
if algorithm == 'ahash': h = imagehash.average_hash(img)
|
|
elif algorithm == 'dhash': h = imagehash.dhash(img)
|
|
else: h = imagehash.phash(img)
|
|
return str(h)
|
|
except Exception as e:
|
|
logger.error(f"Hash generation failed for {filepath}: {e}")
|
|
finally:
|
|
if os.path.exists(tmp_frame_path):
|
|
os.remove(tmp_frame_path)
|
|
return None
|
|
|
|
def generate_contact_sheet(self, video_path):
|
|
"""
|
|
Generates a 3x3 contact sheet for the video.
|
|
Returns (relative_path, phash_str)
|
|
"""
|
|
try:
|
|
duration = self._get_video_duration(video_path)
|
|
if not duration or duration < 10: return None, None
|
|
|
|
# Extract 9 frames at intervals
|
|
interval = duration / 10
|
|
timestamps = [interval * i for i in range(1, 10)]
|
|
|
|
# We use a temp dir to store frames, then stitch
|
|
# ffmpeg tile filter is good but seeking is faster for sparse frames on large files?
|
|
# Actually, `ffmpeg -i ... -vf fps=... tile=...` reads the whole file which is slow over network/SMB.
|
|
# Best to seek.
|
|
|
|
# Since we have the file locally in tmp_path (downloaded), seeking is fast.
|
|
|
|
frames = []
|
|
with tempfile.TemporaryDirectory() as temp_frames_dir:
|
|
for idx, ts in enumerate(timestamps):
|
|
out_frame = os.path.join(temp_frames_dir, f"frame_{idx}.jpg")
|
|
# fast seek
|
|
subprocess.run(
|
|
["ffmpeg", "-ss", str(ts), "-i", video_path, "-vframes", "1", "-q:v", "5", "-vf", "scale=320:-1", "-y", out_frame],
|
|
stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, check=False
|
|
)
|
|
if os.path.exists(out_frame):
|
|
frames.append(Image.open(out_frame))
|
|
|
|
if len(frames) < 4: return None, None # Need at least some frames
|
|
|
|
# Stitch 3x3 (or adaptive)
|
|
# Create blank image
|
|
w, h = frames[0].size
|
|
grid_w = w * 3
|
|
grid_h = h * 3
|
|
contact_sheet = Image.new('RGB', (grid_w, grid_h))
|
|
|
|
for idx, frame in enumerate(frames):
|
|
if idx >= 9: break
|
|
x = (idx % 3) * w
|
|
y = (idx // 3) * h
|
|
contact_sheet.paste(frame, (x, y))
|
|
|
|
# Save
|
|
cache_dir = "resources/cache/thumbnails"
|
|
os.makedirs(cache_dir, exist_ok=True)
|
|
# Use hash of path to ensure uniqueness/retrievability
|
|
filename_hash = imagehash.hex_to_hash(os.path.basename(video_path)) # Just use random or md5
|
|
import hashlib
|
|
file_hash = hashlib.md5(video_path.encode()).hexdigest()
|
|
out_name = f"{file_hash}.jpg"
|
|
out_path = os.path.join(cache_dir, out_name)
|
|
|
|
contact_sheet.save(out_path, "JPEG", quality=80)
|
|
|
|
# Calculate Hash of the SHEET
|
|
sheet_hash = imagehash.phash(contact_sheet)
|
|
|
|
return out_name, str(sheet_hash)
|
|
|
|
except Exception as e:
|
|
logger.error(f"Contact sheet generation failed: {e}")
|
|
return None, None
|
|
|
|
def _process_video_file(self, samba_manager, filepath, filename, size, algorithm='phash', scan_type='fast', log_func=None, cancel_check_func=None):
|
|
if cancel_check_func and cancel_check_func(): return
|
|
|
|
if self._is_excluded(filepath):
|
|
return
|
|
|
|
existing_video = self.db.query(models.VideoFile).filter_by(filepath=filepath).first()
|
|
|
|
# If in scene mode, check if we already have the scene data
|
|
if existing_video:
|
|
if scan_type == 'scene' and not existing_video.contact_sheet_path:
|
|
if log_func: log_func(f"Updating {filename} with contact sheet")
|
|
# Continue to processing
|
|
elif existing_video.size == size:
|
|
return # Skip if unchanged
|
|
else:
|
|
self.db.delete(existing_video)
|
|
self.db.commit()
|
|
existing_video = None
|
|
|
|
if log_func: log_func(f"Processing: {filename} (Mode: {scan_type})")
|
|
|
|
# --- Stash Integration ---
|
|
if self.settings.stash_enabled:
|
|
if cancel_check_func and cancel_check_func(): return
|
|
stash_phash, stash_oshash, scene_id, stash_duration = self.stash_service.get_file_metadata(filepath)
|
|
if stash_phash:
|
|
if log_func: log_func(f"Found Stash metadata for {filename}")
|
|
|
|
sheet_path = None
|
|
if scan_type == 'scene' and stash_oshash:
|
|
if cancel_check_func and cancel_check_func(): return
|
|
remote_sprite = self.stash_service.get_sprite_path(stash_oshash)
|
|
local_sheet_name = f"stash_{stash_oshash}.jpg"
|
|
local_sheet_path = os.path.join("resources/cache/thumbnails", local_sheet_name)
|
|
if not os.path.exists(local_sheet_path):
|
|
try:
|
|
# Ensure directory exists
|
|
os.makedirs(os.path.dirname(local_sheet_path), exist_ok=True)
|
|
|
|
# Use stash_share for sprite download
|
|
stash_samba = samba_manager
|
|
if samba_manager.share_name != self.settings.stash_share:
|
|
stash_samba = SambaManager(
|
|
self.settings.samba_server_ip,
|
|
self.settings.stash_share,
|
|
self.settings.samba_username,
|
|
self.settings.samba_password
|
|
)
|
|
|
|
try:
|
|
with open(local_sheet_path, "wb") as f:
|
|
stash_samba.download_file(remote_sprite, f)
|
|
sheet_path = local_sheet_name
|
|
finally:
|
|
if stash_samba != samba_manager:
|
|
stash_samba.close()
|
|
except Exception as e:
|
|
if log_func: log_func(f"Failed to download Stash sprite for {filename}: {e}")
|
|
else:
|
|
sheet_path = local_sheet_name
|
|
|
|
# Save to DB
|
|
video = models.VideoFile(
|
|
filepath=filepath, filename=filename, size=size,
|
|
duration=stash_duration or 0, phash=stash_phash,
|
|
contact_sheet_path=sheet_path,
|
|
scene_phash=None
|
|
)
|
|
self.db.add(video)
|
|
self.db.commit()
|
|
return
|
|
# --- End Stash Integration ---
|
|
|
|
if cancel_check_func and cancel_check_func(): return
|
|
|
|
# Download to temp
|
|
with tempfile.NamedTemporaryFile(delete=False) as tmp_video:
|
|
try:
|
|
if log_func: log_func(f"Downloading {filename} ({size/1024/1024:.2f} MB)...")
|
|
|
|
# Wrap for cancellation during download
|
|
writer = tmp_video
|
|
if cancel_check_func:
|
|
writer = CancellableWriter(tmp_video, cancel_check_func)
|
|
|
|
samba_manager.download_file(filepath, writer)
|
|
tmp_path = tmp_video.name
|
|
except CancellationException:
|
|
if log_func: log_func(f"Download aborted for {filename}")
|
|
tmp_video.close()
|
|
os.remove(tmp_video.name)
|
|
return
|
|
except Exception as e:
|
|
if log_func: log_func(f"Download failed: {filepath} - {e}")
|
|
tmp_video.close()
|
|
os.remove(tmp_video.name)
|
|
return
|
|
|
|
try:
|
|
if cancel_check_func and cancel_check_func(): return
|
|
|
|
duration = self._get_video_duration(tmp_path)
|
|
phash = self._get_frame_hash(tmp_path, algorithm)
|
|
|
|
if cancel_check_func and cancel_check_func(): return
|
|
|
|
sheet_path = None
|
|
scene_hash = None
|
|
|
|
if scan_type == 'scene':
|
|
sheet_path, scene_hash = self.generate_contact_sheet(tmp_path)
|
|
|
|
if duration is not None: # phash might be None if image generation failed
|
|
if not existing_video:
|
|
video = models.VideoFile(
|
|
filepath=filepath, filename=filename, size=size,
|
|
duration=duration, phash=phash or "",
|
|
contact_sheet_path=sheet_path,
|
|
scene_phash=scene_hash
|
|
)
|
|
self.db.add(video)
|
|
else:
|
|
# Update existing
|
|
existing_video.duration = duration
|
|
existing_video.phash = phash or ""
|
|
if sheet_path: existing_video.contact_sheet_path = sheet_path
|
|
if scene_hash: existing_video.scene_phash = scene_hash
|
|
|
|
self.db.commit()
|
|
finally:
|
|
if os.path.exists(tmp_path):
|
|
try: os.remove(tmp_path)
|
|
except: pass
|
|
|
|
def scan_videos(self, samba_manager, root_paths=["/videos"], algorithm='phash', scan_type='fast'):
|
|
state_file = "logs/duplicates_scan.state"
|
|
progress_file = "logs/duplicates_scan.json"
|
|
cancel_file = "logs/duplicates_scan.cancel"
|
|
log_file = "logs/duplicates_scan.log"
|
|
|
|
# Helper to log to file and console
|
|
def log(msg):
|
|
try:
|
|
with open(log_file, "a") as f:
|
|
f.write(f"{msg}\n")
|
|
except: pass
|
|
logger.info(msg)
|
|
|
|
is_cancelled = False
|
|
# Helper to check cancellation
|
|
def check_cancel():
|
|
nonlocal is_cancelled
|
|
if is_cancelled: return True
|
|
|
|
if os.path.exists(cancel_file):
|
|
log("Scan canceled by user.")
|
|
try:
|
|
os.remove(cancel_file)
|
|
except OSError:
|
|
pass
|
|
with open(progress_file, 'w') as f:
|
|
json.dump({"status": "canceled", "processed": processed_files}, f)
|
|
is_cancelled = True
|
|
return True
|
|
return False
|
|
|
|
# Clear log file
|
|
with open(log_file, "w") as f:
|
|
f.write("Scan started...\n")
|
|
|
|
# Ensure cancel file is gone before starting
|
|
if os.path.exists(cancel_file):
|
|
try: os.remove(cancel_file)
|
|
except OSError: pass
|
|
|
|
dirs_to_scan = list(root_paths)
|
|
processed_files = 0
|
|
|
|
try:
|
|
while dirs_to_scan:
|
|
if check_cancel(): return
|
|
|
|
current_path = dirs_to_scan.pop(0)
|
|
|
|
with open(progress_file, 'w') as f:
|
|
json.dump({"status": "scanning", "current": current_path, "processed": processed_files}, f)
|
|
|
|
try:
|
|
log(f"Scanning directory: {current_path}")
|
|
items = samba_manager.list_path(current_path)
|
|
except Exception as e:
|
|
log(f"Error listing {current_path}: {e}")
|
|
continue
|
|
|
|
for item in items:
|
|
if check_cancel(): return
|
|
|
|
if item['name'] in ['.', '..']: continue
|
|
|
|
if item['is_directory']:
|
|
if not self._is_excluded(item['path']):
|
|
dirs_to_scan.append(item['path'])
|
|
elif any(item['name'].lower().endswith(ext) for ext in ['.mp4', '.mkv', '.avi', '.mov', '.wmv']):
|
|
self._process_video_file(
|
|
samba_manager, item['path'], item['name'], item['size'],
|
|
algorithm, scan_type,
|
|
log_func=log, cancel_check_func=check_cancel
|
|
)
|
|
processed_files += 1
|
|
|
|
with open(progress_file, 'w') as f:
|
|
json.dump({"status": "scanning", "current": item['path'], "processed": processed_files}, f)
|
|
|
|
log("Scan completed.")
|
|
with open(progress_file, 'w') as f:
|
|
json.dump({"status": "completed", "processed": processed_files}, f)
|
|
|
|
except Exception as e:
|
|
log(f"Scan failed: {e}")
|
|
with open(progress_file, 'w') as f:
|
|
json.dump({"status": "failed", "error": str(e)}, f)
|
|
|
|
def find_duplicates(self, threshold=0.95, method='fast'):
|
|
report = models.DuplicateReport(status="running")
|
|
self.db.add(report)
|
|
self.db.commit()
|
|
|
|
videos = self.db.query(models.VideoFile).all()
|
|
groups = []
|
|
processed_ids = set()
|
|
|
|
for i in range(len(videos)):
|
|
if videos[i].id in processed_ids: continue
|
|
|
|
group = [videos[i]]
|
|
scores = []
|
|
|
|
for j in range(i + 1, len(videos)):
|
|
if videos[j].id in processed_ids: continue
|
|
|
|
v1, v2 = videos[i], videos[j]
|
|
|
|
score = 0
|
|
if method == 'scene' and v1.scene_phash and v2.scene_phash:
|
|
# Compare Scene Hashes
|
|
dist = imagehash.hex_to_hash(v1.scene_phash) - imagehash.hex_to_hash(v2.scene_phash)
|
|
score = max(0, 1.0 - (dist / 64.0)) # 64 is typical max distance for 8x8 hash
|
|
else:
|
|
# Standard Comparison
|
|
dist = imagehash.hex_to_hash(v1.phash) - imagehash.hex_to_hash(v2.phash) if v1.phash and v2.phash else 64
|
|
hash_sim = max(0, 1.0 - (dist / 64.0))
|
|
dur_sim = 1.0 - (abs(v1.duration - v2.duration) / max(v1.duration, v2.duration)) if max(v1.duration, v2.duration) > 0 else 1.0
|
|
size_sim = 1.0 - (abs(v1.size - v2.size) / max(v1.size, v2.size)) if max(v1.size, v2.size) > 0 else 1.0
|
|
score = (hash_sim * 0.6) + (dur_sim * 0.3) + (size_sim * 0.1)
|
|
|
|
if score >= threshold:
|
|
group.append(v2)
|
|
scores.append(score)
|
|
processed_ids.add(v2.id)
|
|
|
|
if len(group) > 1:
|
|
processed_ids.add(videos[i].id)
|
|
avg_score = sum(scores) / len(scores)
|
|
|
|
db_group = models.DuplicateFileGroup(report_id=report.id, score=avg_score)
|
|
self.db.add(db_group)
|
|
self.db.commit()
|
|
|
|
for v in group:
|
|
db_file = models.DuplicateFile(group_id=db_group.id, video_file_id=v.id)
|
|
self.db.add(db_file)
|
|
|
|
report.status = "completed"
|
|
self.db.commit()
|
|
return {"report_id": report.id}
|
|
|
|
def get_reports(self):
|
|
return self.db.query(models.DuplicateReport).order_by(models.DuplicateReport.created_at.desc()).all()
|
|
|
|
def get_report(self, report_id):
|
|
report = self.db.query(models.DuplicateReport).filter_by(id=report_id).first()
|
|
if not report: return None
|
|
|
|
res = {
|
|
"id": report.id,
|
|
"status": report.status,
|
|
"date": report.created_at,
|
|
"groups": []
|
|
}
|
|
for g in report.groups:
|
|
files = [{
|
|
"id": f.video_file.id,
|
|
"path": f.video_file.filepath,
|
|
"size": f.video_file.size,
|
|
"duration": f.video_file.duration,
|
|
"contact_sheet": f.video_file.contact_sheet_path
|
|
} for f in g.files]
|
|
res["groups"].append({"id": g.id, "score": g.score, "files": files})
|
|
return res |