├── _env ├── home │ ├── .gitignore │ ├── .sensitive │ │ └── .gitignore │ ├── bin │ │ └── watch-mem │ └── .bashrc ├── .dockerignore ├── __hakase__.jpg ├── arial_monospaced_mt.ttf ├── .gitignore ├── project_config.bashrc ├── machine_config_template.bashrc └── make_config.bashrc ├── _data └── .gitignore ├── make ├── jupyterlab ├── docker_build ├── docker_pull ├── docker_push ├── docker_stop ├── shell_docker ├── shell_singularity ├── singularity_build ├── runnable ├── experiment └── _make.sh ├── .gitignore ├── vtubers-dataset.sublime-project ├── readme.md ├── _scripts └── download_vtubers.py └── _util ├── pytorch_v1.py ├── util_v1.py └── twodee_v1.py /_env/home/.gitignore: -------------------------------------------------------------------------------- 1 | 2 | 3 | -------------------------------------------------------------------------------- /_data/.gitignore: -------------------------------------------------------------------------------- 1 | 2 | * 3 | 4 | !.gitignore 5 | -------------------------------------------------------------------------------- /_env/home/.sensitive/.gitignore: -------------------------------------------------------------------------------- 1 | * 2 | !.gitignore 3 | -------------------------------------------------------------------------------- /_env/.dockerignore: -------------------------------------------------------------------------------- 1 | 2 | 3 | home/ 4 | *.sif 5 | *.bak 6 | -------------------------------------------------------------------------------- /make/jupyterlab: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | ./make/_make.sh jupyterlab 4 | 5 | 6 | -------------------------------------------------------------------------------- /make/docker_build: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | ./make/_make.sh docker_build 4 | 5 | 6 | -------------------------------------------------------------------------------- /make/docker_pull: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | ./make/_make.sh docker_pull 4 | 5 | 6 | -------------------------------------------------------------------------------- /make/docker_push: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | ./make/_make.sh docker_push 4 | 5 | 6 | -------------------------------------------------------------------------------- /make/docker_stop: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | ./make/_make.sh docker_stop 4 | 5 | 6 | -------------------------------------------------------------------------------- /make/shell_docker: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | ./make/_make.sh shell_docker 4 | 5 | 6 | -------------------------------------------------------------------------------- /make/shell_singularity: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | ./make/_make.sh shell_singularity 4 | 5 | 6 | -------------------------------------------------------------------------------- /make/singularity_build: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | ./make/_make.sh singularity_build 4 | 5 | 6 | -------------------------------------------------------------------------------- /_env/__hakase__.jpg: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShuhongChen/vtubers-dataset/HEAD/_env/__hakase__.jpg -------------------------------------------------------------------------------- /_env/home/bin/watch-mem: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | # monitor gpu 4 | watch --no-title -n 1 "nvidia-smi; free -h" 5 | 6 | -------------------------------------------------------------------------------- /_env/arial_monospaced_mt.ttf: -------------------------------------------------------------------------------- https://raw.githubusercontent.com/ShuhongChen/vtubers-dataset/HEAD/_env/arial_monospaced_mt.ttf -------------------------------------------------------------------------------- /_env/.gitignore: -------------------------------------------------------------------------------- 1 | 2 | 3 | *.sif 4 | *.deb 5 | *.tar 6 | *.tar.xz 7 | *.tar.gz 8 | *.zip 9 | *.bak 10 | 11 | machine_config.bashrc 12 | -------------------------------------------------------------------------------- /_env/project_config.bashrc: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | 4 | export PROJECT_NAME=vtubers-dataset 5 | 6 | export DOCKER_USER=shuhongchen 7 | export DOCKER_NAME=panic3d-anime-reconstruction 8 | 9 | -------------------------------------------------------------------------------- /.gitignore: -------------------------------------------------------------------------------- 1 | 2 | data/ 3 | files/ 4 | temp/ 5 | help/ 6 | help 7 | _notes.txt 8 | _notes/ 9 | tmp/ 10 | var_tmp/ 11 | 12 | *.sublime-workspace 13 | **/___* 14 | **/__pycache__ 15 | **/.ipynb_checkpoints 16 | **/.directory 17 | **/.Trash* 18 | **/.empty 19 | **/slurm-* 20 | **/runs 21 | **/lightning_logs 22 | **/slurm 23 | -------------------------------------------------------------------------------- /_env/machine_config_template.bashrc: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | 4 | export MACHINE_NAME= 5 | export PROJECT_DN=/absolute/path/to/repo 6 | 7 | export JUPYTERLAB_PORT=8889 8 | export JUPYTERLAB_HOST=localhost 9 | 10 | export MOUNTS_DOCKER=" 11 | " 12 | export MOUNTS_SINGULARITY=" 13 | " 14 | 15 | export SINGULARITY_CACHEDIR= 16 | export SINGULARITY_TMPDIR= 17 | 18 | -------------------------------------------------------------------------------- /make/runnable: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | source ./_env/project_config.bashrc 4 | source ./_env/machine_config.bashrc 5 | source ./_env/make_config.bashrc 6 | 7 | for fn in ./_env/home/.sensitive/*.bashrc; do 8 | if [ -f $fn ]; then 9 | . $fn 10 | fi 11 | done 12 | 13 | 14 | python3 -m make.writers_runnable.runnable_v0 "$@" 15 | 16 | 17 | 18 | -------------------------------------------------------------------------------- /make/experiment: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | source ./_env/project_config.bashrc 4 | source ./_env/machine_config.bashrc 5 | source ./_env/make_config.bashrc 6 | 7 | for fn in ./_env/home/.sensitive/*.bashrc; do 8 | if [ -f $fn ]; then 9 | . $fn 10 | fi 11 | done 12 | 13 | 14 | python3 -m make.writers_experiment \ 15 | --trainer=$1 \ 16 | --machine=$2 \ 17 | --version=$3 \ 18 | --qos=$4 \ 19 | --gpus=$5 20 | 21 | 22 | 23 | -------------------------------------------------------------------------------- /vtubers-dataset.sublime-project: -------------------------------------------------------------------------------- 1 | { 2 | "folders": 3 | [ 4 | { 5 | "path": ".", 6 | "folder_exclude_patterns": [ 7 | "*/__pycache__", 8 | "*.directory", 9 | "*/.ipynb_checkpoints", 10 | "*/.empty", 11 | ".git/", 12 | "_data/", 13 | "temp/", 14 | "**.pkl", 15 | "*/node_modules" 16 | ], 17 | "follow_symlinks": true, 18 | } 19 | ], 20 | } 21 | -------------------------------------------------------------------------------- /_env/make_config.bashrc: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | 4 | if [ -f ~/project_config.bashrc ] ; then 5 | source ~/project_config.bashrc 6 | source ~/machine_config.bashrc 7 | elif [ -f ./_env/project_config.bashrc ] ; then 8 | source ./_env/project_config.bashrc 9 | source ./_env/machine_config.bashrc 10 | fi 11 | 12 | 13 | export OPTIONS_DOCKER_RUN=" 14 | --rm 15 | --gpus=all 16 | --ipc=host 17 | --net=host 18 | -w $PROJECT_DN 19 | -e DISPLAY=$DISPLAY 20 | -v /dev/shm:/dev/shms 21 | -v $PROJECT_DN:$PROJECT_DN 22 | -v $PROJECT_DN/_env/project_config.bashrc:/root/project_config.bashrc 23 | -v $PROJECT_DN/_env/machine_config.bashrc:/root/machine_config.bashrc 24 | -v $PROJECT_DN/_env/make_config.bashrc:/root/make_config.bashrc 25 | -v $PROJECT_DN/_env/home/bin:/root/bin 26 | -v $PROJECT_DN/_env/home/.bashrc:/root/.bashrc 27 | -v $PROJECT_DN/_env/home/.sensitive:/root/.sensitive 28 | " 29 | 30 | export OPTIONS_SINGULARITY_EXEC=" 31 | --nv 32 | --no-home 33 | --containall 34 | --home /root 35 | -W $PROJECT_DN 36 | --env DISPLAY=$DISPLAY 37 | -B /tmp:/tmp 38 | -B /var/tmp:/var/tmp 39 | -B $PROJECT_DN:$PROJECT_DN 40 | -B $PROJECT_DN/_env/project_config.bashrc:/root/project_config.bashrc 41 | -B $PROJECT_DN/_env/machine_config.bashrc:/root/machine_config.bashrc 42 | -B $PROJECT_DN/_env/make_config.bashrc:/root/make_config.bashrc 43 | -B $PROJECT_DN/_env/home/bin:/root/bin 44 | -B $PROJECT_DN/_env/home/.bashrc:/root/.bashrc 45 | -B $PROJECT_DN/_env/home/.sensitive:/root/.sensitive 46 | " 47 | 48 | 49 | 50 | 51 | 52 | 53 | 54 | 55 | -------------------------------------------------------------------------------- /make/_make.sh: -------------------------------------------------------------------------------- 1 | #!/bin/bash 2 | 3 | 4 | source ./_env/project_config.bashrc 5 | source ./_env/machine_config.bashrc 6 | source ./_env/make_config.bashrc 7 | 8 | 9 | if [[ $1 == shell_docker ]]; then 10 | ([ -z $DISPLAY ] && echo "no display" || xhost +local:docker) \ 11 | && sudo docker run \ 12 | $OPTIONS_DOCKER_RUN \ 13 | $MOUNTS_DOCKER \ 14 | -it $DOCKER_USER/$DOCKER_NAME:latest \ 15 | /bin/bash 16 | 17 | elif [[ $1 == jupyterlab ]]; then 18 | ([ -z $DISPLAY ] && echo "no display" || xhost +local:docker) \ 19 | && sudo docker run \ 20 | $OPTIONS_DOCKER_RUN \ 21 | $MOUNTS_DOCKER \ 22 | -it $DOCKER_USER/$DOCKER_NAME:latest \ 23 | /bin/bash -c " 24 | source /root/.bashrc \ 25 | && python3 -m jupyterlab --notebook-dir / --ip $JUPYTERLAB_HOST --port $JUPYTERLAB_PORT \ 26 | --allow-root --no-browser --ContentsManager.allow_hidden=True 27 | " 28 | 29 | elif [[ $1 == shell_singularity ]]; then 30 | singularity exec \ 31 | $OPTIONS_SINGULARITY_EXEC \ 32 | $MOUNTS_SINGULARITY \ 33 | ./_env/singularity.sif \ 34 | /bin/bash 35 | 36 | elif [[ $1 == docker_build ]]; then 37 | sudo docker build -t $DOCKER_USER/$DOCKER_NAME:latest $PROJECT_DN/_env 38 | elif [[ $1 == docker_push ]]; then 39 | sudo docker push $DOCKER_USER/$DOCKER_NAME:latest 40 | elif [[ $1 == docker_pull ]]; then 41 | sudo docker pull $DOCKER_USER/$DOCKER_NAME:latest 42 | elif [[ $1 == docker_stop ]]; then 43 | sudo docker stop $(sudo docker ps -aq) && sudo docker rm $(sudo docker ps -aq) 44 | 45 | elif [[ $1 == singularity_build ]]; then 46 | sudo -E singularity build \ 47 | $PROJECT_DN/_env/singularity.sif \ 48 | $PROJECT_DN/_env/singularity.def 49 | 50 | fi 51 | 52 | 53 | -------------------------------------------------------------------------------- /readme.md: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | # vtubers-dataset 5 | 6 | This repo downloads the Vtubers illustrations dataset introduced in [PAniC-3D: Stylized Single-view 3D Reconstruction from Portraits of Anime Characters](https://github.com/ShuhongChen/panic3d-anime-reconstruction). As described in that repo, this downloader will add to `./_data/lustrous` 7 | 8 | 9 | ## setup 10 | 11 | Make a copy of `./_env/machine_config.bashrc.template` to `./_env/machine_config.bashrc`, and set `$PROJECT_DN` to the absolute path of this repository folder. The other variables are optional. 12 | 13 | Run these lines from the project directory to pull the image from the main project and enter a container; note these are bash scripts inside the `./make` folder, not `make` commands. 14 | 15 | make/docker_pull 16 | make/shell_docker 17 | 18 | 19 | ## download 20 | 21 | First, download the `panic_data_models_merged.zip` from the project's [drive folder](https://drive.google.com/drive/folders/1Zpt9x_OlGALi-o-TdvBPzUPcvTc7zpuV?usp=share_link), and merge it with this repo's file structure. 22 | 23 | There should be a `./_data/lustrous/raw/vroid/metadata.json`, which the following commands will use to download the models. Note that `metadata.json` also contains all vroid model attributions. 24 | 25 | # run the downloader 26 | python3 -m _scripts.download_vtubers 27 | 28 | 29 | ## citing 30 | 31 | If you use our repo, please cite our work: 32 | 33 | @inproceedings{chen2023panic3d, 34 | title={PAniC-3D: Stylized Single-view 3D Reconstruction from Portraits of Anime Characters}, 35 | author={Chen, Shuhong and Zhang, Kevin and Shi, Yichun and Wang, Heng and Zhu, Yiheng and Song, Guoxian and An, Sizhe and Kristjansson, Janus and Yang, Xiao and Matthias Zwicker}, 36 | booktitle={Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR)}, 37 | year={2023} 38 | } 39 | 40 | 41 | 42 | -------------------------------------------------------------------------------- /_env/home/.bashrc: -------------------------------------------------------------------------------- 1 | # ~/.bashrc: executed by bash(1) for non-login shells. 2 | # see /usr/share/doc/bash/examples/startup-files (in the package bash-doc) 3 | # for examples 4 | 5 | # If not running interactively, don't do anything 6 | # case $- in 7 | # *i*) ;; 8 | # *) return;; 9 | # esac 10 | 11 | # don't put duplicate lines or lines starting with space in the history. 12 | # See bash(1) for more options 13 | HISTCONTROL=ignoreboth 14 | 15 | # append to the history file, don't overwrite it 16 | shopt -s histappend 17 | 18 | # for setting history length see HISTSIZE and HISTFILESIZE in bash(1) 19 | HISTSIZE=1000 20 | HISTFILESIZE=2000 21 | 22 | # check the window size after each command and, if necessary, 23 | # update the values of LINES and COLUMNS. 24 | shopt -s checkwinsize 25 | 26 | # If set, the pattern "**" used in a pathname expansion context will 27 | # match all files and zero or more directories and subdirectories. 28 | #shopt -s globstar 29 | 30 | # make less more friendly for non-text input files, see lesspipe(1) 31 | [ -x /usr/bin/lesspipe ] && eval "$(SHELL=/bin/sh lesspipe)" 32 | 33 | # set variable identifying the chroot you work in (used in the prompt below) 34 | if [ -z "${debian_chroot:-}" ] && [ -r /etc/debian_chroot ]; then 35 | debian_chroot=$(cat /etc/debian_chroot) 36 | fi 37 | 38 | # set a fancy prompt (non-color, unless we know we "want" color) 39 | case "$TERM" in 40 | xterm-color|*-256color) color_prompt=yes;; 41 | esac 42 | 43 | # uncomment for a colored prompt, if the terminal has the capability; turned 44 | # off by default to not distract the user: the focus in a terminal window 45 | # should be on the output of commands, not on the prompt 46 | force_color_prompt=yes 47 | 48 | if [ -n "$force_color_prompt" ]; then 49 | if [ -x /usr/bin/tput ] && tput setaf 1 >&/dev/null; then 50 | # We have color support; assume it's compliant with Ecma-48 51 | # (ISO/IEC-6429). (Lack of such support is extremely rare, and such 52 | # a case would tend to support setf rather than setaf.) 53 | color_prompt=yes 54 | else 55 | color_prompt= 56 | fi 57 | fi 58 | 59 | if [ "$color_prompt" = yes ]; then 60 | PS1='${debian_chroot:+($debian_chroot)}\[\033[01;32m\]\u@\h\[\033[00m\]:\[\033[01;34m\]\w\[\033[00m\]\$ ' 61 | else 62 | PS1='${debian_chroot:+($debian_chroot)}\u@\h:\w\$ ' 63 | fi 64 | unset color_prompt force_color_prompt 65 | 66 | # If this is an xterm set the title to user@host:dir 67 | case "$TERM" in 68 | xterm*|rxvt*) 69 | PS1="\[\e]0;${debian_chroot:+($debian_chroot)}\u@\h: \w\a\]$PS1" 70 | ;; 71 | *) 72 | ;; 73 | esac 74 | 75 | # enable color support of ls and also add handy aliases 76 | if [ -x /usr/bin/dircolors ]; then 77 | test -r ~/.dircolors && eval "$(dircolors -b ~/.dircolors)" || eval "$(dircolors -b)" 78 | alias ls='ls --color=auto' 79 | #alias dir='dir --color=auto' 80 | #alias vdir='vdir --color=auto' 81 | 82 | alias grep='grep --color=auto' 83 | alias fgrep='fgrep --color=auto' 84 | alias egrep='egrep --color=auto' 85 | fi 86 | 87 | # colored GCC warnings and errors 88 | #export GCC_COLORS='error=01;31:warning=01;35:note=01;36:caret=01;32:locus=01:quote=01' 89 | 90 | # some more ls aliases 91 | alias ll='ls -alF' 92 | alias la='ls -A' 93 | alias l='ls -CF' 94 | 95 | # Add an "alert" alias for long running commands. Use like so: 96 | # sleep 10; alert 97 | alias alert='notify-send --urgency=low -i "$([ $? = 0 ] && echo terminal || echo error)" "$(history|tail -n1|sed -e '\''s/^\s*[0-9]\+\s*//;s/[;&|]\s*alert$//'\'')"' 98 | 99 | 100 | # enable programmable completion features (you don't need to enable 101 | # this, if it's already enabled in /etc/bash.bashrc and /etc/profile 102 | # sources /etc/bash.bashrc). 103 | if ! shopt -oq posix; then 104 | if [ -f /usr/share/bash-completion/bash_completion ]; then 105 | . /usr/share/bash-completion/bash_completion 106 | elif [ -f /etc/bash_completion ]; then 107 | . /etc/bash_completion 108 | fi 109 | fi 110 | 111 | 112 | 113 | # Alias definitions. 114 | # You may want to put all your additions into a separate file like 115 | # ~/.bash_aliases, instead of adding them here directly. 116 | # See /usr/share/doc/bash-doc/examples in the bash-doc package. 117 | 118 | # some aesthetic settings 119 | tabs -4 120 | alias dush='du -sh' 121 | alias ll='LC_COLLATE=C ls -hlX --group-directories-first --file-type' 122 | alias lla='LC_COLLATE=C ls -hlXA --group-directories-first --file-type' 123 | alias lls='LC_COLLATE=C ls -hlXS --group-directories-first --file-type' 124 | alias l='LC_COLLATE=C ls -CX --group-directories-first --file-type' 125 | 126 | # rsync utility 127 | alias rsynca='rsync -rltpDvpLK' 128 | 129 | # set PATH so it includes user's private bin if it exists 130 | if [ -d "$HOME/bin" ] ; then 131 | PATH="$HOME/bin:$PATH" 132 | fi 133 | 134 | # set PATH so it includes user's private bin if it exists 135 | if [ -d "$HOME/.local/bin" ] ; then 136 | PATH="$HOME/.local/bin:$PATH" 137 | fi 138 | 139 | # do .sensitive config 140 | if [ -d ~/.sensitive ]; then 141 | for fn in ~/.sensitive/*.bashrc; do 142 | if [ -f $fn ]; then 143 | . $fn 144 | fi 145 | done 146 | fi 147 | 148 | # config project and machine 149 | if [ -f ~/project_config.bashrc ] ; then 150 | . ~/project_config.bashrc 151 | fi 152 | if [ -f ~/machine_config.bashrc ] ; then 153 | . ~/machine_config.bashrc 154 | fi 155 | if [ -f ~/make_config.bashrc ] ; then 156 | . ~/make_config.bashrc 157 | fi 158 | 159 | 160 | 161 | -------------------------------------------------------------------------------- /_scripts/download_vtubers.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | from _util.util_v1 import * ; import _util.util_v1 as uutil 5 | from _util.pytorch_v1 import * ; import _util.pytorch_v1 as utorch 6 | from _util.twodee_v1 import * ; import _util.twodee_v1 as u2d 7 | 8 | 9 | DEBUG = False 10 | 11 | rdn = '/dev/shm' if DEBUG else './_data/lustrous/raw/fandom' 12 | fn_meta = f'{rdn}/virtualyoutuber/metadata.json' 13 | dn_images = mkdir(f'{rdn}/virtualyoutuber/images') 14 | 15 | url_base = 'https://virtualyoutuber.fandom.com' 16 | 17 | 18 | ######## transformation helpers ######## 19 | 20 | # detections 21 | idets = pload(f'./_data/lustrous/preprocessed/facecrop_fandom.pkl') 22 | 23 | # transformations from source image 24 | # hard-coded alignment stats calculated from rutileE 25 | x_mu = 294.9637145996094 26 | x_sig = 8.072593688964844 27 | y_mu = 255.85690307617188 28 | y_sig = 1.85466468334198 29 | a_mu = 13.355611801147461 30 | a_sig = 2.529512643814087 31 | def _get_M(site, franch, idx, view, di, ): 32 | # calc stats 33 | ibn = f'{site}/{franch}/{idx}/{view}' 34 | det = idets[ibn] 35 | dbn = f'{ibn}-{di:04d}' 36 | kpts = det['keypoints'][di,:,:-1] 37 | c = kpts.mean(axis=0) 38 | s = np.linalg.norm(kpts-c, axis=1).std() 39 | with np_seed(dbn): 40 | rands = np.random.normal(size=3) 41 | cnew = np.asarray([ 42 | x_mu + x_sig*rands[0], 43 | y_mu, #+ y_sig*rands[1], 44 | ]) 45 | snew = a_mu #+ a_sig*rands[2] 46 | 47 | # get matrix 48 | sf = snew / s 49 | M = np.asarray([ # recenter to new 50 | [1,0,cnew[0]], 51 | [0,1,cnew[1]], 52 | [0,0,1], 53 | ]) @ np.asarray([ # scale 54 | [sf,0,0], 55 | [0,sf,0], 56 | [0,0,1], 57 | ]) @ np.asarray([ # center to origin 58 | [1,0,-c[0]], 59 | [0,1,-c[1]], 60 | [0,0,1], 61 | ]) 62 | return M 63 | def _apply_M(img, M, size=512): 64 | return I(kornia.geometry.transform.warp_affine( 65 | img.convert('RGBA').bg('w').convert('RGB').t()[None], 66 | torch.tensor(M).float()[[1,0,2]].T[[1,0,2]].T[None,:2], 67 | (size,size), 68 | mode='bilinear', 69 | padding_mode='fill', 70 | align_corners=True, 71 | fill_value=torch.ones(3), 72 | )).alpha_set(I(kornia.geometry.transform.warp_affine( 73 | img['a'].t()[None], 74 | torch.tensor(M).float()[[1,0,2]].T[[1,0,2]].T[None,:2], 75 | (size,size), 76 | mode='bilinear', 77 | padding_mode='fill', 78 | align_corners=True, 79 | fill_value=torch.zeros(3), 80 | ))['r']) 81 | 82 | 83 | ######## load metadata ######## 84 | 85 | bns = uutil.read_bns('./_data/lustrous/subsets/virtualyoutuberE_all.csv') 86 | metas = jread(fn_meta) 87 | asrc2url = {} 88 | for meta in metas: 89 | franch = meta['agency'] 90 | name = meta['name'] 91 | if franch.startswith('.'): 92 | franch = f'dot{franch[1:]}' # fuck this man 93 | if 'portrait' in meta: 94 | asrc2url[f'virtualyoutuber/{franch}/{name}/0000'] = meta['portrait'] 95 | if 'gallery_items' in meta: 96 | for i,item in enumerate(meta['gallery_items']): 97 | asrc2url[f'virtualyoutuber/{franch}/{name}/{i+1:04d}'] = item 98 | 99 | 100 | ######## download images ######## 101 | 102 | ext_allowed = { 103 | 'jpg', 104 | 'png', 105 | 'jpeg', 106 | 'gif', 107 | 'webp', 108 | } 109 | def _get_image_url(c): 110 | if 'src' not in c: return None 111 | if '/revision' not in c['src']: return None 112 | url = c['src'].split('/revision')[0] 113 | ext = fnstrip(url,1).ext.lower() 114 | if ext not in ext_allowed: return None 115 | return Dict(url=url, ext=ext) 116 | for bn in bns: 117 | try: 118 | # a = aligndata[bn] 119 | # asrc = a['source'] 120 | # m = asrc2url[a['source']] 121 | # _,franch,name,idx = asrc.split('/') 122 | # idx = int(idx) 123 | _,_,franch,name,idx_di = bn.split('/') 124 | idx = int(idx_di.split('-')[0]) 125 | m = asrc2url[f'virtualyoutuber/{franch}/{name}/{idx:04d}'] 126 | opn = f'{dn_images}/{franch}/{name}' 127 | 128 | if idx==0: 129 | # download portrait 130 | iurl = _get_image_url(m) 131 | if iurl==None: continue 132 | ofn = f'{opn}/0000.{iurl.ext}' 133 | if os.path.isfile(ofn): continue 134 | try: 135 | urllib.request.urlretrieve(iurl.url, mkfile(ofn)) 136 | except Exception as e: 137 | tbs = traceback.format_exc() 138 | write(f'{isonow()}\n{tbs}', mkfile(f'{ofn}.txt')) 139 | else: 140 | # download from gallery 141 | iurl = _get_image_url(m) 142 | if iurl==None: continue 143 | ofn = f'{opn}/{idx:04d}.{iurl.ext}' 144 | if os.path.isfile(ofn): continue 145 | try: 146 | urllib.request.urlretrieve(iurl.url, mkfile(ofn)) 147 | except Exception as e: 148 | tbs = traceback.format_exc() 149 | write(f'{isonow()}\n{tbs}', mkfile(f'{ofn}.txt')) 150 | # break 151 | except: 152 | print(f'skipped: {bn}') 153 | 154 | 155 | ######## transform + segment images ######## 156 | 157 | for bn in bns: 158 | try: 159 | _,_,franch,name,idx_di = bn.split('/') 160 | idx = int(idx_di.split('-')[0]) 161 | fan,franch,idx,view = f'virtualyoutuber/{franch}/{name}/{idx:04d}'.split('/') 162 | img_src = I(f'{rdn}/{fan}/images/{franch}/{idx}/{view}.png') 163 | img_seg = I(f'./_data/lustrous/renders/{bn.replace("fandom_align","fandom_align_seg")}.png') 164 | M = _get_M(site, franch, idx, view, di) 165 | out = _apply_M(img_src, M).alpha_set(img_seg) 166 | if DEBUG: 167 | out.save(mkfile(f'/dev/shm/renders/{bn}.png')) 168 | else: 169 | out.save(mkfile(f'./_data/lustrous/renders/{bn}.png')) 170 | except: 171 | print(f'skipped: {bn}') 172 | 173 | 174 | 175 | 176 | 177 | -------------------------------------------------------------------------------- /_util/pytorch_v1.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | try: 5 | import _util.util_v1 as uutil 6 | except: 7 | pass 8 | 9 | 10 | try: 11 | import torch 12 | import torch.nn as nn 13 | except: 14 | pass 15 | 16 | try: 17 | import torchvision as tv 18 | import torchvision.transforms as TT 19 | import torchvision.transforms.functional as TF 20 | except: 21 | pass 22 | 23 | try: 24 | import pytorch_lightning as pl 25 | import pytorch_lightning.strategies as _ 26 | import torchmetrics 27 | import torchmetrics.image 28 | except: 29 | pass 30 | 31 | try: 32 | import wandb 33 | except: 34 | pass 35 | 36 | try: 37 | import kornia 38 | except: 39 | pass 40 | 41 | try: 42 | import einops 43 | from einops.layers import torch as _ 44 | except: 45 | pass 46 | 47 | try: 48 | import cupy 49 | except: 50 | pass 51 | 52 | try: 53 | import lpips 54 | except: 55 | pass 56 | 57 | try: 58 | from addict import Dict 59 | except: 60 | Dict = dict 61 | 62 | 63 | #################### UTILITIES #################### 64 | 65 | 66 | import contextlib, threading 67 | @contextlib.contextmanager 68 | def torch_seed(seed): 69 | _torch_seed_lock.acquire() 70 | state = torch.get_rng_state() 71 | was_det = torch.backends.cudnn.deterministic 72 | if seed!=None and not isinstance(seed, int): 73 | seed = zlib.adler32(bytes(seed, encoding='utf-8')) 74 | elif seed==None: 75 | seed = torch.seed() 76 | torch.manual_seed(seed) 77 | torch.cuda.manual_seed(seed) 78 | torch.backends.cudnn.deterministic = True 79 | try: 80 | yield 81 | finally: 82 | torch.set_rng_state(state) 83 | torch.backends.cudnn.deterministic = was_det 84 | _torch_seed_lock.release() 85 | _torch_seed_lock = threading.Lock() 86 | 87 | def torch_seed_all(seed): 88 | _torch_seed_lock.acquire() 89 | if seed!=None and not isinstance(seed, int): 90 | seed = zlib.adler32(bytes(seed, encoding='utf-8')) 91 | elif seed==None: 92 | seed = torch.seed() 93 | torch.manual_seed(seed) 94 | torch.cuda.manual_seed(seed) 95 | torch.backends.cudnn.deterministic = True 96 | _torch_seed_lock.release() 97 | return 98 | 99 | 100 | # @cupy.memoize(for_each_device=True) 101 | def cupy_launch(func, kernel): 102 | return cupy.cuda.compile_with_cache(kernel).get_function(func) 103 | 104 | def reset_parameters(model): 105 | for layer in model.children(): 106 | if hasattr(layer, 'reset_parameters'): 107 | layer.reset_parameters() 108 | return model 109 | 110 | def default_collate(items, device=None): 111 | return to(Dict(torch.utils.data.dataloader.default_collate(items)), device) 112 | 113 | def to(x, device): 114 | if device is None: 115 | return x 116 | if issubclass(x.__class__, dict): 117 | return Dict({ 118 | k: v.to(device) if isinstance(v, torch.Tensor) else v 119 | for k,v in x.items() 120 | }) 121 | if isinstance(x, torch.Tensor): 122 | return x.to(device) 123 | if isinstance(x, np.ndarray): 124 | return torch.tensor(x).to(device) 125 | assert 0, 'data not understood' 126 | 127 | 128 | try: 129 | class IdentityModule(nn.Module): 130 | def __init__(self, *args, **kwargs): 131 | super().__init__() 132 | return 133 | def forward(self, x, *args, **kwargs): 134 | return x 135 | class Tanh10Module(nn.Module): 136 | def __init__(self, *args, **kwargs): 137 | super().__init__() 138 | return 139 | def forward(self, x, *args, **kwargs): 140 | return 10*torch.tanh(x/10) 141 | except: pass 142 | 143 | 144 | def tsuma(t, f=' 2.04f'): 145 | with torch.no_grad(): 146 | s = f' shape ({",".join([str(i) for i in t.shape])})\n' 147 | return s+str(uutil.Table([ 148 | # ['shape::l', ('('+','.join([str(i) for i in t.shape])+')', 'r'), ], 149 | ['min::l', (t.min().item(), f'r:{f}'), ' ', 'std::l', (t.std().item(), f'r:{f}'), ], 150 | ['mean::l', (t.mean().item(), f'r:{f}'), ' ', 'norm::l', (t.norm().item(), f'r:{f}'), ], 151 | ['max::l', (t.max().item(), f'r:{f}'), ], 152 | # ['std::l', (t.std().item(), f'r:{f}'), ], 153 | # ['norm::l', (t.norm().item(), f'r:{f}'), ], 154 | ]))+'\n' 155 | 156 | 157 | #################### METRICS #################### 158 | 159 | class LPIPSLoss(nn.Module): 160 | def __init__(self, net_type='alex', **kwargs): 161 | super().__init__() 162 | self.net_type = net_type 163 | assert self.net_type in ['alex', 'vgg', 'squeeze'] 164 | self.model = lpips.LPIPS(net=self.net_type, **kwargs) 165 | return 166 | def forward(self, preds: torch.Tensor, target: torch.Tensor): 167 | ans = self.model(preds, target).mean((1,2,3)) 168 | return ans 169 | 170 | class LaplacianPyramidLoss(nn.Module): 171 | def __init__(self, n_levels=3, colorspace=None, mode='l1'): 172 | super().__init__() 173 | self.n_levels = n_levels 174 | self.colorspace = colorspace 175 | self.mode = mode 176 | assert self.mode in ['l1', 'l2'] 177 | return 178 | def forward(self, preds, target, force_levels=None, force_mode=None): 179 | if self.colorspace=='lab': 180 | preds = kornia.color.rgb_to_lab(preds.float()) 181 | target = kornia.color.rgb_to_lab(target.float()) 182 | lvls = self.n_levels if force_levels==None else force_levels 183 | preds = kornia.geometry.transform.build_pyramid(preds, lvls) 184 | target = kornia.geometry.transform.build_pyramid(target, lvls) 185 | mode = self.mode if force_mode==None else force_mode 186 | if mode=='l1': 187 | ans = torch.stack([ 188 | (p-t).abs().mean((1,2,3)) 189 | for p,t in zip(preds,target) 190 | ]).mean(0) 191 | elif mode=='l2': 192 | ans = torch.stack([ 193 | (p-t).norm(dim=1, keepdim=True).mean((1,2,3)) 194 | for p,t in zip(preds,target) 195 | ]).mean(0) 196 | else: 197 | assert 0 198 | return ans 199 | 200 | def binclass_metrics(pred, gt, dims=(1,2,3)): 201 | assert pred.dtype==torch.bool 202 | assert gt.dtype==torch.bool 203 | 204 | tp = (pred & gt).sum(dims) 205 | acc = (pred==gt).float().mean(dims) 206 | pre = tp / pred.sum(dims) 207 | rec = tp / gt.sum(dims) 208 | f1 = 2*pre*rec/(pre+rec) 209 | 210 | pzero = pred.sum(dims)==0 211 | gzero = gt.sum(dims)==0 212 | pneqg = pzero!=gzero 213 | pgeqz = pzero & gzero 214 | f1[pneqg | ((pre==0)&(rec==0))] = 0 215 | pre[pneqg] = 0 216 | rec[pneqg] = 0 217 | pre[pgeqz] = 1 218 | rec[pgeqz] = 1 219 | f1[pgeqz] = 1 220 | 221 | return Dict({ 222 | 'f1': f1, 223 | 'precision': pre, 224 | 'recall': rec, 225 | 'accuracy': acc, 226 | }) 227 | 228 | 229 | 230 | 231 | -------------------------------------------------------------------------------- /_util/util_v1.py: -------------------------------------------------------------------------------- 1 | 2 | 3 | 4 | 5 | ################ STANDARD ################ 6 | 7 | import os 8 | import sys 9 | import stat 10 | import traceback 11 | import gc 12 | import io 13 | import shutil 14 | import subprocess 15 | import struct 16 | import sqlite3 17 | import logging 18 | import warnings 19 | import pathlib 20 | import importlib 21 | import tempfile 22 | import inspect 23 | import math 24 | import random 25 | import re 26 | import copy 27 | import threading 28 | import multiprocessing 29 | import urllib 30 | import base64 31 | import zlib 32 | import functools 33 | import time 34 | import contextlib 35 | import typing 36 | import argparse 37 | from argparse import Namespace 38 | from pprint import pprint 39 | from collections import OrderedDict, defaultdict, Counter 40 | from collections.abc import Iterable 41 | from datetime import datetime 42 | 43 | try: 44 | import requests 45 | except: 46 | pass 47 | try: 48 | import psutil 49 | except: 50 | pass 51 | try: 52 | from addict import Dict 53 | except: 54 | Dict = dict 55 | try: 56 | import pyunpack 57 | import zipfile 58 | import tarfile 59 | except: 60 | pass 61 | 62 | 63 | ################ SCIENTIFIC ################ 64 | 65 | try: 66 | import numpy as np 67 | except: 68 | pass 69 | try: 70 | import matplotlib.pyplot as plt 71 | except: 72 | pass 73 | try: 74 | import pandas as pd 75 | except: 76 | pass 77 | try: 78 | import scipy 79 | import scipy.spatial 80 | except: 81 | pass 82 | try: 83 | import sklearn 84 | from sklearn import cluster as _ 85 | from sklearn import decomposition as _ 86 | from sklearn import metrics as _ 87 | from sklearn import neighbors as _ 88 | except: 89 | pass 90 | try: 91 | import PIL 92 | from PIL import Image 93 | ImageFile.LOAD_TRUNCATED_IMAGES = True 94 | except: 95 | pass 96 | try: 97 | import cv2 98 | except: 99 | pass 100 | 101 | @contextlib.contextmanager 102 | def np_seed(seed): 103 | _np_seed_lock.acquire() 104 | state = np.random.get_state() 105 | if seed!=None and not isinstance(seed, int): 106 | seed = zlib.adler32(bytes(seed, encoding='utf-8')) 107 | np.random.seed(seed) 108 | try: 109 | yield 110 | finally: 111 | np.random.set_state(state) 112 | _np_seed_lock.release() 113 | _np_seed_lock = threading.Lock() 114 | 115 | def np_seed_all(seed): 116 | _np_seed_lock.acquire() 117 | if seed!=None and not isinstance(seed, int): 118 | seed = zlib.adler32(bytes(seed, encoding='utf-8')) 119 | np.random.seed(seed) 120 | _np_seed_lock.release() 121 | return 122 | 123 | def np_choice_strat(length, size, shuffle=True): 124 | ans = np.concatenate([ 125 | np.tile(np.arange(length), size//length), 126 | np.random.choice(length, size=size%length, replace=False), 127 | ]) 128 | if shuffle: np.random.shuffle(ans) 129 | return ans 130 | 131 | 132 | ################ FILES ################ 133 | 134 | def mkfile(fn, parents=True, exist_ok=True): 135 | dn = '/'.join(fn.split('/')[:-1]) 136 | mkdir(dn, parents=parents, exist_ok=exist_ok) 137 | return fn 138 | def mkdir(dn, parents=True, exist_ok=True): 139 | pathlib.Path(dn).mkdir(parents=parents, exist_ok=exist_ok) 140 | return dn if (not dn[-1]=='/' or dn=='/') else dn[:-1] 141 | def fnstrip(fn, return_more=False): 142 | dspl = fn.split('/') 143 | dn = '/'.join(dspl[:-1]) if len(dspl)>1 else '.' 144 | fn = dspl[-1] 145 | fspl = fn.split('.') 146 | if len(fspl)==1: 147 | bn = fspl[0] 148 | ext = '' 149 | else: 150 | bn = '.'.join(fspl[:-1]) 151 | ext = fspl[-1] 152 | if return_more: 153 | return Namespace( 154 | dn=dn, 155 | fn=fn, 156 | path=f'{dn}/{fn}', 157 | bn_path=f'{dn}/{bn}', 158 | bn=bn, 159 | ext=ext, 160 | ) 161 | else: 162 | return bn 163 | def patmat(bn, pat, junk=False, strict=True): 164 | # './_data/thing/dtype/asdfg/123-4567890.png', 165 | # './_data/thing/dtype/{sub}/{a}-{b:int}.png', 166 | # 'asdfkjlj,123/stuff/238b9', 167 | # '{a},{b:float}/stuff/{c}', 168 | # '*,{b:float}.{d}/{}/{c}', 169 | if pat is None: 170 | return Dict({0: bn}) if junk else Dict() 171 | try: 172 | bn = bn + '?' 173 | pat = pat + '?' 174 | pat = pat.replace('*', '{}') 175 | t = pat.split('{') 176 | s = [(t[0], None), ] # exact match 177 | for q in t[1:]: 178 | a,b = q.split('}') 179 | a = a.split(':') 180 | if len(a)==1: 181 | a,f = a[0], str 182 | elif len(a)==2: 183 | a,f = a[0], { 184 | 'int': int, 185 | 'float': float, 186 | 'str': str, 187 | }[a[1]] 188 | else: 189 | assert 0 190 | s.append((a, f)) 191 | s.append((b, None)) 192 | ans = Dict() 193 | i = 0 194 | mc = 0 195 | for si,(q,f) in enumerate(s): 196 | if f is None: 197 | i += len(q) 198 | else: 199 | n = bn[i:].index(s[si+1][0]) 200 | v = f(bn[i:i+n]) 201 | if q=='': 202 | if junk: 203 | ans[mc] = v 204 | mc += 1 205 | else: 206 | ans[q] = v 207 | i += n 208 | return ans 209 | except: 210 | if not strict: 211 | return None 212 | else: 213 | assert 0, f'"{bn[:-1]}" does not match "{pat[:-1]}"' 214 | def walk(dn, pat, junk=False, strip=False): 215 | ans = [] 216 | for p,ds,fs in os.walk(dn): 217 | for f in ds+fs: 218 | f = f'{p}/{f}' 219 | q = patmat(f[len(dn)+1:], pat, junk=junk, strict=False) 220 | if q: 221 | if strip: q['fnstrip'] = Dict(vars(fnstrip(f,1))) 222 | ans.append((f, q)) 223 | return sorted(ans, key=lambda x: x[0]) 224 | 225 | # r: open for reading (default) 226 | # w: open for writing, truncating the file first 227 | # x: open for exclusive creation, failing if the file already exists 228 | # a: open for writing, appending to the end of file if it exists 229 | # b: binary mode 230 | # t: text mode (default) 231 | # +: open for updating (reading and writing) 232 | def read(fn, mode='r', buffering=-1, encoding=None, 233 | errors=None, newline=None, closefd=True, opener=None): 234 | with open(fn, mode=mode, buffering=buffering, encoding=encoding, 235 | errors=errors, newline=newline, closefd=closefd, opener=opener) as handle: 236 | return handle.read() 237 | def write(text, fn, mode='w', buffering=-1, encoding=None, 238 | errors=None, newline=None, closefd=True, opener=None): 239 | with open(fn, mode=mode, buffering=buffering, encoding=encoding, 240 | errors=errors, newline=newline, closefd=closefd, opener=opener) as handle: 241 | return handle.write(text) 242 | def read_bns(fn, safe=False, sort=False, sort_key=None): 243 | ans = [ 244 | line 245 | for line in read(fn).split('\n') 246 | if line!='' 247 | ] 248 | if sort: 249 | ans = sorted(ans, key=sort_key) 250 | if safe: 251 | return safe_bns(ans) 252 | else: 253 | return ans 254 | def safe_bns(bns): 255 | return np.array(bns, dtype=np.string_) 256 | def safe_bn(bn): 257 | return np.string_(bn) 258 | def unsafe_bns(bns): 259 | return [unsafe_bn(q) for q in bns] 260 | def unsafe_bn(bn, bns=None): 261 | if np.issubdtype(type(bn), np.integer): 262 | assert bns is not None 263 | return unsafe_bn(bns[bn]) 264 | elif isinstance(bn, np.bytes_): 265 | return str(bn, encoding='utf-8') 266 | elif isinstance(bn, str): 267 | return bn 268 | else: 269 | assert 0, f'bn not understood' 270 | 271 | import pickle 272 | def pload(fn, mode='rb'): 273 | with open(fn, mode) as handle: 274 | return pickle.load(handle) 275 | def pdump(obj, fn, mode='wb'): 276 | with open(fn, mode) as handle: 277 | return pickle.dump(obj, handle) 278 | 279 | import json 280 | def jread(fn, mode='r'): 281 | with open(fn, mode) as handle: 282 | return json.load(handle) 283 | def jwrite(x, fn, mode='w', indent='\t', sort_keys=False): 284 | with open(fn, mode) as handle: 285 | return json.dump(x, handle, indent=indent, sort_keys=sort_keys) 286 | 287 | try: 288 | import yaml 289 | def yread(fn, mode='r'): 290 | with open(fn, mode) as handle: 291 | return yaml.safe_load(handle) 292 | def ywrite(x, fn, mode='w', default_flow_style=False): 293 | if dict!=Dict: 294 | x = Dict(top=x).to_dict()['top'] 295 | with open(fn, mode) as handle: 296 | yaml.dump(x, handle, default_flow_style=default_flow_style) 297 | except: 298 | pass 299 | 300 | 301 | ################ MISC ################ 302 | 303 | try: 304 | hakase = f'{os.environ["PROJECT_DN"]}/_env/__hakase__.jpg' 305 | if not os.path.isfile(hakase): 306 | hakase = f'{os.environ["PROJECT_DN"]}/env/__hakase__.jpg' 307 | if not os.path.isfile(hakase): 308 | hakase = f'{os.environ["PROJECT_DN"]}/__env__/__hakase__.jpg' 309 | if not os.path.isfile(hakase): 310 | hakase = None 311 | except: 312 | hakase = None 313 | 314 | enum = enumerate 315 | 316 | def isonow(): 317 | return datetime.now().isoformat().replace(':','-').replace('.','-') 318 | 319 | def mem(units='m'): 320 | return psutil.Process(os.getpid()).memory_info().rss / { 321 | 'b': 1, 322 | 'k': 1e3, 323 | 'm': 1e6, 324 | 'g': 1e9, 325 | 't': 1e12, 326 | }[units[0].lower()] 327 | def gpumem(units='m', query='used'): 328 | assert query in ['used', 'free', 'total'] 329 | out = subprocess.check_output( 330 | f'nvidia-smi --query-gpu=memory.{query} --format=csv'.split() 331 | ).decode('ascii').split('\n')[:-1][1:] 332 | ans = [int(x.split()[0]) for i, x in enumerate(out)][0] * { 333 | 'b': 1e6, 334 | 'k': 1e3, 335 | 'm': 1, 336 | 'g': 1e-3, 337 | 't': 1e-9, 338 | }[units[0].lower()] 339 | return ans 340 | 341 | def chunk_cols(x, n): 342 | return [x[i:i+n] for i in range(0, len(x), n)] 343 | def chunk_rows(x, n): 344 | return chunk_cols(x, int(math.ceil(len(x)/n))) 345 | 346 | def idxs2list(idxs, n=None): 347 | if isinstance(idxs, list) or isinstance(idxs, tuple): 348 | return idxs 349 | elif isinstance(idxs, slice) or isinstance(idxs, range): 350 | a,b,c = idxs.start, idxs.stop, idxs.step 351 | if a is None: a = 0 352 | if b is None: b = n 353 | if c is None: c = 1 354 | if a<0: a = len(self)+a 355 | if b<0: b = len(self)+b 356 | return list(range(a,b,c)) 357 | assert 0, 'data not understood' 358 | 359 | 360 | ################ AESTHETIC ################ 361 | 362 | try: 363 | from tqdm import tqdm as tqdm_ 364 | from tqdm import trange as trange_ 365 | from tqdm.auto import tqdm, trange 366 | except: 367 | pass 368 | 369 | class Table: 370 | def __init__(self, 371 | table, 372 | delimiter=' ', 373 | orientation='br', 374 | double_colon=True, 375 | ): 376 | self.delimiter = delimiter 377 | self.orientation = orientation 378 | self.t = Table.parse( 379 | table, delimiter, orientation, double_colon, 380 | ) 381 | return 382 | 383 | # rendering 384 | def __str__(self): 385 | return self.render() 386 | def __repr__(self): 387 | return self.render() 388 | def render(self): 389 | # set up empty entry 390 | empty = ('', Table._spec(self.orientation, transpose=False)) 391 | 392 | # calculate table size 393 | t = copy.deepcopy(self.t) 394 | totalrows = len(t) 395 | totalcols = [len(r) for r in t] 396 | assert min(totalcols)==max(totalcols) 397 | totalcols = totalcols[0] 398 | 399 | # string-ify 400 | for i in range(totalrows): 401 | for j in range(totalcols): 402 | x,s = t[i][j] 403 | sp = s[11] 404 | if sp: x = eval(f'f"{{{x}{sp}}}"') 405 | Table._put((str(x),s), t, (i,j), empty) 406 | 407 | # expand delimiters 408 | _repl = lambda s: \ 409 | s[:2] + (1,0,0,0,0) + s[7:10] + (1,) + s[11:] \ 410 | if s[2] else \ 411 | s[:2] + (0,0,0,0,0) + s[7:10] + (1,) + s[11:] 412 | for i,row in enumerate(t): 413 | for j,(x,s_own) in enumerate(row): 414 | # expand delim_up(^) 415 | if s_own[3]: 416 | u,v = i,j 417 | while 0<=u: 418 | _,s = t[u][v] 419 | if (i,j)!=(u,v) and (s[2] and not s[10]): break 420 | Table._put((x, _repl(s)), t, (u,v), empty) 421 | u -= 1 422 | 423 | # expand delim_down(v) 424 | if s_own[4]: 425 | u,v = i,j 426 | while u) 433 | if s_own[5]: 434 | u,v = i,j 435 | while v height=1 463 | heights = [h+1 for h in heights] 464 | 465 | # render table 466 | rend = [] 467 | roff = 0 468 | for i,row in enumerate(t): 469 | for j,(x,s) in enumerate(row): 470 | w,h = widths[j], heights[i] 471 | 472 | # expand fillers and delimiters 473 | if s[2] or s[10]: 474 | xs = x.split('\n') 475 | xw0 = min(len(l) for l in xs) 476 | xw1 = max(len(l) for l in xs) 477 | xh = len(xs) 478 | if (xw0==xw1==w) and (xh==h): 479 | pass 480 | elif xw0==xw1==w: 481 | x = '\n'.join([xs[0],]*h) 482 | elif xh==h: 483 | x = '\n'.join([(l[0] if l else '')*w for l in xs]) 484 | else: 485 | x = x[0] if x else ' ' 486 | x = '\n'.join([x*w,]*h) 487 | 488 | # justify horizontally 489 | x = [ 490 | l.rjust(w) if s[0] else l.ljust(w) 491 | for l in x.split('\n') 492 | ] 493 | 494 | # justify vertically 495 | plus = [' '*w,]*(h-len(x)) 496 | x = plus+x if not s[1] else x+plus 497 | 498 | # input to table 499 | for r,xline in enumerate(x): 500 | Table._put(xline, rend, (roff+r,j), None) 501 | roff += h 502 | 503 | # return rendered string 504 | return '\n'.join([''.join(r) for r in rend]) 505 | 506 | # parsing 507 | def _spec(s, transpose=False): 508 | if ':' in s: 509 | i = s.index(':') 510 | sp = s[i:] 511 | s = s[:i] 512 | else: 513 | sp = '' 514 | s = s.lower() 515 | return ( 516 | int('r' in s), # 0:: 0:left(l) 1:right(r) 517 | int('t' in s), # 1:: 0:bottom(b) 1:top(t) 518 | int(any([i in s for i in ['.','<','>','^','v']])), # 2:: delim_here(.) 519 | int('^' in s if not transpose else '<' in s), # 3:: delim_up(^) 520 | int('v' in s if not transpose else '>' in s), # 4:: delim_down(v) 521 | int('>' in s if not transpose else 'v' in s), # 5:: delim_right(>) 522 | int('<' in s if not transpose else '^' in s), # 6:: delim_left(<) 523 | int('+' in s), # 7:: subtable(+) 524 | int('-' in s if not transpose else '|' in s), # 8:: subtable_horiz(-) 525 | int('|' in s if not transpose else '-' in s), # 9:: subtable_vert(|) 526 | int('_' in s), # 10:: fill(_); if delim, overwrite; else fit 527 | sp, # 11:: special(:) f-string for numbers 528 | ) 529 | def _put(obj, t, ij, empty): 530 | i,j = ij 531 | while i>=len(t): 532 | t.append([]) 533 | while j>=len(t[i]): 534 | t[i].append(empty) 535 | t[i][j] = obj 536 | return 537 | def parse( 538 | table, 539 | delimiter=' ', 540 | orientation='br', 541 | double_colon=True, 542 | ): 543 | # disabling transpose 544 | transpose = False 545 | 546 | # set up empty entry 547 | empty = ('', Table._spec(orientation, transpose)) 548 | 549 | # transpose 550 | t = [] 551 | for i,row in enumerate(table): 552 | for j,item in enumerate(row): 553 | ij = (i,j) if not transpose else (j,i) 554 | if type(item)==tuple and len(item)==2 and type(item[1])==str: 555 | item = (item[0], Table._spec(item[1], transpose)) 556 | elif double_colon and type(item)==str and '::' in item: 557 | x,s = item.split('::') 558 | item = (x, Table._spec(s, transpose)) 559 | else: 560 | item = (item, Table._spec(orientation, transpose)) 561 | Table._put(item, t, ij, empty) 562 | 563 | # normalization 564 | maxcol = 0 565 | maxrow = len(t) 566 | for i,row in enumerate(t): 567 | # take element number into account 568 | maxcol = max(maxcol, len([i for i in row if not i[1][2]])) 569 | 570 | # take subtables into account 571 | for j,(x,s) in enumerate(row): 572 | if s[7]: 573 | r = len(x) 574 | maxrow = max(maxrow, i+r) 575 | c = max(len(q) for q in x) 576 | maxcol = max(maxcol, j+c) 577 | elif s[8]: 578 | c = len(x) 579 | maxcol = max(maxcol, j+c) 580 | elif s[9]: 581 | r = len(x) 582 | maxrow = max(maxrow, i+r) 583 | totalcols = 2*maxcol + 1 584 | totalrows = maxrow 585 | t += [[]]*(totalrows-len(t)) 586 | newt = [] 587 | delim = (delimiter, Table._spec('._'+orientation, transpose)) 588 | for i,row in enumerate(t): 589 | wasd = False 590 | tcount = 0 591 | for j in range(totalcols): 592 | item = t[i][tcount] if tcount thresh).nonzero() 301 | y = (a.amax(dim=0) > thresh).nonzero() 302 | if len(x)==0 or len(y)==0: 303 | return None 304 | ans = t[:, x[0]:x[-1]+1][...,y[0]:y[-1]+1] 305 | return I(ans) 306 | 307 | # resizing 308 | def resize(self, s, resample='nearest', dry=False): 309 | if dry: 310 | return Namespace( 311 | size_original=self.size, 312 | corner_from=(0,0), 313 | size_from=self.size, 314 | size_to=s, 315 | ) 316 | if self.dtype=='pil': 317 | s = I.sizer(self.resize(s, dry=True).size_to) 318 | return I(self.data.resize( 319 | s[::-1], resample=getattr(Image, resample.upper()), 320 | )) 321 | else: 322 | return I(self.pil()).resize(s, resample=resample) 323 | def rescale(self, factor, resample='nearest', dry=False): 324 | if dry: 325 | return Namespace( 326 | size_original=self.size, 327 | corner_from=(0,0), 328 | size_from=self.size, 329 | size_to=tuple(factor*s for s in self.size), 330 | ) 331 | else: 332 | return self.resize( 333 | self.rescale(factor, dry=True).size_to, 334 | resample=resample, 335 | ) 336 | def rmax(self, s=512, resample='nearest', dry=False): 337 | if dry: 338 | h,w = self.size 339 | return Namespace( 340 | size_original=(h,w), 341 | corner_from=(0,0), 342 | size_from=(h,w), 343 | size_to=(s, w*s/h) if w0 else 0 451 | oy = offy[r,c-1] if c>0 else 0 452 | ans.paste(item.pil().convert('RGBA'), (int(oy),int(ox))) 453 | return I(ans) 454 | 455 | # color 456 | _COLOR_DICT = { 457 | 'r': (1,0,0), 458 | 'g': (0,1,0), 459 | 'b': (0,0,1), 460 | 'k': 0, 461 | 'w': 1, 462 | 't': (0,1,1), 463 | 'm': (1,0,1), 464 | 'y': (1,1,0), 465 | 'a': (0,0,0,0), 466 | } 467 | @staticmethod 468 | def c255(c, c255=True): 469 | # color format utility 470 | # returns RGBA as uint tuple 471 | if c is None: 472 | return None 473 | if isinstance(c, str): 474 | c = I._COLOR_DICT[c] 475 | if isinstance(c, list) or isinstance(c, tuple): 476 | if len(c)==3: 477 | c = c + (1,) 478 | elif len(c)==1: 479 | c = (c,c,c,1) 480 | c = tuple(int(255*q) for q in c) 481 | else: 482 | c = int(255*c) 483 | c = (c,c,c,255) 484 | if not c255: 485 | c = tuple(ch/255.0 for ch in c) 486 | return c 487 | @staticmethod 488 | def uniform_colors(n, seed=None, lightness=0.5, saturation=0.5, roll=0, c255=False): 489 | if seed is not None: 490 | with uutil.numpy_seed(seed): 491 | perm = np.random.permutation(n) 492 | else: 493 | perm = np.arange(n) 494 | return [ 495 | I.c255(colorsys.hls_to_rgb(hue, lightness, saturation), c255=c255) 496 | for hue in np.modf(np.linspace(0, 1, n, endpoint=False)+roll)[0][perm] 497 | ] 498 | 499 | # compositing 500 | @staticmethod 501 | def blank(size, c='k', rounding='floor'): 502 | size = I.sizer(size, rounding=rounding) 503 | c = I.c255(c) 504 | return I( 505 | np.ones((4,*size)) * np.asarray(c)[:,None,None]/255.0 506 | ) 507 | def channel_set(self, ch, val=0.0): 508 | d = self.torch().clone() 509 | d[ch] = val 510 | return I(d) 511 | def alpha_set(self, alpha): 512 | alpha = I(alpha) 513 | assert alpha.mode=='L' 514 | alpha = alpha.np() 515 | rgb = (self.convert('RGB') if self.mode=='L' else self).np()[:3] 516 | return I(np.concatenate([rgb, alpha])) 517 | def alpha_composite(self, img, opacity=1): 518 | if img is None or opacity==0: 519 | return I(self) 520 | img = I(img) 521 | assert img.mode=='RGBA' 522 | if opacity!=1: 523 | img = img.np().copy() 524 | img[3] *= opacity 525 | img = I(img) 526 | return I(Image.alpha_composite(self.convert('RGBA').pil(), img.pil())) 527 | def acomp(self, *args, **kwargs): 528 | return self.alpha_composite(*args, **kwargs) 529 | def background_composite(self, img, opacity=1): 530 | return I(img).alpha_composite(self, opacity=opacity) 531 | def bcomp(self, *args, **kwargs): 532 | return self.background_composite(*args, **kwargs) 533 | def background(self, c=0.5): 534 | return I.blank(self.size, c=c).alpha_composite(self) 535 | def bg(self, *args, **kwargs): 536 | return self.background(*args, **kwargs) 537 | def l2rgba(self, c='k'): 538 | # turns bw to solid color, using L as alpha 539 | assert self.mode=='L' 540 | return I.blank(self.size, c=c).alpha_set(self.np()) 541 | @staticmethod 542 | def average(imgs): 543 | imgs = [I(i) for i in imgs] 544 | assert len(set([i.mode for i in imgs]))==1 545 | return I(np.mean(np.stack([i.np() for i in imgs]), axis=0)) 546 | @staticmethod 547 | def avg(*args, **kwargs): 548 | return I.average(*args, **kwargs) 549 | @staticmethod 550 | def mean(*args, **kwargs): 551 | return I.average(*args, **kwargs) 552 | 553 | # drawing 554 | def line(self, a, b, w=1, c='r'): 555 | a = I.sizer(a, rounding='floor') 556 | b = I.sizer(b, rounding='floor') 557 | c = I.c255(c) 558 | w = max(1, math.floor(w)) 559 | ans = self.convert('RGBA').pil().copy() 560 | d = PIL.ImageDraw.Draw(ans) 561 | d.line([a[::-1], b[::-1]], fill=c, width=w) 562 | return I(ans) 563 | def point(self, xy, s=1, c='r'): 564 | c = I.c255(c) 565 | x,y = I.sizer(xy, rounding='floor') 566 | ans = self.convert('RGBA') 567 | if s==0: 568 | ans = ans.uint8() 569 | ans[:,x,y] = c 570 | ans = I(ans) 571 | else: 572 | ans = ans.pil().copy() 573 | d = PIL.ImageDraw.Draw(ans) 574 | d.ellipse( 575 | [(y-s,x-s), (y+s,x+s)], 576 | fill=c, 577 | ) 578 | return I(ans) 579 | def box(self, corner, size, w=1, c='r', f=None): 580 | corner = I.sizer(corner, rounding='floor') 581 | size = I.sizer(size, rounding='floor') 582 | w = max(1, math.floor(w)) 583 | c = I.c255(c) 584 | f = I.c255(f) 585 | ans = self.convert('RGBA').pil().copy() 586 | d = PIL.ImageDraw.Draw(ans) 587 | d.rectangle( 588 | [corner[1], corner[0], corner[1]+size[1]-1, corner[0]+size[0]-1], 589 | fill=f, outline=c, width=w, 590 | ) 591 | return I(ans) 592 | def border(self, w=1, c='r', pad=False): 593 | i,j = self.size 594 | ans = self if not pad else self.crop((-w,-w), (i+2*w,j+2*w)) 595 | return ans.box( 596 | (0, 0), 597 | ans.size, 598 | w=w, c=c, f=None, 599 | ) 600 | def borderp(self, w=1, c=0.5, pad=True): 601 | return self.border(w=w, c=c, pad=pad) 602 | 603 | # text 604 | def annotate(self, text, pos, s=12, anchor='tl', c='m', bg='k', spacing=None, padding=0): 605 | t = I.text( 606 | text, s=s, c=c, bg=bg, 607 | spacing=spacing, padding=padding, 608 | ) 609 | x,y = pos 610 | x = { 611 | 't': x, 612 | 'b': x-t.size[0], 613 | 'c': x-t.size[0]/2, 614 | }[anchor[0].lower()] 615 | y = { 616 | 'l': y, 617 | 'r': y-t.size[1], 618 | 'c': y-t.size[1]/2, 619 | }[anchor[1].lower()] 620 | t = t.convert('RGBA').pil() 621 | ans = self.convert('RGBA').pil().copy() 622 | ans.paste(t, I.sizer((y,x), rounding='floor'), t) 623 | return I(ans) 624 | def annot(self, *args, **kwargs): 625 | return self.annotate(*args, **kwargs) 626 | def caption(self, text, s=24, pos='t', c='w', bg='k', spacing=None, padding=None): 627 | pos = pos[0].lower() 628 | t = I.text(text, s=s, c=c, bg=bg, spacing=spacing, padding=padding) 629 | if pos=='t': 630 | return self.top(t, bg=bg) 631 | elif pos=='b': 632 | return self.bottom(t, bg=bg) 633 | elif pos=='l': 634 | return self.left(t, bg=bg) 635 | elif pos=='r': 636 | return self.right(t, bg=bg) 637 | assert 0, 'data not understood' 638 | def cap(self, *args, **kwargs): 639 | return self.caption(*args, **kwargs) 640 | @staticmethod 641 | def text( 642 | text, 643 | s=24, 644 | facing='right', # write in this direction 645 | pos='left', # align to this position 646 | c='w', 647 | bg='k', 648 | h=None, 649 | w=None, 650 | spacing=None, # between lines 651 | padding=None, # around entire thing 652 | force_size=False, 653 | ): 654 | # text image utility 655 | text = str(text) 656 | s = max(1, math.floor(s)) 657 | spacing = math.ceil(s*4/10) if spacing is None else spacing 658 | padding = math.ceil(s*4/10) if padding is None else padding 659 | facing = facing.lower() 660 | if facing in ['u', 'up', 'd', 'down']: 661 | h,w = w,h 662 | c,bg = I.c255(c), I.c255(bg) 663 | f = PIL.ImageFont.truetype(_FN_ARIAL, s) 664 | 665 | td = PIL.ImageDraw.Draw(Image.new('RGBA', (1,1), (0,0,0,0))) 666 | tw,th = td.multiline_textsize(text, font=f, spacing=spacing) 667 | if not force_size: 668 | if h and h