Spaces:
Sleeping
Sleeping
| import os | |
| import random | |
| import time | |
| from typing import Tuple, Union | |
| import cv2 | |
| import numpy as np | |
| import streamlit as st | |
| from PIL import Image | |
| from torch import nn | |
| num_format = "{:,}".format | |
| def count_parameters(model: nn.Module) -> str: | |
| """Count the number of parameters of a model""" | |
| return num_format(sum(p.numel() for p in model.parameters() if p.requires_grad)) | |
| class FrameRate: | |
| def __init__(self) -> None: | |
| self.c: int = 0 | |
| self.start_time: float = None | |
| self.NO_FRAMES = 100 | |
| self.fps: float = -1 | |
| def reset(self) -> None: | |
| self.start_time = time.time() | |
| self.c = 0 | |
| self.fps = -1 | |
| def count(self) -> None: | |
| self.c += 1 | |
| if self.c % self.NO_FRAMES == 0: | |
| self.c = 0 | |
| end_time = time.time() | |
| self.fps = self.NO_FRAMES / (end_time - self.start_time) | |
| self.start_time = end_time | |
| def show_fps(self, image: np.ndarray) -> np.ndarray: | |
| if self.fps != -1: | |
| return cv2.putText( | |
| image, | |
| f"FPS {self.fps:.0f}", | |
| (50, 50), | |
| cv2.FONT_HERSHEY_SIMPLEX, | |
| fontScale=1, | |
| color=(255, 0, 0), | |
| thickness=2, | |
| ) | |
| else: | |
| return image | |
| class ImgContainer: | |
| img: np.ndarray = None # raw image | |
| frame_rate: FrameRate = FrameRate() | |
| def load_video(video_path: str) -> bytes: | |
| if not os.path.isfile(video_path): | |
| return | |
| with st.spinner(f"Loading video {video_path} ..."): | |
| video_bytes = open(video_path, "rb").read() | |
| st.video(video_bytes, format="video/mp4") | |
| def normalize(data: np.ndarray) -> np.ndarray: | |
| return (data - data.min()) / (data.max() - data.min() + 1e-8) | |
| def get_size(image: Union[Image.Image, np.ndarray]) -> Tuple[int, int]: | |
| """Get resolution (w, h) of an image | |
| An input image can be Pillow Image or CV2 Image | |
| """ | |
| if type(image) == np.ndarray: | |
| return (image.shape[1], image.shape[0]) | |
| else: | |
| return image.size | |
| def random_choice(p: float) -> bool: | |
| """Return True if random float <= p""" | |
| return random.random() <= p | |