Update app.py

app.py

@@ -28,23 +28,18 @@ def remove_watermark_area(original_image, text_mask_path):
     return cleaned_image
 from PIL import Image
 
-def remove_watermark(image_path,saved_path):
-    if isinstance(image_path, str) and os.path.isfile(image_path):
+def remove_watermark(image_path,file_type="",saved_path=""):
+    # file_type="pil"
+    # file_type="opencv"
+    # file_type="filepath"
+    if file_type=="filepath":
         # Load the image using OpenCV
         image = cv2.imread(image_path) 
-    elif isinstance(image_path, np.ndarray):
-        # Directly use OpenCV image (NumPy array)
-        image = image_path
-        if len(image_path.shape) == 3 and image_path.shape[2] == 3:
-            # Assuming it's in RGB format; convert to BGR
-            image = cv2.cvtColor(image_path, cv2.COLOR_RGB2BGR)
-        else:
-            # Otherwise, assume it's already in BGR format
-            image = image_path
-    else:
-        raise TypeError("Invalid image_path format")
-    print(type(image))
-    cv2.imwrite("test.jpg",image)
+    if file_type=="pil":
+        image = cv2.cvtColor(image_path, cv2.COLOR_RGB2BGR)
+    if file_type=="opencv":
+      image=image_path
+    # cv2.imwrite("test.jpg",image)
     image=cv2.resize(image,(1280,1280))
     # Define the area of the watermark (adjust this based on the watermark size)
     height, width, _ = image.shape
@@ -65,7 +60,10 @@ def remove_watermark(image_path,saved_path):
     # cv2.imwrite('cleaned_watermark.jpg', cleaned_image)
     # Paste back the cleaned watermark on the original image
     image[y_start:y_end, x_start:x_end] = cleaned_image
-    cv2.imwrite(saved_path, image)
+    if saved_path=="":
+      pass
+    else:
+      cv2.imwrite(saved_path, image)
     return image
 
 def make_zip(image_list):
@@ -82,7 +80,7 @@ def random_image_name():
     
 def process_file(pil_image):
     saved_path = f"./temp/{random_image_name()}.jpg"
-    remove_watermark(pil_image, saved_path)
+    remove_watermark(pil_image,"pil",saved_path)
     return saved_path, saved_path
 
 
@@ -92,19 +90,59 @@ def process_files(image_files):
         # saved_path = os.path.basename(image_files[0])
         # saved_path = f"./temp/{saved_path}"
         saved_path = f"./temp/{random_image_name()}.jpg"
-        remove_watermark(image_files[0], saved_path)
+        remove_watermark(image_files[0],"filepath", saved_path)
         return saved_path, saved_path
     else:
         for image_path in image_files:
             # saved_path = os.path.basename(image_path)
             # saved_path = f"./temp/{saved_path}"
             saved_path = f"./temp/{random_image_name()}.jpg"
-            remove_watermark(image_path, saved_path)
+            remove_watermark(image_path,"filepath",saved_path)
             image_list.append(saved_path)
         zip_path = make_zip(image_list)
         return zip_path,None
 
 
+import cv2
+import numpy as np
+
+
+
+def process_video(input_video_path):
+    print(input_video_path)
+    # output_video_path=""
+    # if output_video_path=="":
+    output_video_path=f"./temp/{random_image_name()}.mp4"
+    cap = cv2.VideoCapture(input_video_path)
+    
+    if not cap.isOpened():
+        raise ValueError(f"Unable to open video file {input_video_path}")
+
+    fps = cap.get(cv2.CAP_PROP_FPS)
+    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))
+    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))
+    
+    # Create VideoWriter object for output video
+    fourcc = cv2.VideoWriter_fourcc(*'mp4v')  # You might try 'XVID' or 'H264' if issues persist
+    video_writer = cv2.VideoWriter(output_video_path, fourcc, fps, (width, height))
+    
+    # Process frames
+    while cap.isOpened():
+        ret, frame = cap.read()
+        if not ret:
+            break
+        no_watermark_frame=remove_watermark(frame,"opencv")
+        
+        # Ensure the frame has the same size and type
+        if no_watermark_frame.shape[1] != width or no_watermark_frame.shape[0] != height:
+            no_watermark_frame = cv2.resize(no_watermark_frame, (width, height))
+        
+        video_writer.write(no_watermark_frame)
+    
+    cap.release()
+    video_writer.release()
+    return output_video_path,output_video_path
+    
 if not os.path.exists("./temp"):
     os.mkdir("./temp")
 
@@ -114,7 +152,7 @@ meta_examples = ["./images/1.jpg", "./images/2.jpg", "./images/3.jpg", "./images
 gradio_input=[gr.Image(label='Upload an Image')]
 gradio_Output=[gr.File(label='Download Image'),gr.Image(label='Display Image')]
 gradio_interface = gr.Interface(fn=process_file, inputs=gradio_input,outputs=gradio_Output ,
-                              title="Meta Watermark Remover",
+                              title="Meta Watermark Remover From Image",
                               examples=meta_examples)
 # gradio_interface.launch(debug=True)
 
@@ -127,6 +165,15 @@ gradio_multiple_images = gr.Interface(
     cache_examples=True
 )
 
-demo = gr.TabbedInterface([gradio_interface, gradio_multiple_images], ["Meta Watermark Remover","Bluk Meta Watermark Remover"],title="Meta Watermark Remover")
-demo.launch(debug=True)
 
+meta_video_examples = [ "./videos/1.mp4"]
+
+gradio_video_input=[gr.Video(label='Upload an Video')]
+gradio_video_Output=[gr.File(label='Download Video'),gr.Video(label='Display Video')]
+gradio_video_interface = gr.Interface(fn=process_video, inputs=gradio_video_input,outputs=gradio_video_Output ,
+                              title="Meta Watermark Remover from Video",
+                              examples=meta_video_examples)
+
+
+demo = gr.TabbedInterface([gradio_interface, gradio_video_interface,gradio_multiple_images], ["Meta Watermark Remover","Meta Watermark Remover from Video","Bluk Meta Watermark Remover"],title="Meta Watermark Remover")
+demo.launch(debug=True)