ivar / utils /video_processing.py
FabioSantos's picture
Upload 10 files
770a166 verified
import cv2
import tempfile
import os
import time
import base64
from langchain_core.messages import HumanMessage
from langchain_google_genai import ChatGoogleGenerativeAI
os.environ["GOOGLE_API_KEY"] = "SUA_CHAVE_GEMINI_AQUI" # Ajuste essa chave
gemini_model = ChatGoogleGenerativeAI(model="gemini-2.0-flash")
FRAME_FOLDER = "full_frames"
os.makedirs(FRAME_FOLDER, exist_ok=True)
def parse_irregularity(observation):
detected_irregularities = []
lines = observation.split("\n")
for line in lines:
parts = line.split("|")
if len(parts) >= 4:
tipo = parts[1].strip()
observado = parts[2].strip().lower()
descricao = parts[3].strip() or "Não especificado"
if observado == "yes":
detected_irregularities.append({"tipo": tipo, "descricao": descricao})
return detected_irregularities
async def processar_video(video_file):
temp_path = tempfile.mktemp(suffix=".mp4")
with open(temp_path, "wb") as f:
f.write(video_file.file.read())
cap = cv2.VideoCapture(temp_path)
total_frames = int(cap.get(cv2.CAP_PROP_FRAME_COUNT))
fps = cap.get(cv2.CAP_PROP_FPS)
duration = total_frames / fps
SEND_INTERVAL = 1.5 if duration <= 30 else (2 if duration <= 180 else 4)
irregularidades = []
def analyze_with_gemini(image_path):
with open(image_path, "rb") as img_file:
base64_image = base64.b64encode(img_file.read()).decode("utf-8")
message = HumanMessage(content=[
{"type": "text", "text": "Analyze this football play image using FIFA Laws..."},
{"type": "image_url", "image_url": {"url": f"data:image/jpeg;base64,{base64_image}"}}
])
response = gemini_model.invoke([message])
return parse_irregularity(response.content.strip())
frame_count, last_sent_time = 0, 0
while cap.isOpened():
ret, frame = cap.read()
if not ret:
break
frame = cv2.resize(frame, (640, 360))
if time.time() - last_sent_time >= SEND_INTERVAL:
last_sent_time = time.time()
frame_path = os.path.join(FRAME_FOLDER, f"frame_{time.strftime('%Y-%m-%d_%H-%M-%S')}.jpg")
cv2.imwrite(frame_path, frame)
irregularidades.extend(analyze_with_gemini(frame_path))
cap.release()
os.remove(temp_path)
return irregularidades