Upload DeFogify_Main.py

DeFogify_Main.py

@@ -1,7 +1,11 @@
 import cv2
 import numpy as np
 import gradio as gr
+import tempfile
+import os
+from tqdm import tqdm
 
+# Original Functions
 def dark_channel(img, size=15):
     r, g, b = cv2.split(img)
     min_img = cv2.min(r, cv2.min(g, b))
@@ -36,26 +40,123 @@ def guided_filter(p, i, r, e):
 def dehaze(image):
     img = image.astype('float64') / 255
     img_gray = cv2.cvtColor(image, cv2.COLOR_BGR2GRAY).astype('float64') / 255
-
     atom = get_atmo(img)
     trans = get_trans(img, atom)
     trans_guided = guided_filter(trans, img_gray, 20, 0.0001)
     trans_guided = np.maximum(trans_guided, 0.25)  # Ensure trans_guided is not below 0.25
-
     result = np.empty_like(img)
     for i in range(3):
         result[:, :, i] = (img[:, :, i] - atom) / trans_guided + atom
-
-    # Ensure the result is in the range [0, 1]
     result = np.clip(result, 0, 1)
     return (result * 255).astype(np.uint8)
 
-# Save example images for testing
+# Single Image Processing
+def process_single_image(image):
+    dehazed_img = dehaze(image)
+    return dehazed_img
+
+# Batch Processing Function for Multiple Images with Progress Bar
+def process_images(files):
+    temp_dir = tempfile.mkdtemp()
+    output_files = []
+    
+    for file in tqdm(files, desc="Processing Images"):
+        img = cv2.imread(file.name)
+        if img is not None:
+            dehazed_img = dehaze(img)
+            output_path = os.path.join(temp_dir, os.path.basename(file.name))
+            cv2.imwrite(output_path, dehazed_img)
+            output_files.append(output_path)
+    
+    return output_files
+
+# Video Dehazing Function with Gradio Progress Bar and Error Handling
+def dehaze_video(input_video_path, output_video_path, progress=None):
+    try:
+        cap = cv2.VideoCapture(input_video_path)
+        if not cap.isOpened():
+            raise ValueError("Error: Could not open video.")
+
+        fourcc = cv2.VideoWriter_fourcc(*'mp4v')
+        fps = int(cap.get(cv2.CAP_PROP_FPS))
+        frame_width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+        frame_height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+        total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
+
+        if total_frames <= 0:  # Assume a constant count for webcam scenarios
+            total_frames = 1000
+
+        out = cv2.VideoWriter(output_video_path, fourcc, fps, (frame_width, frame_height))
+        frame_count = 0
+
+        if progress is not None:
+            progress(0, desc="Processing Video", unit="frame")
+
+        while cap.isOpened():
+            ret, frame = cap.read()
+            if not ret:
+                break
+            dehazed_frame = dehaze(frame)
+            out.write(dehazed_frame)
+            frame_count += 1
+
+            if progress is not None:
+                progress(frame_count / total_frames)  # Ensure progress is within 0-1 range
+
+        cap.release()
+        out.release()
+        print(f"\nDehazed video saved to: {output_video_path}")
+    except Exception as e:
+        print(f"An error occurred during video processing: {e}")
+
+# Gradio Video Processing Wrapper
+def process_video(file):
+    input_video_path = file  # File is a string representing the path
+    output_video_path = os.path.join(tempfile.mkdtemp(), "dehazed_video.mp4")
+    progress = gr.Progress()
+    dehaze_video(input_video_path, output_video_path, progress)
+    return output_video_path
+
+# Real-Time Webcam Processing with Gradio Progress Bar
+def dehaze_webcam(progress=gr.Progress()):
+    try:
+        cap = cv2.VideoCapture(0)  # Capture from the first webcam
+        if not cap.isOpened():
+            raise ValueError("Unable to open webcam")
+
+        frame_count = 0
+        total_frames = 100  # Arbitrary number for progress bar
+        progress(0, desc="Processing Webcam Feed", unit="frame")
+
+        while frame_count < total_frames:
+            ret, frame = cap.read()
+            if not ret:
+                break
+            dehazed_frame = dehaze(frame)
+            frame_count += 1
+            progress(frame_count / total_frames)  # Ensure progress is within 0-1 range
+
+            cv2.imshow('Dehazed Webcam Feed', dehazed_frame)
+            if cv2.waitKey(1) & 0xFF == ord('q'):
+                break
+
+        cap.release()
+        cv2.destroyAllWindows()
+        progress(1)  # Ensure progress bar reaches 100%
+    except Exception as e:
+        print(f"An error occurred during webcam processing: {e}")
+
+# Gradio Webcam Processing Wrapper
+def process_webcam():
+    progress = gr.Progress()
+    dehaze_webcam(progress)
+    return "Webcam processing completed."
+
+# Example Images for Testing
 example_images = [
     "Sample Images for Testing/ai-generated-9025430_1280.jpg",
     "Sample Images for Testing/meadow-5648849_1280.jpg",
     "Sample Images for Testing/mountains-7662717_1280.jpg",
-    "Sample Images for Testing/mountains-8292685_1280.jpg",
     "Sample Images for Testing/nature-6722031_1280.jpg"
 ]
 
@@ -66,13 +167,37 @@ for i, img_path in enumerate(example_images):
     cv2.imwrite(save_path, img)
     example_paths.append([save_path])
 
-# Create Gradio interface
+# Gradio Interfaces
 PixelDehazer = gr.Interface(
-    fn=dehaze,
+    fn=process_single_image,
     inputs=gr.Image(type="numpy"),
     outputs="image",
     examples=example_paths,
-    cache_examples=False
+    cache_examples=False,
+    description="Upload a single image to remove haze."
+)
+
+BatchDehazer = gr.Interface(
+    fn=process_images,
+    inputs=gr.Files(label="Upload Multiple Images", file_types=["image"]),
+    outputs=gr.Files(label="Download Dehazed Images"),
+    description="Upload multiple images to remove haze. Download the processed dehazed images."
+)
+
+VideoDehazer = gr.Interface(
+    fn=process_video,
+    inputs=gr.Video(label="Upload a Video"),
+    outputs=gr.File(label="Download Dehazed Video"),
+    description="Upload a video to remove haze. Download the processed dehazed video."
+)
+
+# Combined Gradio App
+app = gr.TabbedInterface(
+    [PixelDehazer, BatchDehazer, VideoDehazer],
+    ["Single Image Dehazing", "Batch Image Dehazing", "Video Dehazing"],
+    title="DeFogify App"
 )
 
-PixelDehazer.launch()
+# Launch the Gradio App
+if __name__ == "__main__":
+    app.launch()