File size: 5,807 Bytes
5d346e0 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 |
import gradio as gr
import numpy as np
import cv2
import tempfile
def extract_first_frame(video_file):
# Open the video using OpenCV. The video_file is assumed to be a file-like object.
cap = cv2.VideoCapture(video_file.name if hasattr(video_file, "name") else video_file)
frame = None
while True:
ret, frame = cap.read()
if not ret:
break
if frame is not None and frame.size != 0:
break
cap.release()
if frame is None or frame.size == 0:
return None
# Convert frame from BGR to RGB format.
frame_rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)
return frame_rgb
def process_video(editor_image, epsilon_ratio, video_file):
# This function is analogous to your original "process_image" but now,
# it uses the drawing from the editor (on the first frame) and overlays
# the derived polygon on every frame of the uploaded video.
if editor_image is None:
return "❌ No image provided.", None
composite = editor_image.get("composite")
original = editor_image.get("background")
if composite is None or original is None:
return "⚠️ Please load the first frame and add a drawing layer.", None
composite_np = np.array(composite)
original_np = np.array(original)
# Extract red channel information to detect drawn strokes
r_channel = composite_np[:, :, 0]
g_channel = composite_np[:, :, 1]
b_channel = composite_np[:, :, 2]
# Use a threshold to detect red strokes (assuming user draws with a vivid red)
red_mask = (r_channel > 150) & (g_channel < 100) & (b_channel < 100)
binary_mask = red_mask.astype(np.uint8) * 255
# Find contours from the binary mask
contours, _ = cv2.findContours(binary_mask, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
if not contours:
return "⚠️ No visible drawing found. Please use the brush on a new layer.", None
largest_contour = max(contours, key=cv2.contourArea)
# Approximate contour to polygon using provided epsilon_ratio
epsilon = epsilon_ratio * cv2.arcLength(largest_contour, True)
polygon = cv2.approxPolyDP(largest_contour, epsilon, True)
if polygon is None or len(polygon) < 3:
return "⚠️ Polygon extraction failed. Try drawing a clearer shape.", None
polygon = polygon.astype(np.int32).reshape((-1, 1, 2))
polygon_coords = polygon.reshape(-1, 2).tolist()
# Open the input video for overlaying the polygon on every frame.
cap = cv2.VideoCapture(video_file.name if hasattr(video_file, "name") else video_file)
fourcc = cv2.VideoWriter_fourcc(*'mp4v')
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 a temporary file for saving the output video
temp_output = tempfile.NamedTemporaryFile(delete=False, suffix=".mp4").name
out = cv2.VideoWriter(temp_output, fourcc, fps, (width, height))
# Process each frame and draw the polygon overlay
while True:
ret, frame = cap.read()
if not ret:
break
overlay = frame.copy()
cv2.polylines(overlay, [polygon], isClosed=True, color=(0, 255, 0), thickness=10)
for idx, (x, y) in enumerate(polygon_coords):
cv2.circle(overlay, (x, y), 5, (0, 0, 255), -1)
cv2.putText(overlay, str(idx + 1), (x + 5, y - 5),
cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 255, 255), 2)
out.write(overlay)
cap.release()
out.release()
msg = f"✅ Polygon with {len(polygon_coords)} points (ε={epsilon_ratio}):\n{polygon_coords}"
return msg, temp_output
# Build the Gradio interface using Blocks
with gr.Blocks() as demo:
gr.HTML("<style>body { margin: 0; padding: 0; }</style>")
gr.Markdown("## 🖌️ Accurate Polygon Extraction & Video Overlay")
gr.Markdown(
"""
**Instructions:**
1. Upload a video.
2. Click **Load First Frame to Editor** to extract a frame for annotation.
3. ➕ Add a drawing layer and draw with the brush (use red strokes).
4. Adjust polygon approximation if needed.
5. Click **Process Drawing and Overlay on Video** — the generated video will show the green polygon overlaid on every frame.
"""
)
with gr.Tab("Load Video"):
video_input = gr.Video(label="Upload Video", format="mp4")
load_frame_btn = gr.Button("Load First Frame to Editor")
# The ImageEditor will be preloaded with the extracted frame.
frame_editor = gr.ImageEditor(label="Draw on this frame (Add a layer first!)", type="numpy", width=1920, height=1080)
epsilon_slider = gr.Slider(
label="Polygon Approximation (ε)", minimum=0.001, maximum=0.05, value=0.01, step=0.001
)
with gr.Row():
output_text = gr.Textbox(label="Polygon Coordinates", lines=6)
video_preview = gr.Video(label="Video with Polygon Overlay", format="mp4")
# Function to load the first non-empty frame from the uploaded video.
def load_frame(video_file):
frame = extract_first_frame(video_file)
if frame is None:
return gr.update(value=None), "❌ Failed to extract frame from video."
# Return the frame for the editor and a confirmation message.
return frame, "Frame loaded successfully."
load_frame_btn.click(fn=load_frame, inputs=video_input, outputs=[frame_editor, output_text])
# Process the drawing and overlay the polygon on the video.
process_btn = gr.Button("Process Drawing and Overlay on Video")
process_btn.click(fn=process_video, inputs=[frame_editor, epsilon_slider, video_input],
outputs=[output_text, video_preview])
demo.launch()
|