Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import os
 import cv2
-from modelscope.outputs import OutputKeys
+# Assuming img_colorization as before
 from modelscope.pipelines import pipeline
 from modelscope.utils.constant import Tasks
 import PIL
@@ -9,26 +9,49 @@ import numpy as np
 import uuid
 from gradio_imageslider import ImageSlider
 
+# Your existing colorization pipeline
 img_colorization = pipeline(Tasks.image_colorization, model='iic/cv_ddcolor_image-colorization')
-img_path = 'input.png'
-##result = img_colorization(img_path)
-##cv2.imwrite('result.png', result[OutputKeys.OUTPUT_IMG])
-def color(image):
-    output = img_colorization(image[...,::-1])
-    result = output[OutputKeys.OUTPUT_IMG].astype(np.uint8)
-    # result = result[...,::-1]
-    # Generate a unique filename using UUID
+
+# Additional pipelines for image enhancement and repair
+img_enhancement = pipeline(Tasks.image_super_resolution, model='your_super_resolution_model')
+img_repair = pipeline(Tasks.image_inpainting, model='your_image_repair_model')
+
+def process_image(image, enhance_quality, repair_damage):
+    # Convert image to proper format for processing
+    image = image[..., ::-1]
+    
+    # Repair the image if requested
+    if repair_damage:
+        image = img_repair(image)[OutputKeys.OUTPUT_IMG].astype(np.uint8)
+    
+    # Enhance image quality if requested
+    if enhance_quality:
+        image = img_enhancement(image)[OutputKeys.OUTPUT_IMG].astype(np.uint8)
+    
+    # Colorize the image
+    colored_image = img_colorization(image)[OutputKeys.OUTPUT_IMG].astype(np.uint8)
+    
+    # Save the processed image with a unique filename
     unique_imgfilename = str(uuid.uuid4()) + '.png'
-    cv2.imwrite(unique_imgfilename, result)
-    print('infer finished!')
-    return (image, unique_imgfilename)
+    cv2.imwrite(unique_imgfilename, colored_image)
+    print('Infer finished!')
     
+    # Return both original and processed image paths for the slider
+    return image[..., ::-1], unique_imgfilename
 
-title = "old_photo_restoration"
-description = "upload old photo, ddcolor image colorization"
+title = "Old Photo Restoration"
+description = "Upload old photos for colorization, quality enhancement, and damage repair."
 examples = [['./input.jpg'],]
 
-demo = gr.Interface(fn=color,inputs="image",outputs=ImageSlider(position=0.5,label='Colored image with slider-view'),examples=examples,title=title,description=description)
+inputs = [
+    gr.inputs.Image(shape=(512, 512)),
+    gr.inputs.Checkbox(label="Enhance Quality"),
+    gr.inputs.Checkbox(label="Repair Damage")
+]
+
+outputs = gr.outputs.ImageSlider(position=0.5, label='Processed image with slider-view')
+
+demo = gr.Interface(fn=process_image, inputs=inputs, outputs=outputs, examples=examples, title=title, description=description)
 
 if __name__ == "__main__":
-    demo.launch(share=False)
+    demo.launch()