dkescape commited on
Commit
d8668a4
·
verified ·
1 Parent(s): cda6b5b

Create app.py

Browse files
Files changed (1) hide show
  1. app.py +64 -0
app.py ADDED
@@ -0,0 +1,64 @@
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
+ import os
2
+ import uuid
3
+ import cv2
4
+ import numpy as np
5
+ import gradio as gr
6
+ from modelscope.pipelines import pipeline
7
+ from modelscope.utils.constant import Tasks
8
+ from modelscope.outputs import OutputKeys
9
+ from gradio_imageslider import ImageSlider
10
+
11
+ # Initialize the image colorization pipeline with a specified model revision
12
+ img_colorization = pipeline(
13
+ task=Tasks.image_colorization,
14
+ model='iic/cv_ddcolor_image-colorization',
15
+ model_revision='v1.02' # Specify the model revision to avoid warnings
16
+ )
17
+
18
+ def colorize_image(input_image):
19
+ """
20
+ Colorizes a grayscale image using the DDColor model.
21
+
22
+ Args:
23
+ input_image (numpy.ndarray): Grayscale input image.
24
+
25
+ Returns:
26
+ tuple: Original and colorized images.
27
+ """
28
+ # Convert RGB to BGR for OpenCV compatibility
29
+ bgr_image = cv2.cvtColor(input_image, cv2.COLOR_RGB2BGR)
30
+
31
+ # Perform colorization
32
+ output = img_colorization(bgr_image)
33
+ colorized_bgr = output[OutputKeys.OUTPUT_IMG].astype(np.uint8)
34
+
35
+ # Convert BGR back to RGB for display
36
+ colorized_rgb = cv2.cvtColor(colorized_bgr, cv2.COLOR_BGR2RGB)
37
+
38
+ return input_image, colorized_rgb
39
+
40
+ # Define the Gradio interface
41
+ title = "DDColor: Photo-Realistic Image Colorization"
42
+ description = (
43
+ "Upload a black & white image to colorize it using the DDColor model. "
44
+ "This application utilizes a dual-decoder architecture to produce vivid and natural colorizations."
45
+ )
46
+
47
+ # Example images (ensure these paths are correct or adjust accordingly)
48
+ examples = [['./input.jpg']]
49
+
50
+ # Create the Gradio interface
51
+ demo = gr.Interface(
52
+ fn=colorize_image,
53
+ inputs=gr.Image(label="Upload Grayscale Image"),
54
+ outputs=ImageSlider(position=0.5, label='Colorization Comparison'),
55
+ examples=examples,
56
+ title=title,
57
+ description=description,
58
+ allow_flagging="never"
59
+ )
60
+
61
+ if __name__ == "__main__":
62
+ # Launch the Gradio app with production settings
63
+ demo.launch(server_name="0.0.0.0", server_port=7860, share=False)
64
+