Create app.py

app.py

@@ -0,0 +1,191 @@
+import streamlit as st
+import os
+import py7zr
+import requests
+from huggingface_hub import HfApi
+import torch
+from torch.utils.data import DataLoader
+import shutil
+from pathlib import Path
+from typing import Optional
+import sys
+import io
+
+# Import the denoising code (assuming it's in a file called denoising_model.py)
+from denoising_model import DenoisingModel, DenoiseDataset, get_optimal_threads
+
+class StreamCapture:
+    def __init__(self):
+        self.logs = []
+    
+    def write(self, text):
+        self.logs.append(text)
+        st.warning(text)
+    
+    def flush(self):
+        pass
+
+def download_and_extract_7z(url: str, extract_to: str = '.') -> Optional[str]:
+    """Downloads a 7z file and extracts it"""
+    try:
+        st.warning(f"Downloading file from {url}...")
+        response = requests.get(url, stream=True)
+        response.raise_for_status()
+        
+        archive_path = os.path.join(extract_to, 'dataset.7z')
+        with open(archive_path, 'wb') as f:
+            for chunk in response.iter_content(chunk_size=8192):
+                f.write(chunk)
+        
+        st.warning("Extracting 7z archive...")
+        with py7zr.SevenZipFile(archive_path, mode='r') as z:
+            z.extractall(extract_to)
+        
+        # Handle directory renaming
+        output_images_path = os.path.join(extract_to, 'output_images')
+        if os.path.exists(output_images_path):
+            # Move and rename directories
+            source_noisy = os.path.join(output_images_path, 'images_noisy')
+            source_target = os.path.join(output_images_path, 'images_target')
+            
+            if os.path.exists('noisy_images'):
+                shutil.rmtree('noisy_images')
+            if os.path.exists('target_images'):
+                shutil.rmtree('target_images')
+                
+            shutil.move(source_noisy, 'noisy_images')
+            shutil.move(source_target, 'target_images')
+            
+            # Clean up
+            if os.path.exists(output_images_path):
+                shutil.rmtree(output_images_path)
+            
+        os.remove(archive_path)
+        st.warning("Download and extraction completed successfully.")
+        return None
+        
+    except Exception as e:
+        return f"Error: {str(e)}"
+
+def upload_to_huggingface(file_path: str, repo_id: str, path_in_repo: str):
+    """Uploads a file to Hugging Face"""
+    try:
+        api = HfApi()
+        api.upload_file(
+            path_or_fileobj=file_path,
+            path_in_repo=path_in_repo,
+            repo_id=repo_id,
+            repo_type="space"
+        )
+        st.warning(f"Successfully uploaded {file_path} to {repo_id}")
+    except Exception as e:
+        st.error(f"Error uploading to Hugging Face: {str(e)}")
+
+def train_model_with_upload(epochs, batch_size, learning_rate, save_interval, num_workers, repo_id):
+    """Modified training function that uploads checkpoints to Hugging Face"""
+    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+    st.warning(f"Using device: {device}")
+    
+    # Create temporary directory for checkpoints
+    checkpoint_dir = "temp_checkpoints"
+    os.makedirs(checkpoint_dir, exist_ok=True)
+    
+    try:
+        dataset = DenoiseDataset('noisy_images', 'target_images')
+        dataloader = DataLoader(
+            dataset,
+            batch_size=batch_size,
+            shuffle=True,
+            num_workers=num_workers,
+            pin_memory=True if torch.cuda.is_available() else False
+        )
+        
+        model = DenoisingModel().to(device)
+        criterion = torch.nn.MSELoss()
+        optimizer = torch.optim.Adam(model.parameters(), lr=learning_rate)
+        
+        for epoch in range(epochs):
+            st.warning(f"Starting epoch {epoch+1}/{epochs}")
+            for batch_idx, (noisy_patches, target_patches) in enumerate(dataloader):
+                noisy_patches = noisy_patches.to(device)
+                target_patches = target_patches.to(device)
+                
+                outputs = model(noisy_patches)
+                loss = criterion(outputs, target_patches)
+                
+                optimizer.zero_grad()
+                loss.backward()
+                optimizer.step()
+                
+                if (batch_idx + 1) % 10 == 0:
+                    st.warning(f"Epoch [{epoch+1}/{epochs}], Batch [{batch_idx+1}], Loss: {loss.item():.6f}")
+                
+                if (batch_idx + 1) % save_interval == 0:
+                    checkpoint_path = os.path.join(checkpoint_dir, f"checkpoint_epoch{epoch+1}_batch{batch_idx+1}.pth")
+                    torch.save(model.state_dict(), checkpoint_path)
+                    
+                    # Upload checkpoint to Hugging Face
+                    upload_to_huggingface(
+                        checkpoint_path,
+                        repo_id,
+                        f"checkpoints/checkpoint_epoch{epoch+1}_batch{batch_idx+1}.pth"
+                    )
+        
+        # Save and upload final model
+        final_model_path = os.path.join(checkpoint_dir, "final_model.pth")
+        torch.save(model.state_dict(), final_model_path)
+        upload_to_huggingface(final_model_path, repo_id, "model/final_model.pth")
+        
+    finally:
+        # Clean up temporary directory
+        if os.path.exists(checkpoint_dir):
+            shutil.rmtree(checkpoint_dir)
+
+def main():
+    st.title("Image Denoising Model Training")
+    
+    # Redirect stdout to capture print statements
+    sys.stdout = StreamCapture()
+    
+    # Input for Hugging Face token
+    hf_token = st.text_input("Enter your Hugging Face token:", type="password")
+    if hf_token:
+        os.environ["HF_TOKEN"] = hf_token
+    
+    # Input for repository ID
+    repo_id = st.text_input("Enter your Hugging Face repository ID (username/repo):")
+    
+    # Download and extract dataset button
+    if st.button("Download and Extract Dataset"):
+        url = "https://huggingface.co/spaces/vericudebuget/ok4231/resolve/main/output_images.7z"
+        error = download_and_extract_7z(url)
+        if error:
+            st.error(error)
+    
+    # Training parameters
+    col1, col2 = st.columns(2)
+    with col1:
+        epochs = st.number_input("Number of epochs", min_value=1, value=10)
+        batch_size = st.number_input("Batch size", min_value=1, value=4)
+        learning_rate = st.number_input("Learning rate", min_value=0.0001, value=0.001, format="%.4f")
+    
+    with col2:
+        save_interval = st.number_input("Save interval (batches)", min_value=1, value=1000)
+        num_workers = st.number_input("Number of workers", min_value=1, value=get_optimal_threads())
+    
+    # Start training button
+    if st.button("Start Training"):
+        if not hf_token:
+            st.error("Please enter your Hugging Face token")
+            return
+        if not repo_id:
+            st.error("Please enter your repository ID")
+            return
+        if not os.path.exists("noisy_images") or not os.path.exists("target_images"):
+            st.error("Dataset not found. Please download and extract it first.")
+            return
+            
+        train_model_with_upload(epochs, batch_size, learning_rate, save_interval, num_workers, repo_id)
+
+if __name__ == "__main__":
+    main()