File size: 9,675 Bytes
6ead25b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
'''
conda create --name animeins python=3.10
conda activate animeins
pip install ipykernel
python -m ipykernel install --user --name animeins --display-name "animeins"
pip install -r requirements.txt

pip install torch==2.1.1 torchvision
pip install mmcv==2.1.0 -f https://download.openmmlab.com/mmcv/dist/cu121/torch2.1/index.html
pip install mmdet
pip install "numpy<2.0.0"
pip install moviepy==1.0.3
pip install "httpx[socks]"
'''

import gradio as gr
import os
import cv2
import numpy as np
from PIL import Image
from typing import Literal
import pathlib
from animeinsseg import AnimeInsSeg, AnimeInstances
from animeinsseg.anime_instances import get_color

# Install required packages
os.system("mim install mmengine")
os.system('mim install mmcv==2.1.0')
os.system("mim install mmdet==3.2.0")

# Download model if not exists
if not os.path.exists("models"):
    os.mkdir("models")
    os.system("huggingface-cli lfs-enable-largefiles .")
    os.system("git clone https://huggingface.co/dreMaz/AnimeInstanceSegmentation models/AnimeInstanceSegmentation")

# Initialize segmentation model
ckpt = r'models/AnimeInstanceSegmentation/rtmdetl_e60.ckpt'
mask_thres = 0.3
instance_thres = 0.3
refine_kwargs = {'refine_method': 'refinenet_isnet'}
net = AnimeInsSeg(ckpt, mask_thr=mask_thres, refine_kwargs=refine_kwargs)

def image_to_sketch(image: np.ndarray) -> np.ndarray:
    """Convert image to pencil sketch"""
    gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY)
    inverted = 255 - gray
    blurred = cv2.GaussianBlur(inverted, (21, 21), 0)
    inverted_blurred = 255 - blurred
    sketch = cv2.divide(gray, inverted_blurred, scale=256.0)
    return cv2.cvtColor(sketch, cv2.COLOR_GRAY2BGR)  # Return 3-channel image

def generate_segmentation_video(
    original_image: np.ndarray,
    depth_map: np.ndarray,
    render_order: str = 'character_first',
    duration_sec: float = 3.0,
    frame_rate: int = 30,
    depth_blur: int = 15,
    debug_visualize: bool = False
) -> str:
    """
    Generate transition video with different rendering approaches:
    - 'character_first': Segmented instances transition first, then depth-based
    - 'near_to_far': Transition from nearest to farthest based on depth
    - 'far_to_near': Transition from farthest to nearest based on depth
    """
    # Convert images to proper format
    original = cv2.cvtColor(original_image, cv2.COLOR_RGB2BGR)
    depth_map = cv2.cvtColor(depth_map, cv2.COLOR_RGB2GRAY)
    
    # Get sketch version
    sketch = image_to_sketch(original)
    h, w = original.shape[:2]
    
    # Perform instance segmentation
    instances: AnimeInstances = net.infer(
        original,
        output_type='numpy',
        pred_score_thr=instance_thres
    )
    
    # Prepare depth map
    depth_map = cv2.resize(depth_map, (w, h))
    depth_map = cv2.GaussianBlur(depth_map, (depth_blur, depth_blur), 0)
    depth_map = depth_map.astype(np.float32) / 255.0
    
    # Create layer masks and their depths
    layer_masks = []
    layer_depths = []
    
    # Process segmented instances (for all modes)
    if instances.bboxes is not None:
        for mask in instances.masks:
            # Calculate average depth for this instance
            instance_depth = np.mean(depth_map[mask.astype(bool)])
            
            if render_order == 'character_first':
                # For character-first mode, we'll process characters separately
                layer_masks.append(mask.astype(np.float32))
                layer_depths.append(0)  # Depth doesn't matter for character-first
            else:
                # For depth-based modes, use the actual depth
                if render_order == 'near_to_far':
                    instance_depth = 1.0 - instance_depth
                layer_masks.append(mask.astype(np.float32))
                layer_depths.append(instance_depth)
    
    # Create a full mask for the remaining areas
    if layer_masks:
        full_mask = 1.0 - np.clip(np.sum(layer_masks, axis=0), 0, 1)
    else:
        full_mask = np.ones((h, w), dtype=np.float32)
    
    # Process remaining areas based on the selected mode
    if render_order == 'character_first':
        # For character-first mode, add the remaining areas as one layer
        if np.sum(full_mask) > 0:
            layer_masks.append(full_mask)
            layer_depths.append(1)  # Background comes last
    else:
        # For depth-based modes, divide remaining areas into depth bands
        remaining_depth = depth_map * full_mask
        num_depth_bands = 10  # Number of depth bands for non-segmented areas
        
        min_depth = np.min(remaining_depth[full_mask > 0]) if np.sum(full_mask) > 0 else 0
        max_depth = np.max(remaining_depth[full_mask > 0]) if np.sum(full_mask) > 0 else 1
        depth_bands = np.linspace(min_depth, max_depth, num_depth_bands + 1)
        
        for i in range(num_depth_bands):
            lower = depth_bands[i]
            upper = depth_bands[i+1]
            band_mask = ((remaining_depth >= lower) & (remaining_depth < upper)).astype(np.float32)
            
            if np.sum(band_mask) > 0:
                band_depth = np.mean(remaining_depth[band_mask.astype(bool)])
                if render_order == 'near_to_far':
                    band_depth = 1.0 - band_depth
                
                layer_masks.append(band_mask)
                layer_depths.append(band_depth)
    
    # Sort layers based on the selected mode
    if render_order == 'character_first':
        # Characters first, then background
        pass  # Already in correct order
    else:
        # Sort by depth for depth-based modes
        if layer_masks:
            sorted_indices = np.argsort(layer_depths)
            layer_masks = [layer_masks[i] for i in sorted_indices]
    
    # Generate video
    output_path = "output_video.mp4"
    fourcc = cv2.VideoWriter_fourcc(*'mp4v')
    video = cv2.VideoWriter(output_path, fourcc, frame_rate, (w, h))
    total_frames = int(duration_sec * frame_rate)
    num_layers = len(layer_masks)
    layer_duration = duration_sec / num_layers if num_layers > 0 else duration_sec
    
    for frame_idx in range(total_frames):
        current_time = frame_idx / frame_rate
        blended = original.copy().astype(np.float32)
        
        for layer_idx, layer_mask in enumerate(layer_masks):
            # Calculate current layer progress
            layer_start = layer_idx * layer_duration
            layer_progress = np.clip((current_time - layer_start) / layer_duration, 0, 1)
            
            # Generate blending mask
            layer_alpha = layer_mask * (1 - layer_progress)
            layer_alpha = np.repeat(layer_alpha[..., np.newaxis], 3, axis=2)
            
            # Blend with sketch
            blended = blended * (1 - layer_alpha) + sketch.astype(np.float32) * layer_alpha
        
        blended = np.clip(blended, 0, 255).astype(np.uint8)
        
        if debug_visualize:
            cv2.imshow('Blended', blended)
            if cv2.waitKey(1) == 27:
                break
        
        video.write(blended)
    
    video.release()
    if debug_visualize:
        cv2.destroyAllWindows()
    
    return output_path

def process_images(original_image, depth_map, render_order, duration):
    # Convert PIL Images to numpy arrays
    original_np = np.array(original_image)
    depth_np = np.array(depth_map)
    
    # Generate video
    video_path = generate_segmentation_video(
        original_image=original_np,
        depth_map=depth_np,
        render_order=render_order,
        duration_sec=float(duration),
        debug_visualize=False
    )
    
    return video_path

# Prepare example images
genshin_impact_exps = []
if os.path.exists("Genshin_Impact_Images"):
    genshin_impact_exps = list(map(str, pathlib.Path("Genshin_Impact_Images").rglob("*.png")))

# Create Gradio interface
with gr.Blocks() as demo:
    gr.Markdown("# Colored Anime Image from Sketch to Color Video Generator")
    gr.Markdown("Upload an image and its depth map to generate a depth-aware transition video from sketch to original.")
    
    with gr.Row():
        with gr.Column():
            original_image = gr.Image(label="Original Image", type="pil")
            depth_map = gr.Image(label="Depth Map", type="pil")
            render_order = gr.Radio(
                choices=["character_first", "near_to_far", "far_to_near"],
                value="character_first",
                label="Render Order",
                info="'character_first' shows characters first, others are depth-based"
            )
            duration = gr.Slider(1, 10, value=3, step=0.5, label="Duration (seconds)")
            submit_btn = gr.Button("Generate Video")
        
        with gr.Column():
            output_video = gr.Video(label="Output Video")
       
    submit_btn.click(
        fn=process_images,
        inputs=[original_image, depth_map, render_order, duration],
        outputs=output_video
    )

    # Add examples if available 
    gr.Examples(
        [
            ["化物语封面.jpeg", "化物语封面深度.png", "character_first",],
            ["化物语封面.jpeg", "化物语封面深度.png", "far_to_near",],
            ["可莉风景.png", "可莉风景_depth.png", "near_to_far",],
            ["竹林万叶.jpg", "竹林万叶_depth.png", "character_first",],
            ["竹林万叶.jpg", "竹林万叶_depth.png", "near_to_far",],
            ["重云行秋.jpg", "重云行秋_depth.png", "character_first",],
        ],
        inputs = [original_image, depth_map, render_order]
    )

if __name__ == "__main__":
    demo.launch(share=True)