|
import gradio as gr |
|
import cv2 |
|
import torch |
|
from transformers import AutoModelForImageTranslation |
|
|
|
|
|
model = AutoModelForImageTranslation.from_pretrained("username/gaze_correction_model") |
|
|
|
|
|
def correct_gaze_in_video(video_path): |
|
|
|
cap = cv2.VideoCapture(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)) |
|
|
|
|
|
output_path = "corrected_gaze_output.mp4" |
|
fourcc = cv2.VideoWriter_fourcc(*'mp4v') |
|
out = cv2.VideoWriter(output_path, fourcc, fps, (width, height)) |
|
|
|
while cap.isOpened(): |
|
ret, frame = cap.read() |
|
if not ret: |
|
break |
|
|
|
|
|
inputs = preprocess_image(frame) |
|
corrected_frame = model(inputs) |
|
|
|
|
|
out.write(corrected_frame) |
|
|
|
|
|
cap.release() |
|
out.release() |
|
|
|
return output_path |
|
|
|
|
|
iface = gr.Interface( |
|
fn=correct_gaze_in_video, |
|
inputs="file", |
|
outputs="file", |
|
title="Gaze Correction in Video", |
|
description="Upload a video to correct gaze direction." |
|
) |
|
|
|
iface.launch() |
|
|