main.py

from transnetv2_pytorch import TransNetV2
from typing import Optional
import torch
import os
import numpy as np
from PIL import Image, ImageDraw
import argparse
from tqdm import tqdm

try:
    import ffmpeg
except ModuleNotFoundError:
    raise ModuleNotFoundError("For `predict_video` function `ffmpeg` needs to be installed in order to extract "
                                "individual frames from video file. Install `ffmpeg` command line tool and then "
                                "install python wrapper by `pip install ffmpeg-python`.")


class TransNetV2Torch:
    def __init__(self, model_path: Optional[str] = None):
        weights_path = model_path or os.path.join(os.path.dirname(__file__), "transnetv2-pytorch-weights.pth")
        if not os.path.isfile(weights_path):
            raise FileNotFoundError(f"[TransNetV2] ERROR: weights file not found at {weights_path}.")
        else:
            print(f"[TransNetV2] Using weights from {weights_path}.")

        self._input_size = (27, 48, 3)
        self.model = TransNetV2()
        try:
            self.model.load_state_dict(torch.load(weights_path))
        except Exception as exc:
            raise IOError(f"[TransNetV2] Could not load weights from {weights_path}.") from exc
        self.model.eval()
        self.device = "cuda" if torch.cuda.is_available() else "cpu"
        self.model.to(self.device)


    def predict_raw(self, frames: np.ndarray):
        assert len(frames.shape) == 5 and frames.shape[2:] == self._input_size, \
            "[TransNetV2] Input shape must be [batch, frames, height, width, 3]."
        
        frames_tensor = torch.from_numpy(frames)
        with torch.no_grad():
            single_frame_pred, all_frames_pred = self.model(frames_tensor.to(self.device))
            
            single_frame_pred = torch.sigmoid(single_frame_pred).cpu().numpy()
            all_frames_pred = torch.sigmoid(all_frames_pred["many_hot"]).cpu().numpy()

        return single_frame_pred, all_frames_pred

    def predict_frames(self, frames: np.ndarray):
        assert len(frames.shape) == 4 and frames.shape[1:] == self._input_size, \
            "[TransNetV2] Input shape must be [frames, height, width, 3]."

        total = len(frames)

        def input_iterator():
            # return windows of size 100 where the first/last 25 frames are from the previous/next batch
            # the first and last window must be padded by copies of the first and last frame of the video
            no_padded_frames_start = 25
            no_padded_frames_end = 25 + 50 - (total % 50 if total % 50 != 0 else 50)  # 25 - 74

            start_frame = np.expand_dims(frames[0], 0)
            end_frame = np.expand_dims(frames[-1], 0)
            padded_inputs = np.concatenate(
                [start_frame] * no_padded_frames_start + [frames] + [end_frame] * no_padded_frames_end, 0
            )

            ptr = 0
            while ptr + 100 <= len(padded_inputs):
                out = padded_inputs[ptr:ptr + 100]
                ptr += 50
                yield out[np.newaxis]

        predictions = []

        with tqdm(total=total, desc="[TransNetV2] Processing video frames", unit="frames") as pbar:
            for inp in input_iterator():
                single_frame_pred, all_frames_pred = self.predict_raw(inp)
                predictions.append((single_frame_pred[0, 25:75, 0],
                                    all_frames_pred[0, 25:75, 0]))

                processed = min(len(predictions) * 50, total)
                pbar.n = processed
                pbar.last_print_n = processed
                pbar.refresh()

        single_frame_pred = np.concatenate([single_ for single_, _ in predictions])
        all_frames_pred = np.concatenate([all_ for _, all_ in predictions])

        return single_frame_pred[:total], all_frames_pred[:total]


    def predict_video(self, video_fn: str):
        print("[TransNetV2] Extracting frames from {}".format(video_fn))
        video_stream, _ = ffmpeg.input(video_fn).output(
            "pipe:", format="rawvideo", pix_fmt="rgb24", s="48x27"
        ).run(capture_stdout=True, capture_stderr=True)

        video = np.frombuffer(video_stream, np.uint8).reshape([-1, 27, 48, 3])
        return (video, *self.predict_frames(video))

    @staticmethod
    def predictions_to_scenes(predictions: np.ndarray, threshold: float = 0.5):
        predictions = (predictions > threshold).astype(np.uint8)

        scenes = []
        t_prev, start = 0, 0
        for i, t in enumerate(predictions):
            if t_prev == 1 and t == 0:
                start = i
            if t_prev == 0 and t == 1 and i != 0:
                scenes.append([start, i])
            t_prev = t
        if t == 0:
            scenes.append([start, i])
        if len(scenes) == 0: # just fix if all predictions are 1
            return np.array([[0, len(predictions) - 1]], dtype=np.int32)

        return np.array(scenes, dtype=np.int32)

    @staticmethod
    def visualize_predictions(frames: np.ndarray, predictions):

        if isinstance(predictions, np.ndarray):
            predictions = [predictions]

        ih, iw, ic = frames.shape[1:]
        width = 25

        # pad frames so that length of the video is divisible by width
        # pad frames also by len(predictions) pixels in width in order to show predictions
        pad_with = width - len(frames) % width if len(frames) % width != 0 else 0
        frames = np.pad(frames, [(0, pad_with), (0, 1), (0, len(predictions)), (0, 0)])

        predictions = [np.pad(x, (0, pad_with)) for x in predictions]
        height = len(frames) // width

        img = frames.reshape([height, width, ih + 1, iw + len(predictions), ic])
        img = np.concatenate(np.split(
            np.concatenate(np.split(img, height), axis=2)[0], width
        ), axis=2)[0, :-1]

        img = Image.fromarray(img)
        draw = ImageDraw.Draw(img)

        for i, pred in enumerate(zip(*predictions)):
            x, y = i % width, i // width
            x, y = x * (iw + len(predictions)) + iw, y * (ih + 1) + ih - 1

            # we can visualize multiple predictions per single frame
            for j, p in enumerate(pred):
                color = [0, 0, 0]
                color[(j + 1) % 3] = 255

                value = round(p * (ih - 1))
                if value != 0:
                    draw.line((x + j, y, x + j, y - value), fill=tuple(color), width=1)
        return img

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--files", type=str, help="path to video files to process")
    parser.add_argument("--weights", type=str, default=None,
                        help="path to TransNet V2 weights, tries to infer the location if not specified")
    parser.add_argument('--visualize', action="store_true",
                        help="save a png file with prediction visualization for each extracted video")
    args = parser.parse_args()
    
    return args

def main(args):
    model = TransNetV2Torch(args.weights)

    files = []
    if os.path.isdir(args.files):
        for f in os.listdir(args.files):
            if f.lower().endswith(".mp4"):
                files.append(os.path.join(args.files, f))
    else:
        files = [args.files]

    for file in files:
        video_frames, single_frame_predictions, all_frames_predictions = \
            model.predict_video(file)

        predictions = np.stack([single_frame_predictions, all_frames_predictions], 1)
        np.savetxt(file + ".predictions.txt", predictions, fmt="%.6f")

        scenes = model.predictions_to_scenes(single_frame_predictions)
        np.savetxt(file + ".scenes.txt", scenes, fmt="%d")

        if args.visualize:
            pil_image = model.visualize_predictions(
                video_frames, predictions=(single_frame_predictions, all_frames_predictions))
            pil_image.save(file + ".vis.png")

if __name__ == "__main__":
    args = parse_args()
    main(args)