Update app.py

app.py

@@ -5,11 +5,20 @@ import os
 from PIL import Image
 import time
 import app_config as cfg
+import threading
 
 
 model = SwinTExCo(weights_path=cfg.ckpt_path)
 
+stop_thread = False
+
+def stop_process():
+    global stop_thread
+    stop_thread = True
+
 def video_colorization(video_path, ref_image, progress=gr.Progress()):
+    global stop_thread
+    
     # Initialize video reader
     video_reader = cv2.VideoCapture(video_path)
     fps = video_reader.get(cv2.CAP_PROP_FPS)
@@ -17,34 +26,107 @@ def video_colorization(video_path, ref_image, progress=gr.Progress()):
     width = int(video_reader.get(cv2.CAP_PROP_FRAME_WIDTH))
     num_frames = int(video_reader.get(cv2.CAP_PROP_FRAME_COUNT))
     
+    if not video_reader.isOpened():
+        gr.Warning("Please upload a valid video.")
+    if ref_image is None:
+        gr.Warning("Please upload a valid reference image.")
+    
     # Initialize reference image
     ref_image = Image.fromarray(ref_image)
     
     # Initialize video writer
-    output_path = os.path.join(os.path.dirname(video_path), os.path.basename(video_path).split('.')[0] + '_colorized.mp4')
+    output_path = os.path.join(os.path.dirname(video_path), os.path.basename(video_path).split('.')[0] + f'_colorized_{time.time_ns()}.mp4')
     video_writer = cv2.VideoWriter(output_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (width, height))
     
-    # Init progress bar
-    
     for colorized_frame, _ in zip(model.predict_video(video_reader, ref_image), progress.tqdm(range(num_frames), desc="Colorizing video", unit="frames")):
-        colorized_frame = cv2.cvtColor(colorized_frame, cv2.COLOR_RGB2BGR)
-        video_writer.write(colorized_frame)
-    
-    # for i in progress.tqdm(range(1000)):
-    #     time.sleep(0.5)
+        if stop_thread:
+            stop_thread = False
+            break
+        else:
+            colorized_frame = cv2.cvtColor(colorized_frame, cv2.COLOR_RGB2BGR)
+            video_writer.write(colorized_frame)
         
     video_writer.release()
     
     return output_path
 
-app = gr.Interface(
-    fn=video_colorization,
-    inputs=[gr.Video(format="mp4", sources="upload", label="Input video (grayscale)", interactive=True), 
-            gr.Image(sources="upload", label="Reference image (color)")],
-    outputs=gr.Video(label="Output video (colorized)"),
-    title=cfg.TITLE,
-    description=cfg.DESCRIPTION
-).queue()
+def image_colorization(image, ref_image):
+    if image is None:
+        gr.Warning("Please upload a valid image.")
+    if ref_image is None:
+        gr.Warning("Please upload a valid reference image.")
+    
+    # Initialize image
+    image = Image.fromarray(image)
+    ref_image = Image.fromarray(ref_image)
+    
+    colorized_image = model.predict_image(image, ref_image)
+    
+    return colorized_image
+
+# app = gr.Interface(
+#     fn=video_colorization,
+#     inputs=[gr.Video(format="mp4", sources="upload", label="Input video (grayscale)", interactive=True), 
+#             gr.Image(sources="upload", label="Reference image (color)")],
+#     outputs=gr.Video(label="Output video (colorized)"),
+#     title=cfg.TITLE,
+#     description=cfg.DESCRIPTION,
+#     allow_flagging='never'
+# )
+
+with gr.Blocks() as app:
+    # Title
+    gr.Markdown(cfg.CONTENT)
+    
+    # Video tab
+    with gr.Tab("📹 Video"):
+        with gr.Row():
+            with gr.Column(scale=1):
+                input_video_comp = gr.Video(format="mp4", sources="upload", label="Input video (grayscale)", interactive=True)
+                ref_image_comp = gr.Image(sources="upload", label="Reference image (color)", height=300)
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        clear_btn = gr.ClearButton(value="Clear input", variant=['secondary'])
+                        clear_btn.add([input_video_comp, ref_image_comp])
+                    with gr.Column(scale=1):
+                        start_btn = gr.Button(value="Start!", variant=['primary'])
+            with gr.Column(scale=1):
+                output_video_comp = gr.Video(label="Output video (colorized)")
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        clear_output_btn = gr.ClearButton(value="Clear output", variant=['secondary'])
+                        clear_output_btn.add([output_video_comp])
+                    with gr.Column(scale=1):
+                        stop_btn = gr.Button(value="Stop!", variant=['stop'])
+    
+            start_event = start_btn.click(video_colorization, inputs=[input_video_comp, ref_image_comp], outputs=[output_video_comp])
+            stop_btn.click(fn=None, cancels=[start_event])
+                
+    # Image tab
+    with gr.Tab("🖼️ Image"):
+        with gr.Row():
+            with gr.Column(scale=1):
+                input_image_comp = gr.Image(sources="upload", label="Input image (grayscale)", height=300)
+                ref_image_comp = gr.Image(sources="upload", label="Reference image (color)", height=300)
+                with gr.Row():
+                    with gr.Column(scale=1):
+                        clear_input_btn = gr.ClearButton(value="Clear input", variant=['secondary'])
+                        clear_input_btn.add([input_image_comp, ref_image_comp])
+                    with gr.Column(scale=1):
+                        start_btn = gr.Button(value="Start!", variant=['primary'])
+            with gr.Column(scale=1):
+                output_image_comp = gr.Image(label="Output image (colorized)", height=300)
+                
+                with gr.Row():
+                    with gr.Column():
+                        clear_output_btn = gr.ClearButton(value="Clear output", variant=['secondary'])
+                        clear_output_btn.add([output_image_comp])
+                    with gr.Column():
+                        stop_btn = gr.Button(value="Stop!", variant=['stop'])
+            
+            start_event = start_btn.click(image_colorization, inputs=[input_image_comp, ref_image_comp], outputs=[output_image_comp])
+            stop_btn.click(fn=None, cancels=[start_event])
 
+    gr.Markdown(cfg.APPENDIX)
 
-app.launch()
+app.launch(auth=('admin', 'admin'))