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