Blame for app.py

"""
Motion-Track Shape Censor  (CPU 전용, GPU 불필요)
클릭한 지점을 OpenCV 광학흐름으로 추적해서 흰 원/네모를 씌우는 경량 검열 도구.
SAM 계열의 형태 백색 마스킹이 못 잡는 부분을 '도형'으로 보조 검열하는 용도.
허깅페이스 무료 CPU Basic Space에 올리면 상시 요금 0원으로 항상 사용 가능.
"""

import os
import tempfile
import cv2
import numpy as np
import gradio as gr

MAX_SECONDS = 12  # 처리 상한 (짧은 클립 전용)

# ────────────────────────────────────────────────
#  상태
# ────────────────────────────────────────────────
STATE = {
    "video": None,
    "frames": [],   # RGB 프레임 리스트
    "gray": [],     # 그레이스케일(광학흐름용)
    "fps": 24.0,
    "w": 0,
    "h": 0,
}


def _load_video(video_path):
    if not video_path:
        return None, gr.update(maximum=1, value=0), [], gr.update(open=True)
    cap = cv2.VideoCapture(video_path)
    fps = cap.get(cv2.CAP_PROP_FPS) or 24.0
    frames, grays = [], []
    max_frames = int(fps * MAX_SECONDS)
    while len(frames) < max_frames:
        ok, fr = cap.read()
        if not ok:
            break
        rgb = cv2.cvtColor(fr, cv2.COLOR_BGR2RGB)
        frames.append(rgb)
        grays.append(cv2.cvtColor(fr, cv2.COLOR_BGR2GRAY))
    cap.release()
    if not frames:
        return None, gr.update(maximum=1, value=0), [], gr.update(open=True)
    STATE.update(video=video_path, frames=frames, gray=grays, fps=fps,
                 w=frames[0].shape[1], h=frames[0].shape[0])
    # 영상 로드되면 업로드 패널 접기
    return frames[0], gr.update(maximum=len(frames) - 1, value=0), [], gr.update(open=False)


def _draw_points(idx, clicks):
    """현재 프레임에 클릭 지점 표시."""
    if not STATE["frames"]:
        return None
    idx = int(np.clip(idx, 0, len(STATE["frames"]) - 1))
    img = STATE["frames"][idx].copy()
    for (x, y) in (clicks or []):
        cv2.circle(img, (int(x), int(y)), 6, (0, 200, 0), -1)
        cv2.circle(img, (int(x), int(y)), 6, (255, 255, 255), 1)
    return img


def _click(idx, clicks, evt: gr.SelectData):
    x, y = evt.index[0], evt.index[1]
    clicks = (clicks or []) + [[x, y]]
    return _draw_points(idx, clicks), clicks


def _reset(idx):
    return _draw_points(idx, []), []


def _preview(idx, clicks, mask_mode, shape_size):
    """현재 프레임에 도형을 씌운 미리보기."""
    if not STATE["frames"]:
        return None, "영상을 업로드하세요."
    if not clicks:
        return None, "가릴 지점을 클릭하세요."
    idx = int(np.clip(idx, 0, len(STATE["frames"]) - 1))
    img = STATE["frames"][idx].copy()
    s = int(shape_size)
    for (x, y) in clicks:
        x, y = int(x), int(y)
        if mask_mode == "흰 원":
            cv2.circle(img, (x, y), s, (255, 255, 255), -1)
        else:
            cv2.rectangle(img, (x - s, y - s), (x + s, y + s), (255, 255, 255), -1)
    return img, "흰 도형이 가려질 위치입니다. 위치·크기를 확인하고 처리 실행을 누르세요."


def _track_points(start_idx, points, error_threshold=30.0):
    """
    start_idx 프레임의 points를 광학흐름으로 전체 프레임에 대해 추적.
    반환: {frame_idx: [(x, y, visible), ...]} — 각 프레임에서의 점 위치 + 가시성.
    앞으로(→끝) + 뒤로(→처음) 양방향 추적.
    visible=False면 그 프레임에서 추적을 놓친 것(부위가 가려짐/사라짐) → 도형 숨김.
    """
    grays = STATE["gray"]
    n = len(grays)
    npts = len(points)
    tracked = {start_idx: [(float(p[0]), float(p[1]), True) for p in points]}

    lk_params = dict(winSize=(21, 21), maxLevel=3,
                     criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 20, 0.03))

    def run(direction):
        pts = np.array(points, dtype=np.float32).reshape(-1, 1, 2)
        rng = range(start_idx + 1, n) if direction > 0 else range(start_idx - 1, -1, -1)
        prev_pts = pts
        prev_idx = start_idx
        # 한 번 놓친 점은 그 방향으로 계속 숨김 (탁 사라진 뒤 유지)
        lost = [False] * npts
        for f in rng:
            new_pts, st, err = cv2.calcOpticalFlowPyrLK(
                grays[prev_idx], grays[f], prev_pts, None, **lk_params)
            if new_pts is None:
                new_pts = prev_pts
                st = np.zeros((npts, 1), dtype=np.uint8)
                err = np.full((npts, 1), 999.0, dtype=np.float32)
            row = []
            for i in range(npts):
                ok = (st is not None and int(st[i][0]) == 1)
                e = float(err[i][0]) if err is not None else 999.0
                # 추적 실패 또는 오차가 임계값 초과 → 이 점은 사라진 것으로 간주
                if (not ok) or (e > error_threshold):
                    lost[i] = True
                x = float(np.clip(new_pts[i, 0, 0], 0, STATE["w"] - 1))
                y = float(np.clip(new_pts[i, 0, 1], 0, STATE["h"] - 1))
                visible = not lost[i]
                row.append((x, y, visible))
                new_pts[i, 0, 0] = x
                new_pts[i, 0, 1] = y
            # 기존 값이 있으면(반대 방향에서 채워진 경우) 가시성은 OR로 합침
            if f in tracked:
                merged = []
                for a, b in zip(tracked[f], row):
                    merged.append((b[0], b[1], a[2] or b[2]))
                tracked[f] = merged
            else:
                tracked[f] = row
            prev_pts = new_pts
            prev_idx = f

    run(+1)   # 앞으로
    run(-1)   # 뒤로 (양방향)
    return tracked


def _render(start_idx, clicks, mask_mode, shape_size, sensitivity):
    if not STATE["frames"]:
        return None, "영상을 업로드하세요."
    if not clicks:
        return None, "가릴 지점을 클릭하세요."

    start_idx = int(np.clip(start_idx, 0, len(STATE["frames"]) - 1))
    points = [[float(x), float(y)] for (x, y) in clicks]

    # 민감도 → 오차 임계값 변환 (민감도 높을수록 쉽게 사라짐 = 임계값 낮음)
    # 슬라이더 1(둔감)~100(민감) → threshold 60~8
    error_threshold = 60.0 - (float(sensitivity) - 1) * (52.0 / 99.0)

    # 점 추적 (가시성 포함)
    tracked = _track_points(start_idx, points, error_threshold=error_threshold)

    # 렌더
    tmp = tempfile.mkdtemp()
    # 고유 파일명: 타임스탬프 + 짧은 랜덤 (분류/겹침방지 편의)
    import datetime, uuid
    stamp = datetime.datetime.now().strftime("%Y%m%d_%H%M%S")
    rand = uuid.uuid4().hex[:4]
    out_path = os.path.join(tmp, f"censored_{stamp}_{rand}.mp4")
    fourcc = cv2.VideoWriter_fourcc(*"mp4v")
    vw = cv2.VideoWriter(out_path, fourcc, STATE["fps"], (STATE["w"], STATE["h"]))

    s = int(shape_size)
    for f in range(len(STATE["frames"])):
        # BGR로 변환해서 기록
        frame_bgr = cv2.cvtColor(STATE["frames"][f], cv2.COLOR_RGB2BGR).copy()
        pts = tracked.get(f, [])
        for pt in pts:
            x, y, visible = int(pt[0]), int(pt[1]), pt[2]
            if not visible:
                continue  # 추적 놓친 프레임 → 도형 숨김 (탁 사라짐)
            if mask_mode == "흰 원":
                cv2.circle(frame_bgr, (x, y), s, (255, 255, 255), -1)
            else:
                cv2.rectangle(frame_bgr, (x - s, y - s), (x + s, y + s), (255, 255, 255), -1)
        vw.write(frame_bgr)
    vw.release()

    return out_path, f"완료! {len(clicks)}개 지점 추적. 가려지면 도형이 자동으로 사라집니다."


# ────────────────────────────────────────────────
#  UI
# ────────────────────────────────────────────────
with gr.Blocks(title="Motion-Track Shape Censor") as demo:
    gr.Markdown("""
    ### Motion-Track Shape Censor (도형 검열 · CPU 전용)
    클릭한 지점을 광학흐름으로 추적해서 흰 원/네모를 씌웁니다.
    형태 백색 마스킹이 못 잡는 부분을 도형으로 보조 검열하는 용도. **GPU 불필요.**
    """)

    clicks_st = gr.State([])

    with gr.Row():
        with gr.Column():
            with gr.Accordion("입력 영상 (업로드/교체하려면 펼치기)", open=True) as upload_acc:
                vid = gr.Video(label="입력 영상 (mp4)")
            frame_img = gr.Image(label="프레임 — 가릴 지점 클릭 (여러 개 가능)", interactive=True)
            frame_slider = gr.Slider(0, 1, value=0, step=1,
                                     label="프레임 선택 (대상이 잘 보이는 지점)")
            reset_btn = gr.Button("점 초기화")
        with gr.Column():
            preview_img = gr.Image(label="미리보기 (흰 도형 = 가려질 위치)")
            mask_mode = gr.Radio(["흰 원", "흰 네모"], value="흰 원", label="도형 종류")
            shape_size = gr.Slider(10, 200, value=50, step=5,
                                   label="도형 크기 (원=반지름 / 네모=반쪽 변, px)")
            sensitivity = gr.Slider(1, 100, value=50, step=1,
                                    label="사라짐 민감도 (높을수록 가려질 때 빨리 사라짐)")
            with gr.Row():
                preview_btn = gr.Button("미리보기")
                go_btn = gr.Button("처리 실행", variant="primary")
            status = gr.Textbox(label="상태", interactive=False)

    result_video = gr.Video(label="결과 영상 (원본 + 흰 도형)", height=280)

    vid.change(_load_video, inputs=vid,
               outputs=[frame_img, frame_slider, clicks_st, upload_acc])
    frame_slider.change(_draw_points, inputs=[frame_slider, clicks_st], outputs=frame_img)
    frame_img.select(_click, inputs=[frame_slider, clicks_st], outputs=[frame_img, clicks_st])
    reset_btn.click(_reset, inputs=frame_slider, outputs=[frame_img, clicks_st])
    preview_btn.click(_preview, inputs=[frame_slider, clicks_st, mask_mode, shape_size],
                      outputs=[preview_img, status])
    go_btn.click(_render, inputs=[frame_slider, clicks_st, mask_mode, shape_size, sensitivity],
                 outputs=[result_video, status])

    gr.Markdown("""
    **사용법**: 영상 업로드 → 프레임 이동 → 가릴 지점 클릭(여러 개 OK) → 미리보기로 확인 → 처리 실행
    · 각 점을 앞뒤로 추적해서 도형이 대상을 따라갑니다.
    · **부위가 손 등으로 가려지면 도형이 탁 사라집니다** — 잘 안 되면 '사라짐 민감도'를 조절하세요.
    · 대상이 급회전하거나 화면을 벗어나면 추적이 흔들릴 수 있어요(짧은 클립 권장).
    """)

if __name__ == "__main__":
    demo.queue(max_size=10).launch()