Update app.py

app.py

@@ -1,6 +1,3 @@
-'''
-    Gradio demo (almost the same code as the one used in Huggingface space)
-'''
 import os, sys
 import cv2
 import time
@@ -9,7 +6,11 @@ import gradio as gr
 import torch
 import numpy as np
 from torchvision.utils import save_image
-
+import json
+import threading
+from queue import Queue
+from pathlib import Path
+import shutil
 
 # Import files from the local folder
 root_path = os.path.abspath('.')
@@ -17,6 +18,18 @@ sys.path.append(root_path)
 from test_code.inference import super_resolve_img
 from test_code.test_utils import load_grl, load_rrdb, load_dat
 
+# Global configuration
+OUTPUT_DIR = "outputs"
+HISTORY_FILE = "history.json"
+VIDEO_QUEUE_FILE = "video_queue.json"
+video_queue = Queue()
+processing_status = {}
+
+# Initialize directories
+os.makedirs(OUTPUT_DIR, exist_ok=True)
+os.makedirs(os.path.join(OUTPUT_DIR, "images"), exist_ok=True)
+os.makedirs(os.path.join(OUTPUT_DIR, "videos"), exist_ok=True)
+
 
 def auto_download_if_needed(weight_path):
     if os.path.exists(weight_path):
@@ -40,62 +53,254 @@ def auto_download_if_needed(weight_path):
     if weight_path == "pretrained/4x_APISR_DAT_GAN_generator.pth":
         os.system("wget https://github.com/Kiteretsu77/APISR/releases/download/v0.3.0/4x_APISR_DAT_GAN_generator.pth")
         os.system("mv 4x_APISR_DAT_GAN_generator.pth pretrained")
+
+
+def load_history():
+    """Load processing history from JSON file"""
+    if os.path.exists(HISTORY_FILE):
+        with open(HISTORY_FILE, 'r') as f:
+            return json.load(f)
+    return []
+
+
+def save_history(history):
+    """Save processing history to JSON file"""
+    with open(HISTORY_FILE, 'w') as f:
+        json.dump(history, f, indent=2)
+
+
+def add_to_history(input_path, output_path, model_name, process_type, status="completed"):
+    """Add a record to history"""
+    history = load_history()
+    record = {
+        "timestamp": datetime.datetime.now().isoformat(),
+        "input_path": input_path,
+        "output_path": output_path,
+        "model_name": model_name,
+        "process_type": process_type,
+        "status": status
+    }
+    history.insert(0, record)  # Add to beginning
+    save_history(history)
+
+
+def load_generator(model_name):
+    """Load the appropriate model"""
+    if model_name == "4xGRL":
+        weight_path = "pretrained/4x_APISR_GRL_GAN_generator.pth"
+        auto_download_if_needed(weight_path)
+        generator = load_grl(weight_path, scale=4)
+        
+    elif model_name == "4xRRDB":
+        weight_path = "pretrained/4x_APISR_RRDB_GAN_generator.pth"
+        auto_download_if_needed(weight_path)
+        generator = load_rrdb(weight_path, scale=4)
+        
+    elif model_name == "2xRRDB":
+        weight_path = "pretrained/2x_APISR_RRDB_GAN_generator.pth"
+        auto_download_if_needed(weight_path)
+        generator = load_rrdb(weight_path, scale=2)
+        
+    elif model_name == "4xDAT":
+        weight_path = "pretrained/4x_APISR_DAT_GAN_generator.pth"
+        auto_download_if_needed(weight_path)
+        generator = load_dat(weight_path, scale=4)
+    else:
+        raise ValueError(f"Model {model_name} not supported")
     
+    return generator.to(device='cpu')
 
 
-def inference(img_path, model_name):
+def inference_image(img_path, model_name):
+    """Process a single image"""
+    try:
+        generator = load_generator(model_name)
+        
+        print("Processing image:", img_path)
+        print("Time:", datetime.datetime.now(pytz.timezone('US/Eastern')))
+
+        # Process image
+        super_resolved_img = super_resolve_img(
+            generator, img_path, output_path=None, 
+            downsample_threshold=720, crop_for_4x=True
+        )
+        
+        # Save output
+        timestamp = int(time.time() * 1000)
+        output_name = f"image_{timestamp}.png"
+        output_path = os.path.join(OUTPUT_DIR, "images", output_name)
+        save_image(super_resolved_img, output_path)
+        
+        # Load and convert for display
+        outputs = cv2.imread(output_path)
+        outputs = cv2.cvtColor(outputs, cv2.COLOR_BGR2RGB)
+        
+        # Add to history
+        add_to_history(img_path, output_path, model_name, "image")
+        
+        return outputs, f"✅ Saved to: {output_path}"
     
+    except Exception as error:
+        raise gr.Error(f"Error processing image: {error}")
+
+
+def process_video_frame_by_frame(video_path, model_name, task_id):
+    """Process video frame by frame"""
     try:
-        weight_dtype = torch.float32
+        processing_status[task_id] = {"status": "processing", "progress": 0}
         
-        # Load the model
-        if model_name == "4xGRL":
-            weight_path = "pretrained/4x_APISR_GRL_GAN_generator.pth"
-            auto_download_if_needed(weight_path)
-            generator = load_grl(weight_path, scale=4)  # Directly use default way now
+        # Load model
+        generator = load_generator(model_name)
+        
+        # Open video
+        cap = cv2.VideoCapture(video_path)
+        if not cap.isOpened():
+            raise ValueError("Cannot open video file")
+        
+        # Get video properties
+        fps = cap.get(cv2.CAP_PROP_FPS)
+        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+        width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        
+        # Prepare output
+        timestamp = int(time.time() * 1000)
+        output_name = f"video_{timestamp}.mp4"
+        output_path = os.path.join(OUTPUT_DIR, "videos", output_name)
+        
+        # Create temporary directory for frames
+        temp_dir = f"temp_frames_{timestamp}"
+        os.makedirs(temp_dir, exist_ok=True)
+        
+        # Process frames
+        frame_count = 0
+        while True:
+            ret, frame = cap.read()
+            if not ret:
+                break
+            
+            # Save frame temporarily
+            temp_frame_path = os.path.join(temp_dir, f"frame_{frame_count:06d}.png")
+            cv2.imwrite(temp_frame_path, frame)
             
-        elif model_name == "4xRRDB":
-            weight_path = "pretrained/4x_APISR_RRDB_GAN_generator.pth"
-            auto_download_if_needed(weight_path)
-            generator = load_rrdb(weight_path, scale=4)  # Directly use default way now
+            # Super resolve frame
+            super_resolved_img = super_resolve_img(
+                generator, temp_frame_path, output_path=None,
+                downsample_threshold=720, crop_for_4x=True
+            )
             
-        elif model_name == "2xRRDB":
-            weight_path = "pretrained/2x_APISR_RRDB_GAN_generator.pth"
-            auto_download_if_needed(weight_path)
-            generator = load_rrdb(weight_path, scale=2) # Directly use default way now
+            # Save processed frame
+            output_frame_path = os.path.join(temp_dir, f"output_{frame_count:06d}.png")
+            save_image(super_resolved_img, output_frame_path)
             
-        elif model_name == "4xDAT":
-            weight_path = "pretrained/4x_APISR_DAT_GAN_generator.pth"
-            auto_download_if_needed(weight_path)
-            generator = load_dat(weight_path, scale=4) # Directly use default way now
+            frame_count += 1
+            progress = int((frame_count / total_frames) * 100)
+            processing_status[task_id] = {"status": "processing", "progress": progress}
             
-        else:
-            raise gr.Error("We don't support such Model")
+            print(f"Task {task_id}: Processed frame {frame_count}/{total_frames} ({progress}%)")
+        
+        cap.release()
         
-        generator = generator.to(dtype=weight_dtype)
+        # Combine frames into video using ffmpeg
+        print(f"Task {task_id}: Combining frames into video...")
+        processing_status[task_id] = {"status": "encoding", "progress": 100}
+        
+        os.system(f"ffmpeg -framerate {fps} -i {temp_dir}/output_%06d.png -c:v libx264 -pix_fmt yuv420p {output_path}")
+        
+        # Clean up
+        shutil.rmtree(temp_dir)
+        
+        processing_status[task_id] = {"status": "completed", "progress": 100, "output": output_path}
+        add_to_history(video_path, output_path, model_name, "video")
+        
+        print(f"Task {task_id}: Completed! Output: {output_path}")
+        
+    except Exception as error:
+        processing_status[task_id] = {"status": "error", "error": str(error)}
+        print(f"Task {task_id}: Error - {error}")
 
 
-        print("We are processing ", img_path)
-        print("The time now is ", datetime.datetime.now(pytz.timezone('US/Eastern')))
+def video_queue_worker():
+    """Background worker to process video queue"""
+    print("Video queue worker started...")
+    while True:
+        try:
+            task = video_queue.get()
+            if task is None:  # Poison pill to stop worker
+                break
+            
+            task_id, video_path, model_name = task
+            print(f"Starting task {task_id}...")
+            process_video_frame_by_frame(video_path, model_name, task_id)
+            
+        except Exception as e:
+            print(f"Worker error: {e}")
+        finally:
+            video_queue.task_done()
+
 
-        # In default, we will automatically use crop to match 4x size
-        super_resolved_img = super_resolve_img(generator, img_path, output_path=None, weight_dtype=weight_dtype, downsample_threshold=720, crop_for_4x=True)
-        store_name = str(time.time()) + ".png"
-        save_image(super_resolved_img, store_name)
-        outputs = cv2.imread(store_name)
-        outputs = cv2.cvtColor(outputs, cv2.COLOR_RGB2BGR)
-        os.remove(store_name)
-        
-        return outputs
+def submit_video(video_path, model_name):
+    """Submit video to processing queue"""
+    if video_path is None:
+        return None, "❌ Please upload a video first"
     
+    task_id = f"task_{int(time.time() * 1000)}"
+    video_queue.put((task_id, video_path, model_name))
+    processing_status[task_id] = {"status": "queued", "progress": 0}
     
-    except Exception as error:
-        raise gr.Error(f"global exception: {error}")
+    return None, f"✅ Video submitted to queue! Task ID: {task_id}\nCheck status in the monitoring section."
+
 
+def get_queue_status():
+    """Get current queue status"""
+    status_text = "📊 **Queue Status**\n\n"
+    status_text += f"Videos in queue: {video_queue.qsize()}\n\n"
+    
+    if processing_status:
+        status_text += "**Active Tasks:**\n"
+        for task_id, status in processing_status.items():
+            status_text += f"\n🎬 {task_id}:\n"
+            status_text += f"  Status: {status['status']}\n"
+            status_text += f"  Progress: {status.get('progress', 0)}%\n"
+            if 'output' in status:
+                status_text += f"  Output: {status['output']}\n"
+            if 'error' in status:
+                status_text += f"  Error: {status['error']}\n"
+    else:
+        status_text += "No active tasks"
+    
+    return status_text
+
+
+def get_history_display():
+    """Get formatted history for display"""
+    history = load_history()
+    if not history:
+        return "No history available"
+    
+    history_text = "📜 **Processing History**\n\n"
+    for idx, record in enumerate(history[:50]):  # Show last 50
+        history_text += f"**{idx + 1}. {record['process_type'].upper()}** - {record['timestamp']}\n"
+        history_text += f"   Model: {record['model_name']}\n"
+        history_text += f"   Status: {record['status']}\n"
+        history_text += f"   Output: {record['output_path']}\n\n"
+    
+    return history_text
+
+
+def clear_history():
+    """Clear all history"""
+    if os.path.exists(HISTORY_FILE):
+        os.remove(HISTORY_FILE)
+    return "✅ History cleared!"
 
 
 if __name__ == '__main__':
     
+    # Start background worker thread
+    worker_thread = threading.Thread(target=video_queue_worker, daemon=True)
+    worker_thread.start()
+    
     MARKDOWN = \
     """
     ## <p style='text-align: center'> APISR: Anime Production Inspired Real-World Anime Super-Resolution (CVPR 2024) </p>
@@ -104,50 +309,104 @@ if __name__ == '__main__':
 
     APISR aims at restoring and enhancing low-quality low-resolution **anime** images and video sources with various degradations from real-world scenarios.
     
-    ### Note: Due to memory restriction, all images whose short side is over 720 pixel will be downsampled to 720 pixel with the same aspect ratio.  E.g., 1920x1080 -> 1280x720
-    ### Note: Please check [Model Zoo](https://github.com/Kiteretsu77/APISR/blob/main/docs/model_zoo.md) for the description of each weight and [Here](https://imgsli.com/MjU0MjI0) for model comparisons.
-    
-    ### If APISR is helpful, please help star the [GitHub Repo](https://github.com/Kiteretsu77/APISR). Thanks! ###
+    ### ⚠️ Note: Images with short side > 720px will be downsampled to 720px (e.g., 1920x1080 → 1280x720)
+    ### 📹 New: Video processing runs in background queue - you can close the browser and it continues!
     """
 
     block = gr.Blocks().queue(max_size=10)
     with block:
-        with gr.Row():
-            gr.Markdown(MARKDOWN)
-        with gr.Row(elem_classes=["container"]):
-            with gr.Column(scale=2):
-                input_image = gr.Image(type="filepath", label="Input")
-                model_name = gr.Dropdown(
-                    [
-                        "2xRRDB",
-                        "4xRRDB",
-                        "4xGRL",
-                        "4xDAT",
-                    ],
-                    type="value",
-                    value="4xGRL",
-                    label="model",
+        gr.Markdown(MARKDOWN)
+        
+        with gr.Tabs():
+            # Tab 1: Image Processing
+            with gr.Tab("🖼️ Image Processing"):
+                with gr.Row():
+                    with gr.Column(scale=2):
+                        input_image = gr.Image(type="filepath", label="Input Image")
+                        image_model = gr.Dropdown(
+                            ["2xRRDB", "4xRRDB", "4xGRL", "4xDAT"],
+                            value="4xGRL",
+                            label="Model"
+                        )
+                        image_btn = gr.Button("🚀 Process Image", variant="primary")
+                    
+                    with gr.Column(scale=3):
+                        output_image = gr.Image(type="numpy", label="Output Image")
+                        image_status = gr.Textbox(label="Status", lines=2)
+                
+                with gr.Row():
+                    gr.Examples(
+                        [
+                            ["__assets__/lr_inputs/image-00277.png"],
+                            ["__assets__/lr_inputs/image-00542.png"],
+                            ["__assets__/lr_inputs/41.png"],
+                            ["__assets__/lr_inputs/f91.jpg"],
+                        ],
+                        [input_image],
+                    )
+                
+                image_btn.click(
+                    inference_image, 
+                    inputs=[input_image, image_model], 
+                    outputs=[output_image, image_status]
                 )
-                run_btn = gr.Button(value="Submit")
-
-            with gr.Column(scale=3):
-                output_image = gr.Image(type="numpy", label="Output image")
-
-        with gr.Row(elem_classes=["container"]):
-            gr.Examples(
-                [
-                    ["__assets__/lr_inputs/image-00277.png"],
-                    ["__assets__/lr_inputs/image-00542.png"],
-                    ["__assets__/lr_inputs/41.png"],
-                    ["__assets__/lr_inputs/f91.jpg"],
-                    ["__assets__/lr_inputs/image-00440.png"],
-                    ["__assets__/lr_inputs/image-00164.jpg"],
-                    ["__assets__/lr_inputs/img_eva.jpeg"],
-                    ["__assets__/lr_inputs/naruto.jpg"],
-                ],
-                [input_image],
-            )
-
-        run_btn.click(inference, inputs=[input_image, model_name], outputs=[output_image])
+            
+            # Tab 2: Video Processing
+            with gr.Tab("🎬 Video Processing"):
+                gr.Markdown("""
+                ### Video Processing Queue
+                Videos are processed in the background. You can submit multiple videos and close the browser - processing continues!
+                """)
+                
+                with gr.Row():
+                    with gr.Column():
+                        input_video = gr.Video(label="Input Video")
+                        video_model = gr.Dropdown(
+                            ["2xRRDB", "4xRRDB", "4xGRL", "4xDAT"],
+                            value="4xGRL",
+                            label="Model"
+                        )
+                        video_btn = gr.Button("📤 Submit to Queue", variant="primary")
+                        video_status = gr.Textbox(label="Submission Status", lines=3)
+                    
+                    with gr.Column():
+                        gr.Markdown("### 📊 Queue Monitor")
+                        queue_status = gr.Textbox(label="Queue Status", lines=15)
+                        refresh_btn = gr.Button("🔄 Refresh Status")
+                
+                video_btn.click(
+                    submit_video,
+                    inputs=[input_video, video_model],
+                    outputs=[input_video, video_status]
+                )
+                
+                refresh_btn.click(
+                    get_queue_status,
+                    outputs=[queue_status]
+                )
+            
+            # Tab 3: History
+            with gr.Tab("📜 History"):
+                gr.Markdown("### Processing History")
+                
+                with gr.Row():
+                    refresh_history_btn = gr.Button("🔄 Refresh History")
+                    clear_history_btn = gr.Button("🗑️ Clear History", variant="stop")
+                
+                history_display = gr.Textbox(label="History", lines=20)
+                clear_status = gr.Textbox(label="Status", lines=1)
+                
+                refresh_history_btn.click(
+                    get_history_display,
+                    outputs=[history_display]
+                )
+                
+                clear_history_btn.click(
+                    clear_history,
+                    outputs=[clear_status]
+                )
+                
+                # Auto-load history on tab open
+                block.load(get_history_display, outputs=[history_display])
 
-    block.launch()
+    block.launch(server_name="0.0.0.0", server_port=7860)