update files

.gitignore

@@ -217,13 +217,15 @@ outputs
 yolo11n.pt
 yolo11s.pt
 best.pt
+**/best.pt
 
 .claude
-best.pt
 
 runs
 .idea
 working_dir
 data
 upload_to_huggingface.py
-resources/best.pt
+resources/best.pt
+resources/model_version.json
+.web

.streamlit/config.toml

@@ -0,0 +1,3 @@
+[server]
+maxUploadSize=4096 
+# 4GB as maximum

app.py

@@ -3,6 +3,7 @@ os.environ["HOME"] = "/tmp"
 os.environ["STREAMLIT_CONFIG_DIR"] = "/tmp/.streamlit"
 os.makedirs("/tmp/.streamlit", exist_ok=True)
 
+
 import shutil
 import tempfile
 from pathlib import Path
@@ -17,9 +18,75 @@ def main():
         page_title="Sora Watermark Cleaner", page_icon="🎬", layout="centered"
     )
 
-    st.title("🎬 Sora Watermark Cleaner")
-    st.markdown("Remove watermarks from Sora-generated videos with ease")
-
+    # Header section with improved layout
+    st.markdown(
+        """
+        <div style='text-align: center; padding: 1rem 0;'>
+            <h1 style='margin-bottom: 0.5rem;'>
+                🎬 Sora Watermark Cleaner
+            </h1>
+            <p style='font-size: 1.2rem; color: #666; margin-bottom: 1rem;'>
+                Remove watermarks from Sora-generated videos with AI-powered precision
+            </p>
+        </div>
+        """,
+        unsafe_allow_html=True,
+    )
+    
+    # # Feature badges
+    # col1, col2, col3 = st.columns(3)
+    # with col1:
+    #     st.markdown(
+    #         """
+    #         <div style='text-align: center; padding: 0.8rem; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); 
+    #                     border-radius: 10px; color: white;'>
+    #             <div style='font-size: 1.5rem;'>⚡</div>
+    #             <div style='font-weight: bold;'>Fast Processing</div>
+    #             <div style='font-size: 0.85rem; opacity: 0.9;'>GPU Accelerated</div>
+    #         </div>
+    #         """,
+    #         unsafe_allow_html=True,
+    #     )
+    # with col2:
+    #     st.markdown(
+    #         """
+    #         <div style='text-align: center; padding: 0.8rem; background: linear-gradient(135deg, #f093fb 0%, #f5576c 100%); 
+    #                     border-radius: 10px; color: white;'>
+    #             <div style='font-size: 1.5rem;'>🎯</div>
+    #             <div style='font-weight: bold;'>High Precision</div>
+    #             <div style='font-size: 0.85rem; opacity: 0.9;'>AI-Powered</div>
+    #         </div>
+    #         """,
+    #         unsafe_allow_html=True,
+    #     )
+    # with col3:
+    #     st.markdown(
+    #         """
+    #         <div style='text-align: center; padding: 0.8rem; background: linear-gradient(135deg, #4facfe 0%, #00f2fe 100%); 
+    #                     border-radius: 10px; color: white;'>
+    #             <div style='font-size: 1.5rem;'>📦</div>
+    #             <div style='font-weight: bold;'>Batch Support</div>
+    #             <div style='font-size: 0.85rem; opacity: 0.9;'>Process Multiple</div>
+    #         </div>
+    #         """,
+    #         unsafe_allow_html=True,
+    #     )
+    
+    # Footer info
+    st.markdown(
+        """
+        <div style='text-align: center; padding: 1rem 0; margin-top: 1rem;'>
+            <p style='color: #888; font-size: 0.9rem;'>
+                Built with ❤️ using Streamlit and AI | 
+                <a href='https://github.com/linkedlist771/SoraWatermarkCleaner' 
+                   target='_blank' style='color: #667eea; text-decoration: none;'>
+                    ⭐ Star on GitHub
+                </a>
+            </p>
+        </div>
+        """,
+        unsafe_allow_html=True,
+    )
     # Initialize SoraWM
     if "sora_wm" not in st.session_state:
         with st.spinner("Loading AI models..."):
@@ -27,84 +94,232 @@ def main():
 
     st.markdown("---")
 
-    # File uploader
-    uploaded_file = st.file_uploader(
-        "Upload your video",
-        type=["mp4", "avi", "mov", "mkv"],
-        help="Select a video file to remove watermarks",
+    # Mode selection
+    mode = st.radio(
+        "Select input mode:",
+        ["📁 Upload Video File", "🗂️ Process Folder"],
+        horizontal=True,
     )
 
-    if uploaded_file is not None:
-        # Display video info
-        st.success(f"✅ Uploaded: {uploaded_file.name}")
-        st.video(uploaded_file)
-
-        # Process button
-        if st.button("🚀 Remove Watermark", type="primary", use_container_width=True):
-            with tempfile.TemporaryDirectory() as tmp_dir:
-                tmp_path = Path(tmp_dir)
-
-                # Save uploaded file
-                input_path = tmp_path / uploaded_file.name
-                with open(input_path, "wb") as f:
-                    f.write(uploaded_file.read())
-
-                # Process video
-                output_path = tmp_path / f"cleaned_{uploaded_file.name}"
-
-                try:
-                    # Create progress bar and status text
-                    progress_bar = st.progress(0)
-                    status_text = st.empty()
-
-                    def update_progress(progress: int):
-                        progress_bar.progress(progress / 100)
-                        if progress < 50:
-                            status_text.text(f"🔍 Detecting watermarks... {progress}%")
-                        elif progress < 95:
-                            status_text.text(f"🧹 Removing watermarks... {progress}%")
-                        else:
-                            status_text.text(f"🎵 Merging audio... {progress}%")
-
-                    # Run the watermark removal with progress callback
-                    st.session_state.sora_wm.run(
-                        input_path, output_path, progress_callback=update_progress
-                    )
-
-                    # Complete the progress bar
-                    progress_bar.progress(100)
-                    status_text.text("✅ Processing complete!")
-
-                    st.success("✅ Watermark removed successfully!")
-
-                    # Display result
-                    st.markdown("### Result")
-                    st.video(str(output_path))
-
-                    # Download button
-                    with open(output_path, "rb") as f:
-                        st.download_button(
-                            label="⬇️ Download Cleaned Video",
-                            data=f,
-                            file_name=f"cleaned_{uploaded_file.name}",
-                            mime="video/mp4",
-                            use_container_width=True,
+    if mode == "📁 Upload Video File":
+        # File uploader
+        uploaded_file = st.file_uploader(
+            "Upload your video",
+            type=["mp4", "avi", "mov", "mkv"],
+            accept_multiple_files=False,
+            help="Select a video file to remove watermark",
+        )
+
+        if uploaded_file:
+            # Clear previous processed video if a new file is uploaded
+            if "current_file_name" not in st.session_state or st.session_state.current_file_name != uploaded_file.name:
+                st.session_state.current_file_name = uploaded_file.name
+                if "processed_video_data" in st.session_state:
+                    del st.session_state.processed_video_data
+                if "processed_video_path" in st.session_state:
+                    del st.session_state.processed_video_path
+                if "processed_video_name" in st.session_state:
+                    del st.session_state.processed_video_name
+            
+            # Display video info
+            st.success(f"✅ Uploaded: {uploaded_file.name}")
+            
+            # Create two columns for before/after comparison
+            col_left, col_right = st.columns(2)
+            
+            with col_left:
+                st.markdown("### 📥 Original Video")
+                st.video(uploaded_file)
+            
+            with col_right:
+                st.markdown("### 🎬 Processed Video")
+                # Placeholder for processed video
+                if "processed_video_data" not in st.session_state:
+                    st.info("Click 'Remove Watermark' to process the video")
+                else:
+                    st.video(st.session_state.processed_video_data)
+
+            # Process button
+            if st.button("🚀 Remove Watermark", type="primary", use_container_width=True):
+                with tempfile.TemporaryDirectory() as tmp_dir:
+                    tmp_path = Path(tmp_dir)
+
+                    try:
+                        # Create progress bar and status text
+                        progress_bar = st.progress(0)
+                        status_text = st.empty()
+
+                        def update_progress(progress: int):
+                            progress_bar.progress(progress / 100)
+                            if progress < 50:
+                                status_text.text(f"🔍 Detecting watermarks... {progress}%")
+                            elif progress < 95:
+                                status_text.text(f"🧹 Removing watermarks... {progress}%")
+                            else:
+                                status_text.text(f"🎵 Merging audio... {progress}%")
+
+                        # Single file processing
+                        input_path = tmp_path / uploaded_file.name
+                        with open(input_path, "wb") as f:
+                            f.write(uploaded_file.read())
+
+                        output_path = tmp_path / f"cleaned_{uploaded_file.name}"
+
+                        st.session_state.sora_wm.run(
+                            input_path, output_path, progress_callback=update_progress
                         )
 
-                except Exception as e:
-                    st.error(f"❌ Error processing video: {str(e)}")
+                        progress_bar.progress(100)
+                        status_text.text("✅ Processing complete!")
+                        st.success("✅ Watermark removed successfully!")
 
-    # Footer
-    st.markdown("---")
-    st.markdown(
-        """
-        <div style='text-align: center'>
-            <p>Built with ❤️ using Streamlit and AI</p>
-            <p><a href='https://github.com/linkedlist771/SoraWatermarkCleaner'>GitHub Repository</a></p>
-        </div>
-        """,
-        unsafe_allow_html=True,
-    )
+                        # Store processed video path and read video data
+                        with open(output_path, "rb") as f:
+                            video_data = f.read()
+                        
+                        st.session_state.processed_video_path = output_path
+                        st.session_state.processed_video_data = video_data
+                        st.session_state.processed_video_name = f"cleaned_{uploaded_file.name}"
+                        
+                        # Rerun to show the video in the right column
+                        st.rerun()
+
+                    except Exception as e:
+                        st.error(f"❌ Error processing video: {str(e)}")
+            
+            # Download button (show only if video is processed)
+            if "processed_video_data" in st.session_state:
+                st.download_button(
+                    label="⬇️ Download Cleaned Video",
+                    data=st.session_state.processed_video_data,
+                    file_name=st.session_state.processed_video_name,
+                    mime="video/mp4",
+                    use_container_width=True,
+                )
+
+    else:  # Folder mode
+        st.info("💡 Drag and drop your video folder here, or click to browse and select multiple video files")
+        
+        # File uploader for multiple files (supports folder drag & drop)
+        uploaded_files = st.file_uploader(
+            "Upload videos from folder",
+            type=["mp4", "avi", "mov", "mkv"],
+            accept_multiple_files=True,
+            help="You can drag & drop an entire folder here, or select multiple video files",
+            key="folder_uploader"
+        )
+        
+        if uploaded_files:
+            # Display uploaded files info
+            video_count = len(uploaded_files)
+            st.success(f"✅ {video_count} video file(s) uploaded")
+            
+            # Show file list in an expander
+            with st.expander("📋 View uploaded files", expanded=False):
+                for i, file in enumerate(uploaded_files, 1):
+                    file_size_mb = file.size / (1024 * 1024)
+                    st.text(f"{i}. {file.name} ({file_size_mb:.2f} MB)")
+            
+            # Process button
+            if st.button("🚀 Process All Videos", type="primary", use_container_width=True):
+                with tempfile.TemporaryDirectory() as tmp_dir:
+                    tmp_path = Path(tmp_dir)
+                    input_folder = tmp_path / "input"
+                    output_folder = tmp_path / "output"
+                    input_folder.mkdir(exist_ok=True)
+                    output_folder.mkdir(exist_ok=True)
+                    
+                    try:
+                        # Save all uploaded files to temp folder
+                        status_text = st.empty()
+                        status_text.text("📥 Saving uploaded files...")
+                        
+                        for uploaded_file in uploaded_files:
+                            # Preserve folder structure if file.name contains subdirectories
+                            file_path = input_folder / uploaded_file.name
+                            file_path.parent.mkdir(parents=True, exist_ok=True)
+                            with open(file_path, "wb") as f:
+                                f.write(uploaded_file.read())
+                        
+                        # Create progress tracking
+                        progress_bar = st.progress(0)
+                        current_file_text = st.empty()
+                        processed_count = 0
+                        
+                        def update_progress(progress: int):
+                            # Calculate overall progress
+                            overall_progress = (processed_count * 100 + progress) / video_count / 100
+                            progress_bar.progress(overall_progress)
+                            
+                            if progress < 50:
+                                current_file_text.text(f"🔍 Processing file {processed_count + 1}/{video_count}: Detecting watermarks... {progress}%")
+                            elif progress < 95:
+                                current_file_text.text(f"🧹 Processing file {processed_count + 1}/{video_count}: Removing watermarks... {progress}%")
+                            else:
+                                current_file_text.text(f"🎵 Processing file {processed_count + 1}/{video_count}: Merging audio... {progress}%")
+                        
+                        # Process each video file
+                        for video_file in input_folder.rglob("*"):
+                            if video_file.is_file() and video_file.suffix.lower() in [".mp4", ".avi", ".mov", ".mkv"]:
+                                # Determine output path maintaining folder structure
+                                rel_path = video_file.relative_to(input_folder)
+                                output_path = output_folder / rel_path.parent / f"cleaned_{rel_path.name}"
+                                output_path.parent.mkdir(parents=True, exist_ok=True)
+                                
+                                # Process the video
+                                st.session_state.sora_wm.run(
+                                    video_file, output_path, progress_callback=update_progress
+                                )
+                                processed_count += 1
+                        
+                        progress_bar.progress(100)
+                        current_file_text.text("✅ All videos processed!")
+                        st.success(f"✅ {video_count} video(s) processed successfully!")
+                        
+                        # Create download option for processed videos
+                        st.markdown("### 📦 Download Processed Videos")
+                        
+                        # Store processed files info in session state
+                        if "batch_processed_files" not in st.session_state:
+                            st.session_state.batch_processed_files = []
+                        
+                        st.session_state.batch_processed_files.clear()
+                        
+                        for processed_file in output_folder.rglob("*"):
+                            if processed_file.is_file():
+                                with open(processed_file, "rb") as f:
+                                    video_data = f.read()
+                                rel_path = processed_file.relative_to(output_folder)
+                                st.session_state.batch_processed_files.append({
+                                    "name": str(rel_path),
+                                    "data": video_data
+                                })
+                        
+                        st.rerun()
+                        
+                    except Exception as e:
+                        st.error(f"❌ Error processing videos: {str(e)}")
+                        import traceback
+                        st.error(f"Details: {traceback.format_exc()}")
+            
+            # Show download buttons for processed files
+            if "batch_processed_files" in st.session_state and st.session_state.batch_processed_files:
+                st.markdown("---")
+                st.markdown("### ⬇️ Download Processed Videos")
+                
+                for file_info in st.session_state.batch_processed_files:
+                    col1, col2 = st.columns([3, 1])
+                    with col1:
+                        st.text(f"📹 {file_info['name']}")
+                    with col2:
+                        st.download_button(
+                            label="⬇️ Download",
+                            data=file_info['data'],
+                            file_name=file_info['name'],
+                            mime="video/mp4",
+                            key=f"download_{file_info['name']}",
+                            use_container_width=True
+                        )
 
 
 if __name__ == "__main__":

datasets/make_yolo_images.py

@@ -2,48 +2,58 @@ from pathlib import Path
 
 import cv2
 from tqdm import tqdm
-
+from sorawm.watermark_detector import SoraWaterMarkDetector
 from sorawm.configs import ROOT
 
 videos_dir = ROOT / "videos"
 datasets_dir = ROOT / "datasets"
 images_dir = datasets_dir / "images"
 images_dir.mkdir(exist_ok=True, parents=True)
+detector = SoraWaterMarkDetector()
+
 
 if __name__ == "__main__":
-    fps_save_interval = 1  # Save every 1th frame
+    fps_save_interval = 1  # Save every 5th frame
 
-    idx = 0
+    video_idx = 0
+    image_idx = 0  # 全局图片索引
+    total_failed = 0  # 检测失败的总数
+    
     for video_path in tqdm(list(videos_dir.rglob("*.mp4"))):
         # Open the video file
         cap = cv2.VideoCapture(str(video_path))
-
+        video_name = video_path.name
         if not cap.isOpened():
             print(f"Error opening video: {video_path}")
             continue
 
         frame_count = 0
 
-        while True:
-            ret, frame = cap.read()
-
-            # Break if no more frames
-            if not ret:
-                break
-
-            # Save frame at the specified interval
-            if frame_count % fps_save_interval == 0:
-                # Create filename: image_idx_framecount.jpg
-                image_filename = f"image_{idx:06d}_frame_{frame_count:06d}.jpg"
-                image_path = images_dir / image_filename
-
-                # Save the frame
-                cv2.imwrite(str(image_path), frame)
-
-            frame_count += 1
-
-        # Release the video capture object
-        cap.release()
-        idx += 1
-
-    print(f"Processed {idx} videos, extracted frames saved to {images_dir}")
+        try:
+            while True:
+                ret, frame = cap.read()
+
+                # Break if no more frames
+                if not ret:
+                    break
+
+                # Save frame at the specified interval
+                if frame_count % fps_save_interval == 0:
+                    if not detector.detect(frame)["detected"]:
+                        # Create filename: image_idx_framecount.jpg
+                        image_filename = f"{video_name}_failed_image_frame_{frame_count:06d}.jpg"
+                        image_path = images_dir / image_filename
+                        # Save the frame
+                        cv2.imwrite(str(image_path), frame)
+                        image_idx += 1
+                        total_failed += 1
+
+                frame_count += 1
+
+        finally:
+            # Release the video capture object
+            cap.release()
+        
+        video_idx += 1
+
+    print(f"Processed {video_idx} videos, extracted {total_failed} failed detection frames to {images_dir}")

mds/reward.md

@@ -0,0 +1 @@
+![](../assests/wechat_reward.jpg)

model_version.json

@@ -0,0 +1 @@
+{"sha256": "79b44170111bd206d4964966b3b35adef1b3b15e7acf6427a95d35a2c715f987"}

pyproject.toml

@@ -12,6 +12,7 @@ dependencies = [
     "fastapi==0.108.0",
     "ffmpeg-python>=0.2.0",
     "fire>=0.7.1",
+	"greenlet>=3.2.4",
     "httpx>=0.28.1",
     "huggingface-hub>=0.35.3",
     "jupyter>=1.1.1",

resources/19700121_1645_68e0a027836c8191a50bea3717ea7485.mp4

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:b7dffeb423765b5fa4ff4dc3571c5c83d45baee0c6abb86d512733fa75b27a72
-size 4336615

resources/dog_vs_sam.mp4

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:bb22d5365c17e2bb73a9bed0598e057fc836ef60aa997982fea885e53e73283b
-size 5603713

resources/puppies.mp4

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:7dd8a9d43d43cd5b901946c7ba14e8018e18d47e75f0f774157327ba7e9181c5
-size 10866692

resources/sora_watermark_removed.mp4

@@ -1,3 +0,0 @@
-version https://git-lfs.github.com/spec/v1
-oid sha256:a7c1776fdb29087a57c73a7d66e02684f7c32f5245aa442cb2b76c64c5f9f7a1
-size 1570689

sorawm/configs.py

@@ -8,6 +8,9 @@ WATER_MARK_TEMPLATE_IMAGE_PATH = RESOURCES_DIR / "watermark_template.png"
 
 WATER_MARK_DETECT_YOLO_WEIGHTS = RESOURCES_DIR / "best.pt"
 
+WATER_MARK_DETECT_YOLO_WEIGHTS_HASH_JSON = RESOURCES_DIR / "model_version.json"
+
+
 
 OUTPUT_DIR = Path("/tmp/output")
 OUTPUT_DIR.mkdir(exist_ok=True, parents=True)

sorawm/core.py

@@ -15,12 +15,39 @@ from sorawm.utils.imputation_utils import (
     find_idxs_interval,
 )
 
+VIDEO_EXTENSIONS = [".mp4", ".avi", ".mov", ".mkv", ".flv", ".wmv", ".webm"]
 
 class SoraWM:
     def __init__(self):
         self.detector = SoraWaterMarkDetector()
         self.cleaner = WaterMarkCleaner()
 
+    def run_batch(self, input_video_dir_path: Path,
+        output_video_dir_path: Path | None = None,
+        progress_callback: Callable[[int], None] | None = None,
+        ):
+        if output_video_dir_path is None:
+            output_video_dir_path = input_video_dir_path.parent / "watermark_removed"
+            logger.warning(f"output_video_dir_path is not set, using {output_video_dir_path} as output_video_dir_path")
+        output_video_dir_path.mkdir(parents=True, exist_ok=True)        
+        input_video_paths = []
+        for ext in VIDEO_EXTENSIONS:
+            input_video_paths.extend(input_video_dir_path.rglob(f"*{ext}"))
+        
+        video_lengths = len(input_video_paths)
+        logger.info(f"Found {video_lengths} video(s) to process")
+        
+        for idx, input_video_path in enumerate(tqdm(input_video_paths, desc="Processing videos")):
+            output_video_path = output_video_dir_path / input_video_path.name            
+            if progress_callback:
+                def batch_progress_callback(single_video_progress: int):
+                    overall_progress = int((idx / video_lengths) * 100 + (single_video_progress / video_lengths))
+                    progress_callback(min(overall_progress, 100))
+                
+                self.run(input_video_path, output_video_path, progress_callback=batch_progress_callback)
+            else:
+                self.run(input_video_path, output_video_path, progress_callback=None)
+
     def run(
         self,
         input_video_path: Path,
@@ -62,7 +89,7 @@ class SoraWM:
             .run_async(pipe_stdin=True)
         )
 
-        frame_and_mask = {}
+        frame_bboxes = {}
         detect_missed = []
         bbox_centers = []
         bboxes = []
@@ -75,13 +102,13 @@ class SoraWM:
         ):
             detection_result = self.detector.detect(frame)
             if detection_result["detected"]:
-                frame_and_mask[idx] = {"frame": frame, "bbox": detection_result["bbox"]}
+                frame_bboxes[idx] = { "bbox": detection_result["bbox"]}
                 x1, y1, x2, y2 = detection_result["bbox"]
                 bbox_centers.append((int((x1 + x2) / 2), int((y1 + y2) / 2)))
                 bboxes.append((x1, y1, x2, y2))
 
             else:
-                frame_and_mask[idx] = {"frame": frame, "bbox": None}
+                frame_bboxes[idx] = {"bbox": None}
                 detect_missed.append(idx)
                 bbox_centers.append(None)
                 bboxes.append(None)
@@ -115,28 +142,30 @@ class SoraWM:
                     interval_idx < len(interval_bboxes)
                     and interval_bboxes[interval_idx] is not None
                 ):
-                    frame_and_mask[missed_idx]["bbox"] = interval_bboxes[interval_idx]
+                    frame_bboxes[missed_idx]["bbox"] = interval_bboxes[interval_idx]
                     logger.debug(f"Filled missed frame {missed_idx} with bbox:\n"
                     f" {interval_bboxes[interval_idx]}")
                 else:
                     # if the interval has no valid bbox, use the previous and next frame to complete (fallback strategy)
                     before = max(missed_idx - 1, 0)
                     after = min(missed_idx + 1, total_frames - 1)
-                    before_box = frame_and_mask[before]["bbox"]
-                    after_box = frame_and_mask[after]["bbox"]
+                    before_box = frame_bboxes[before]["bbox"]
+                    after_box = frame_bboxes[after]["bbox"]
                     if before_box:
-                        frame_and_mask[missed_idx]["bbox"] = before_box
+                        frame_bboxes[missed_idx]["bbox"] = before_box
                     elif after_box:
-                        frame_and_mask[missed_idx]["bbox"] = after_box
+                        frame_bboxes[missed_idx]["bbox"] = after_box
         else:
             del bboxes
             del bbox_centers
             del detect_missed
+        
+        input_video_loader = VideoLoader(input_video_path)
 
-        for idx in tqdm(range(total_frames), desc="Remove watermarks"):
-            frame_info = frame_and_mask[idx]
-            frame = frame_info["frame"]
-            bbox = frame_info["bbox"]
+        for idx, frame in enumerate(tqdm(input_video_loader, total=total_frames, desc="Remove watermarks")):
+        # for idx in tqdm(range(total_frames), desc="Remove watermarks"):
+            # frame_info = 
+            bbox = frame_bboxes[idx]["bbox"]
             if bbox is not None:
                 x1, y1, x2, y2 = bbox
                 mask = np.zeros((height, width), dtype=np.uint8)

sorawm/utils/download_utils.py

@@ -3,33 +3,90 @@ from pathlib import Path
 import requests
 from loguru import logger
 from tqdm import tqdm
-
-from sorawm.configs import WATER_MARK_DETECT_YOLO_WEIGHTS
+import hashlib
+import json
+from sorawm.configs import WATER_MARK_DETECT_YOLO_WEIGHTS, WATER_MARK_DETECT_YOLO_WEIGHTS_HASH_JSON
 
 DETECTOR_URL = "https://github.com/linkedlist771/SoraWatermarkCleaner/releases/download/V0.0.1/best.pt"
+REMOTE_MODEL_VERSION_URL = "https://raw.githubusercontent.com/linkedlist771/SoraWatermarkCleaner/refs/heads/main/model_version.json"
 
+def generate_sha256_hash(file_path: Path) -> str:
+    with open(file_path, "rb") as f:
+        return hashlib.sha256(f.read()).hexdigest()
 
-def download_detector_weights():
-    if not WATER_MARK_DETECT_YOLO_WEIGHTS.exists():
-        logger.debug(f"llama weights not found, downloading from {DETECTOR_URL}")
-        WATER_MARK_DETECT_YOLO_WEIGHTS.parent.mkdir(parents=True, exist_ok=True)
-
+def download_detector_weights(force_download: bool = False):
+    ## 1. check if model exists and if we need to download
+    if not WATER_MARK_DETECT_YOLO_WEIGHTS.exists() or force_download:
+        logger.debug(f"Downloading weights from {DETECTOR_URL}")
+        WATER_MARK_DETECT_YOLO_WEIGHTS.parent.mkdir(parents=True, exist_ok=True)        
+        temp_file = WATER_MARK_DETECT_YOLO_WEIGHTS.with_suffix(".tmp")
+        
         try:
             response = requests.get(DETECTOR_URL, stream=True, timeout=300)
             response.raise_for_status()
-            total_size = int(response.headers.get("content-length", 0))
-            with open(WATER_MARK_DETECT_YOLO_WEIGHTS, "wb") as f:
+            total_size = int(response.headers.get("content-length", 0))            
+            with open(temp_file, "wb") as f:
                 with tqdm(
                     total=total_size, unit="B", unit_scale=True, desc="Downloading"
                 ) as pbar:
                     for chunk in response.iter_content(chunk_size=8192):
                         if chunk:
                             f.write(chunk)
-                            pbar.update(len(chunk))
-
+                            pbar.update(len(chunk))            
+            if WATER_MARK_DETECT_YOLO_WEIGHTS.exists():
+                WATER_MARK_DETECT_YOLO_WEIGHTS.unlink() 
+            temp_file.rename(WATER_MARK_DETECT_YOLO_WEIGHTS)  
+            
             logger.success(f"✓ Weights downloaded: {WATER_MARK_DETECT_YOLO_WEIGHTS}")
-
+            new_hash = generate_sha256_hash(WATER_MARK_DETECT_YOLO_WEIGHTS)
+            WATER_MARK_DETECT_YOLO_WEIGHTS_HASH_JSON.parent.mkdir(parents=True, exist_ok=True)
+            with WATER_MARK_DETECT_YOLO_WEIGHTS_HASH_JSON.open("w") as f:
+                json.dump({"sha256": new_hash}, f)
+            logger.debug(f"Hash updated: {new_hash[:8]}...")
+            return 
+            
         except requests.exceptions.RequestException as e:
-            if WATER_MARK_DETECT_YOLO_WEIGHTS.exists():
-                WATER_MARK_DETECT_YOLO_WEIGHTS.unlink()
-            raise RuntimeError(f"Downing failed: {e}")
+            if temp_file.exists():
+                temp_file.unlink()
+            raise RuntimeError(f"Download failed: {e}")
+
+    ## 2. check the local hash, if it exits, compare it with the remote one(with timeout)
+    ## if not, generate it then compare.
+    local_sha256_hash = None 
+    # WATER_MARK_DETECT_YOLO_WEIGHTS_HASH_JSON
+    if not WATER_MARK_DETECT_YOLO_WEIGHTS_HASH_JSON.exists():
+        pass 
+    else:
+        # if it has not hash,
+        with WATER_MARK_DETECT_YOLO_WEIGHTS_HASH_JSON.open("r") as f:
+            hash_data = json.load(f)
+            local_sha256_hash = hash_data.get("sha256", None)
+
+    if local_sha256_hash is None:
+        
+        # generate the hash and update the config
+        logger.info(f"Generating sha256 hash for {WATER_MARK_DETECT_YOLO_WEIGHTS}")
+        local_sha256_hash = generate_sha256_hash(WATER_MARK_DETECT_YOLO_WEIGHTS)
+        WATER_MARK_DETECT_YOLO_WEIGHTS_HASH_JSON.parent.mkdir(parents=True, exist_ok=True)
+        with WATER_MARK_DETECT_YOLO_WEIGHTS_HASH_JSON.open("w") as f:
+            json.dump({"sha256": local_sha256_hash}, f)
+    remote_sha256_hash = None 
+    try:
+        response = requests.get(REMOTE_MODEL_VERSION_URL, timeout=10)
+        response.raise_for_status()
+        remote_sha256_hash = response.json().get("sha256", None)
+    except requests.exceptions.RequestException as e:
+        logger.error(f"Failed to get remote sha256 hash: {e}")
+        remote_sha256_hash = None
+
+    ## 3. after the compare, if there is a new version, download it and replace the local and 
+    ## update the hash
+    logger.debug(f"Local hash: {local_sha256_hash}, Remote hash: {remote_sha256_hash}")
+    if remote_sha256_hash is None:
+       pass 
+    else:
+        if local_sha256_hash != remote_sha256_hash:
+            logger.info(f"Hash mismatch detected, updating model...")
+            download_detector_weights(force_download=True)
+        else:
+            logger.debug("Model is up-to-date") 

sorawm/watermark_cleaner.py

@@ -36,294 +36,3 @@ class WaterMarkCleaner:
         inpaint_result = cv2.cvtColor(inpaint_result, cv2.COLOR_BGR2RGB)
         return inpaint_result
 
-
-if __name__ == "__main__":
-    from pathlib import Path
-
-    import cv2
-    import numpy as np
-    from tqdm import tqdm
-
-    # ========= 配置 =========
-    video_path = Path("resources/puppies.mp4")
-    save_video = True
-    out_path = Path("outputs/dog_vs_sam_detected.mp4")
-    window = "Sora watermark (threshold+morph+shape + tracking)"
-
-    # 追踪/回退策略参数
-    PREV_ROI_EXPAND = 2.2  # 上一框宽高的膨胀倍数（>1）
-    AREA1 = (1000, 2000)  # 主检测面积范围
-    AREA2 = (600, 4000)  # 回退阶段面积范围
-    # =======================
-
-    cleaner = SoraWaterMarkCleaner(video_path, video_path)
-
-    # 预取一帧确定尺寸/FPS
-    first_frame = None
-    for first_frame in cleaner.input_video_loader:
-        break
-    assert first_frame is not None, "无法读取视频帧"
-    H, W = first_frame.shape[:2]
-    fps = getattr(cleaner.input_video_loader, "fps", 30)
-
-    # 输出视频（原 | bw | all-contours | vis 四联画）
-    writer = None
-    if save_video:
-        out_path.parent.mkdir(parents=True, exist_ok=True)
-        fourcc = cv2.VideoWriter_fourcc(*"avc1")
-        writer = cv2.VideoWriter(str(out_path), fourcc, fps, (W * 4, H))
-        if not writer.isOpened():
-            fourcc = cv2.VideoWriter_fourcc(*"MJPG")
-            writer = cv2.VideoWriter(str(out_path), fourcc, fps, (W * 4, H))
-        assert writer.isOpened(), "无法创建输出视频文件"
-
-    cv2.namedWindow(window, cv2.WINDOW_NORMAL)
-
-    # ---- 工具函数 ----
-    def _clip_rect(x0, y0, x1, y1, w_img, h_img):
-        x0 = max(0, min(x0, w_img - 1))
-        x1 = max(0, min(x1, w_img))
-        y0 = max(0, min(y0, h_img - 1))
-        y1 = max(0, min(y1, h_img))
-        if x1 <= x0:
-            x1 = x0 + 1
-        if y1 <= y0:
-            y1 = y0 + 1
-        return x0, y0, x1, y1
-
-    def _cnt_bbox(cnt):
-        x, y, w, h = cv2.boundingRect(cnt)
-        return (x, y, x + w, y + h)
-
-    def _bbox_center(b):
-        x0, y0, x1, y1 = b
-        return ((x0 + x1) // 2, (y0 + y1) // 2)
-
-    def detect_flower_like(image, prev_bbox=None):
-        """
-        识别流程：
-        灰度范围 → 自适应阈值 → 仅在 3 个区域 + (可选)上一帧膨胀ROI 内找轮廓
-        三个区域：1) 左上20%  2) 左下20%  3) 中间水平带 y∈[0.4H, 0.6H], x∈[0,W]
-        返回: bw_region, best_cnt, contours_region, region_boxes, prev_roi_box
-        """
-        gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
-
-        # 208 ± 20% 亮度范围
-        low, high = int(round(208 * 0.9)), int(round(208 * 1.1))
-        mask = ((gray >= low) & (gray <= high)).astype(np.uint8) * 255
-
-        # 自适应阈值并限制到亮度范围
-        bw = cv2.adaptiveThreshold(
-            gray, 255, cv2.ADAPTIVE_THRESH_MEAN_C, cv2.THRESH_BINARY, 31, -5
-        )
-        bw = cv2.bitwise_and(bw, mask)
-
-        # -------- 三个区域：左上/左下/中间带 --------
-        h_img, w_img = gray.shape[:2]
-        r_top_left = (0, 0, int(0.2 * w_img), int(0.2 * h_img))
-        r_bot_left = (0, int(0.8 * h_img), int(0.2 * w_img), h_img)
-        y0, y1 = int(0.40 * h_img), int(0.60 * h_img)  # 中间带
-        r_mid_band = (0, y0, w_img, y1)
-
-        region_mask = np.zeros_like(bw, dtype=np.uint8)
-        for x0, ys, x1, ye in (r_top_left, r_bot_left):
-            region_mask[ys:ye, x0:x1] = 255
-        region_mask[y0:y1, :] = 255
-
-        # -------- 追加：上一帧膨胀ROI --------
-        prev_roi_box = None
-        if prev_bbox is not None:
-            px0, py0, px1, py1 = prev_bbox
-            pw, ph = (px1 - px0), (py1 - py0)
-            cx, cy = _bbox_center(prev_bbox)
-            rw = int(pw * PREV_ROI_EXPAND)
-            rh = int(ph * PREV_ROI_EXPAND)
-            rx0, ry0 = cx - rw // 2, cy - rh // 2
-            rx1, ry1 = cx + rw // 2, cy + rh // 2
-            rx0, ry0, rx1, ry1 = _clip_rect(rx0, ry0, rx1, ry1, w_img, h_img)
-            region_mask[ry0:ry1, rx0:rx1] = 255
-            prev_roi_box = (rx0, ry0, rx1, ry1)
-
-        bw_region = cv2.bitwise_and(bw, region_mask)
-
-        # -------- 轮廓 + 形状筛选 --------
-        def select_candidates(bw_bin, area_rng):
-            contours, _ = cv2.findContours(
-                bw_bin, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE
-            )
-            cand = []
-            for cnt in contours:
-                area = cv2.contourArea(cnt)
-                if area < area_rng[0] or area > area_rng[1]:
-                    continue
-                peri = cv2.arcLength(cnt, True)
-                if peri == 0:
-                    continue
-                circularity = 4.0 * np.pi * area / (peri * peri)
-                if 0.55 <= circularity <= 0.95:
-                    cand.append(cnt)
-            return contours, cand
-
-        contours_region, cand1 = select_candidates(bw_region, AREA1)
-
-        best_cnt = None
-        if cand1:
-            # 若有上一帧，用“离上一帧中心最近”优先；否则取面积最大
-            if prev_bbox is None:
-                best_cnt = max(cand1, key=lambda c: cv2.contourArea(c))
-            else:
-                pcx, pcy = _bbox_center(prev_bbox)
-                best_cnt = max(
-                    cand1,
-                    key=lambda c: -(
-                        (((_cnt_bbox(c)[0] + _cnt_bbox(c)[2]) // 2 - pcx) ** 2)
-                        + (((_cnt_bbox(c)[1] + _cnt_bbox(c)[3]) // 2 - pcy) ** 2)
-                    ),
-                )
-        else:
-            # 回退1：仅在上一帧 ROI 内放宽面积
-            if prev_roi_box is not None:
-                rx0, ry0, rx1, ry1 = prev_roi_box
-                roi = np.zeros_like(bw_region)
-                roi[ry0:ry1, rx0:rx1] = bw_region[ry0:ry1, rx0:rx1]
-                _, cand2 = select_candidates(roi, AREA2)
-                if cand2:
-                    if prev_bbox is None:
-                        best_cnt = max(cand2, key=lambda c: cv2.contourArea(c))
-                    else:
-                        pcx, pcy = _bbox_center(prev_bbox)
-                        best_cnt = max(
-                            cand2,
-                            key=lambda c: -(
-                                (((_cnt_bbox(c)[0] + _cnt_bbox(c)[2]) // 2 - pcx) ** 2)
-                                + (
-                                    ((_cnt_bbox(c)[1] + _cnt_bbox(c)[3]) // 2 - pcy)
-                                    ** 2
-                                )
-                            ),
-                        )
-                else:
-                    # 回退2：全区域 cand，选最近中心
-                    if prev_bbox is not None:
-                        _, cand3 = select_candidates(bw_region, AREA2)
-                        if cand3:
-                            pcx, pcy = _bbox_center(prev_bbox)
-                            best_cnt = max(
-                                cand3,
-                                key=lambda c: -(
-                                    (
-                                        ((_cnt_bbox(c)[0] + _cnt_bbox(c)[2]) // 2 - pcx)
-                                        ** 2
-                                    )
-                                    + (
-                                        ((_cnt_bbox(c)[1] + _cnt_bbox(c)[3]) // 2 - pcy)
-                                        ** 2
-                                    )
-                                ),
-                            )
-
-        region_boxes = (r_top_left, r_bot_left, r_mid_band, (y0, y1))
-        return bw_region, best_cnt, contours_region, region_boxes, prev_roi_box
-
-    # ---- 时序追踪状态（用字典避免 nonlocal/global） ----
-    state = {"bbox": None}  # 保存上一帧外接框 (x0,y0,x1,y1)
-
-    def process_and_show(frame, idx):
-        img = frame.copy()
-        bw, best, contours, region_boxes, prev_roi_box = detect_flower_like(
-            img, state["bbox"]
-        )
-        r_top_left, r_bot_left, r_mid_band, (y0, y1) = region_boxes
-
-        # 所有轮廓（黄）
-        allc = img.copy()
-        if contours:
-            cv2.drawContours(allc, contours, -1, (0, 255, 255), 1)
-
-        # 画三个区域：红框 + 中间带上下红线
-        def draw_rect(im, rect, color=(0, 0, 255), th=2):
-            x0, y0r, x1, y1r = rect
-            cv2.rectangle(im, (x0, y0r), (x1, y1r), color, th)
-
-        draw_rect(allc, r_top_left)
-        draw_rect(allc, r_bot_left)
-        draw_rect(allc, (r_mid_band[0], r_mid_band[1], r_mid_band[2], r_mid_band[3]))
-        cv2.line(allc, (0, y0), (img.shape[1], y0), (0, 0, 255), 2)
-        cv2.line(allc, (0, y1), (img.shape[1], y1), (0, 0, 255), 2)
-
-        # 画上一帧的膨胀 ROI（青色）
-        if prev_roi_box is not None:
-            x0, y0r, x1, y1r = prev_roi_box
-            cv2.rectangle(allc, (x0, y0r), (x1, y1r), (255, 255, 0), 2)
-
-        # 最终检测
-        vis = img.copy()
-        title = "no-detect"
-        if best is not None:
-            cv2.drawContours(vis, [best], -1, (0, 255, 0), 2)
-            x0, y0r, x1, y1r = _cnt_bbox(best)
-            state["bbox"] = (x0, y0r, x1, y1r)  # 更新追踪状态
-            M = cv2.moments(best)
-            if M["m00"] > 0:
-                cx, cy = int(M["m10"] / M["m00"]), int(M["m01"] / M["m00"])
-                cv2.circle(vis, (cx, cy), 4, (0, 0, 255), -1)
-            title = "detected"
-        else:
-            # 若仍未检测，维持上一状态
-            cv2.putText(
-                vis,
-                "No detection (kept last state)",
-                (12, 28),
-                cv2.FONT_HERSHEY_SIMPLEX,
-                0.8,
-                (0, 0, 255),
-                2,
-            )
-            if state["bbox"] is not None:
-                x0, y0r, x1, y1r = state["bbox"]
-                cv2.rectangle(vis, (x0, y0r), (x1, y1r), (255, 255, 0), 2)
-
-        # 四联画：原图 | 区域内bw | 所有轮廓 | 最终检测
-        panel = np.hstack([img, cv2.cvtColor(bw, cv2.COLOR_GRAY2BGR), allc, vis])
-        cv2.putText(
-            panel,
-            f"Frame {idx} | {title}",
-            (12, 28),
-            cv2.FONT_HERSHEY_SIMPLEX,
-            0.9,
-            (255, 255, 255),
-            2,
-        )
-
-        cv2.imshow(window, panel)
-        if writer is not None:
-            if panel.shape[:2] != (H, W * 4):
-                panel = cv2.resize(panel, (W * 4, H), interpolation=cv2.INTER_AREA)
-            writer.write(panel)
-
-    # 先处理已取出的第一帧
-    process_and_show(first_frame, 0)
-
-    # 按你的遍历方式继续
-    for idx, frame in enumerate(
-        tqdm(cleaner.input_video_loader, desc="Processing frames", initial=1, unit="f")
-    ):
-        process_and_show(frame, idx)
-        key = cv2.waitKey(max(1, int(1000 / max(1, int(fps))))) & 0xFF
-        if key == ord("q"):
-            break
-        elif key == ord(" "):
-            while True:
-                k = cv2.waitKey(50) & 0xFF
-                if k in (ord(" "), ord("q")):
-                    if k == ord("q"):
-                        idx = 10**9
-                    break
-            if idx >= 10**9:
-                break
-
-    if writer is not None:
-        writer.release()
-        print(f"[OK] 可视化视频已保存: {out_path}")
-
-    cv2.destroyAllWindows()

uv.lock

@@ -2661,6 +2661,7 @@ dependencies = [
     { name = "fastapi" },
     { name = "ffmpeg-python" },
     { name = "fire" },
+	{ name = "greenlet" },
     { name = "httpx" },
     { name = "huggingface-hub" },
     { name = "jupyter" },
@@ -2695,6 +2696,7 @@ requires-dist = [
     { name = "fastapi", specifier = "==0.108.0" },
     { name = "ffmpeg-python", specifier = ">=0.2.0" },
     { name = "fire", specifier = ">=0.7.1" },
+	{ name = "greenlet", specifier = ">=3.2.4" },
     { name = "httpx", specifier = ">=0.28.1" },
     { name = "huggingface-hub", specifier = ">=0.35.3" },
     { name = "jupyter", specifier = ">=1.1.1" },

version/model_version.json

@@ -0,0 +1 @@
+{}