Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,20 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+EvalDataset/clips/bear/output_video.mp4 filter=lfs diff=lfs merge=lfs -text
+EvalDataset/clips/bear/output_video_gray.mp4 filter=lfs diff=lfs merge=lfs -text
+EvalDataset/clips/boat/output_video_gray.mp4 filter=lfs diff=lfs merge=lfs -text
+EvalDataset/clips/cows/output_video.mp4 filter=lfs diff=lfs merge=lfs -text
+EvalDataset/clips/cows/output_video_gray.mp4 filter=lfs diff=lfs merge=lfs -text
+EvalDataset/clips/dog/output_video.mp4 filter=lfs diff=lfs merge=lfs -text
+EvalDataset/clips/flamingo/output_video_gray.mp4 filter=lfs diff=lfs merge=lfs -text
+EvalDataset/ref/goat/0000.jpg filter=lfs diff=lfs merge=lfs -text
+EvalDataset/ref/hockey/0000.jpg filter=lfs diff=lfs merge=lfs -text
+EvalDataset/ref/horsejump-high/0000.jpg filter=lfs diff=lfs merge=lfs -text
+EvalDataset/ref/motorbike/0000.jpg filter=lfs diff=lfs merge=lfs -text
+EvalDataset/ref/surf/0000.jpg filter=lfs diff=lfs merge=lfs -text
+examples/bear/video.mp4 filter=lfs diff=lfs merge=lfs -text
+examples/cows/video.mp4 filter=lfs diff=lfs merge=lfs -text
+examples/flamingo/video.mp4 filter=lfs diff=lfs merge=lfs -text
+gradio_cached_examples/13/output/003c3114319372a78bf2f812ebaf0041afa280fb/output_video.mp4 filter=lfs diff=lfs merge=lfs -text
+gradio_cached_examples/13/output/7969adca8ae38cb3b38ff8e7bb54688d942c7bc8/output_video.mp4 filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,137 @@
+checkpoints/
+wandb/
+.vscode
+.DS_Store
+*ckpt*/
+# Custom
+*.pt
+data/local
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/

README.md

@@ -1,12 +1,6 @@
 ---
 title: ViTExCo
-emoji: 👀
-colorFrom: gray
-colorTo: green
+app_file: app.py
 sdk: gradio
 sdk_version: 3.40.1
-app_file: app.py
-pinned: false
 ---
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference

UI.py

@@ -0,0 +1,81 @@
+import streamlit as st
+from PIL import Image
+import torchvision.transforms as transforms
+from streamlit_image_comparison import image_comparison
+import numpy as np
+import torch
+import torchvision
+
+######################################### Utils ########################################
+video_extensions = ["mp4"]
+image_extensions = ["png", "jpg"]
+
+
+def check_type(file_name: str):
+    for image_extension in image_extensions:
+        if file_name.endswith(image_extension):
+            return "image"
+    for video_extension in video_extensions:
+        if file_name.endswith(video_extension):
+            return "video"
+    return None
+
+
+transform = transforms.Compose(
+    [transforms.Resize((256, 256)), transforms.ToTensor(), transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225))]
+)
+
+
+###################################### Load model ######################################
+@st.cache_resource
+def load_model():
+    model = torchvision.models.segmentation.deeplabv3_resnet101(pretrained=True)
+    model.eval()
+    return model
+
+
+model = load_model()
+########################################## UI ##########################################
+st.title("Colorization")
+
+uploaded_file = st.file_uploader("Upload grayscale image or video", type=image_extensions + video_extensions)
+if uploaded_file:
+    # Image
+    if check_type(file_name=uploaded_file.name) == "image":
+        image = np.array(Image.open(uploaded_file), dtype=np.float32)
+
+        input_tensor = torchvision.transforms.functional.normalize(
+            torch.tensor(image).permute(2, 0, 1),
+            mean=[0.485, 0.456, 0.406],
+            std=[0.229, 0.224, 0.225],
+        ).unsqueeze(0)
+        process_button = st.button("Process")
+        if process_button:
+            with st.spinner("Từ từ coi..."):
+                prediction = model(input_tensor)
+                segment = prediction["out"][0].permute(1, 2, 0)
+                segment = segment.detach().numpy()
+
+                st.image(segment)
+                st.image(image)
+
+                image_comparison(
+                    img1=image,
+                    img2=np.array(segment),
+                    label1="Grayscale",
+                    label2="Colorized",
+                    make_responsive=True,
+                    show_labels=True,
+                )
+    # Video
+    else:
+        # video = open(uploaded_file.name)
+        st.video("https://youtu.be/dQw4w9WgXcQ")
+
+hide_menu_style = """
+        <style>
+        #MainMenu {visibility: hidden; }
+        footer {visibility: hidden;}
+        </style>
+        """
+st.markdown(hide_menu_style, unsafe_allow_html=True)

app.py

@@ -0,0 +1,215 @@
+import numpy as np
+import shutil
+import os
+import argparse
+import torch
+import glob
+from tqdm import tqdm
+from PIL import Image
+from collections import OrderedDict
+from src.models.vit.config import load_config
+import torchvision.transforms as transforms
+import cv2
+from skimage import io
+
+from src.models.CNN.ColorVidNet import GeneralColorVidNet
+from src.models.vit.embed import GeneralEmbedModel
+from src.models.CNN.NonlocalNet import GeneralWarpNet
+from src.models.CNN.FrameColor import frame_colorization
+from src.utils import (
+    RGB2Lab,
+    ToTensor,
+    Normalize,
+    uncenter_l,
+    tensor_lab2rgb,
+    SquaredPadding,
+    UnpaddingSquare
+)
+
+import gradio as gr
+
+def load_params(ckpt_file):
+    params = torch.load(ckpt_file, map_location=device)
+    new_params = []
+    for key, value in params.items():
+        new_params.append((key, value))
+    return OrderedDict(new_params)
+
+def custom_transform(transforms, img):
+    for transform in transforms:
+        if isinstance(transform, SquaredPadding):
+            img,padding=transform(img, return_paddings=True)
+        else:
+            img = transform(img)
+    return img.to(device), padding
+
+def save_frames(predicted_rgb, video_name, frame_name):
+    if predicted_rgb is not None:
+        predicted_rgb = np.clip(predicted_rgb, 0, 255).astype(np.uint8)
+        # frame_path_parts = frame_path.split(os.sep)
+        # if os.path.exists(os.path.join(OUTPUT_RESULT_PATH, frame_path_parts[-2])):
+        #   shutil.rmtree(os.path.join(OUTPUT_RESULT_PATH, frame_path_parts[-2]))
+        # os.makedirs(os.path.join(OUTPUT_RESULT_PATH, frame_path_parts[-2]), exist_ok=True)
+        predicted_rgb = np.transpose(predicted_rgb, (1,2,0))
+        pil_img = Image.fromarray(predicted_rgb)
+        pil_img.save(os.path.join(OUTPUT_RESULT_PATH, video_name, frame_name))
+
+def extract_frames_from_video(video_path):
+  cap = cv2.VideoCapture(video_path)
+  fps = cap.get(cv2.CAP_PROP_FPS)
+
+  # remove if exists folder
+  output_frames_path = os.path.join(INPUT_VIDEO_FRAMES_PATH, os.path.basename(video_path))
+  if os.path.exists(output_frames_path):
+    shutil.rmtree(output_frames_path)
+
+  # make new folder
+  os.makedirs(output_frames_path)
+
+  currentframe = 0
+  frame_path_list = []
+  while(True):
+
+      # reading from frame
+      ret,frame = cap.read()
+
+      if ret:
+          name = os.path.join(output_frames_path, f'{currentframe:09d}.jpg')
+          frame_path_list.append(name)
+          cv2.imwrite(name, frame)
+          currentframe += 1
+      else:
+          break
+
+  cap.release()
+  cv2.destroyAllWindows()
+
+  return frame_path_list, fps
+
+def combine_frames_from_folder(frames_list_path, fps = 30):
+  frames_list = glob.glob(f'{frames_list_path}/*.jpg')
+  frames_list.sort()
+
+  sample_shape = cv2.imread(frames_list[0]).shape
+
+  output_video_path = os.path.join(frames_list_path, 'output_video.mp4')
+  out = cv2.VideoWriter(output_video_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, (sample_shape[1], sample_shape[0]))
+  for filename in frames_list:
+      img = cv2.imread(filename)
+      out.write(img)
+
+  out.release()
+  return output_video_path
+
+
+def upscale_image(I_current_rgb, I_current_ab_predict):
+    H, W = I_current_rgb.size
+    high_lab_transforms = [
+      SquaredPadding(target_size=max(H,W)),
+      RGB2Lab(),
+      ToTensor(),
+      Normalize()
+    ]
+    # current_frame_pil_rgb = Image.fromarray(np.clip(I_current_rgb.squeeze(0).permute(1,2,0).cpu().numpy() * 255, 0, 255).astype('uint8'))
+    high_lab_current, paddings = custom_transform(high_lab_transforms, I_current_rgb)
+    high_lab_current = torch.unsqueeze(high_lab_current,dim=0).to(device)
+    high_l_current = high_lab_current[:, 0:1, :, :]
+    high_ab_current = high_lab_current[:, 1:3, :, :]
+    upsampler = torch.nn.Upsample(scale_factor=max(H,W)/224,mode="bilinear")
+    high_ab_predict = upsampler(I_current_ab_predict)
+    I_predict_rgb = tensor_lab2rgb(torch.cat((uncenter_l(high_l_current), high_ab_predict), dim=1))
+    upadded = UnpaddingSquare()
+    I_predict_rgb = upadded(I_predict_rgb, paddings)
+    return I_predict_rgb
+
+def colorize_video(video_path, ref_np):
+    frames_list, fps = extract_frames_from_video(video_path)
+
+    frame_ref = Image.fromarray(ref_np).convert("RGB")
+    I_last_lab_predict = None
+    IB_lab, IB_paddings = custom_transform(transforms, frame_ref)
+    IB_lab = IB_lab.unsqueeze(0).to(device)
+    IB_l = IB_lab[:, 0:1, :, :]
+    IB_ab = IB_lab[:, 1:3, :, :]
+
+    with torch.no_grad():
+        I_reference_lab = IB_lab
+        I_reference_l = I_reference_lab[:, 0:1, :, :]
+        I_reference_ab = I_reference_lab[:, 1:3, :, :]
+        I_reference_rgb = tensor_lab2rgb(torch.cat((uncenter_l(I_reference_l), I_reference_ab), dim=1)).to(device)
+        features_B = embed_net(I_reference_rgb)
+
+    video_path_parts = frames_list[0].split(os.sep)
+
+    if os.path.exists(os.path.join(OUTPUT_RESULT_PATH, video_path_parts[-2])):
+      shutil.rmtree(os.path.join(OUTPUT_RESULT_PATH, video_path_parts[-2]))
+    os.makedirs(os.path.join(OUTPUT_RESULT_PATH, video_path_parts[-2]), exist_ok=True)
+
+    for frame_path in tqdm(frames_list):
+        curr_frame = Image.open(frame_path).convert("RGB")
+        IA_lab, IA_paddings = custom_transform(transforms, curr_frame)
+        IA_lab = IA_lab.unsqueeze(0).to(device)
+        IA_l = IA_lab[:, 0:1, :, :]
+        IA_ab = IA_lab[:, 1:3, :, :]
+
+        if I_last_lab_predict is None:
+            I_last_lab_predict = torch.zeros_like(IA_lab).to(device)
+
+        with torch.no_grad():
+            I_current_lab = IA_lab
+            I_current_ab_predict, _ = frame_colorization(
+                IA_l,
+                I_reference_lab,
+                I_last_lab_predict,
+                features_B,
+                embed_net,
+                nonlocal_net,
+                colornet,
+                luminance_noise=0,
+                temperature=1e-10,
+                joint_training=False
+            )
+            I_last_lab_predict = torch.cat((IA_l, I_current_ab_predict), dim=1)
+
+        # IA_predict_rgb = tensor_lab2rgb(torch.cat((uncenter_l(IA_l), I_current_ab_predict), dim=1))
+        IA_predict_rgb = upscale_image(curr_frame, I_current_ab_predict)
+        #IA_predict_rgb = torch.nn.functional.upsample_bilinear(IA_predict_rgb, scale_factor=2)
+        save_frames(IA_predict_rgb.squeeze(0).cpu().numpy() * 255, video_path_parts[-2], os.path.basename(frame_path))
+    return combine_frames_from_folder(os.path.join(OUTPUT_RESULT_PATH, video_path_parts[-2]), fps)
+
+if __name__ == '__main__':
+  # Init global variables
+  device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+  INPUT_VIDEO_FRAMES_PATH = 'inputs'
+  OUTPUT_RESULT_PATH = 'outputs'
+  weight_path = 'checkpoints'
+
+  embed_net=GeneralEmbedModel(pretrained_model="swin-tiny", device=device).to(device)
+  nonlocal_net = GeneralWarpNet(feature_channel=128).to(device)
+  colornet=GeneralColorVidNet(7).to(device)
+
+  embed_net.eval()
+  nonlocal_net.eval()
+  colornet.eval()
+
+  # Load weights
+  # embed_net_params = load_params(os.path.join(weight_path, "embed_net.pth"))
+  nonlocal_net_params = load_params(os.path.join(weight_path, "nonlocal_net.pth"))
+  colornet_params = load_params(os.path.join(weight_path, "colornet.pth"))
+
+  # embed_net.load_state_dict(embed_net_params, strict=True)
+  nonlocal_net.load_state_dict(nonlocal_net_params, strict=True)
+  colornet.load_state_dict(colornet_params, strict=True)
+
+  transforms = [SquaredPadding(target_size=224),
+                RGB2Lab(),
+                ToTensor(),
+                Normalize()]
+
+  examples = [[vid, ref] for vid, ref in zip(sorted(glob.glob('examples/*/*.mp4')), sorted(glob.glob('examples/*/*.jpg')))]
+  demo = gr.Interface(colorize_video,
+                      inputs=[gr.Video(), gr.Image()],
+                      outputs="playable_video",
+                      examples=examples,
+                      cache_examples=True)
+  demo.launch()

checkpoints/colornet.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5257ae325e292cd5fb2eff47095e1c4e4815455bd5fb6dc5ed2ee2b923172875
+size 131239411

checkpoints/embed_net.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:fc711755a75c43025dabe9407cbd11d164eaa9e21f26430d0c16c7493410d902
+size 110352261

checkpoints/nonlocal_net.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b94c6990f20088bc3cc3fe0b29a6d52e6e746b915c506f0cd349fc6ad6197e72
+size 73189765

cmd.txt

@@ -0,0 +1,21 @@
+python train.py --video_data_root_list datasets/images/images \
+                --flow_data_root_list datasets/flow_fp16/flow_fp16 \
+                --mask_data_root_list datasets/pgm/pgm \
+                --data_root_imagenet datasets/imgnet \
+                --annotation_file_path datasets/final_annot.csv \
+                --imagenet_pairs_file datasets/pairs.txt \
+                --gpu_ids 0 \
+                --workers 12 \
+                --batch_size 2 \
+                --real_reference_probability 0.99 \
+                --weight_contextual 1 \
+                --weight_perceptual 0.1 \
+                --weight_smoothness 5 \
+                --weight_gan 0.9 \
+                --weight_consistent 0.1 \
+                --use_wandb True \
+                --wandb_token "f05d31e6b15339b1cfc5ee1c77fe51f66fc3ea9e" \
+                --wandb_name "vit_tiny_patch16_384_nofeat" \
+                --checkpoint_step 500 \
+                --epoch_train_discriminator 3 \
+                --epoch 20

cmd_ddp.txt

@@ -0,0 +1,20 @@
+!torchrun --nnodes=1 --nproc_per_node=2 train_ddp.py --video_data_root_list $video_data_root_list \
+                --flow_data_root_list $flow_data_root_list \
+                --mask_data_root_list $mask_data_root_list \
+                --data_root_imagenet $data_root_imagenet \
+                --annotation_file_path $annotation_file_path \
+                --imagenet_pairs_file $imagenet_pairs_file \
+                --gpu_ids "0,1" \
+                --workers 2 \
+                --batch_size 2 \
+                --real_reference_probability 0.99 \
+                --weight_contextual 1 \
+                --weight_perceptual 0.1 \
+                --weight_smoothness 5 \
+                --weight_gan 0.9 \
+                --weight_consistent 0.1 \
+                --wandb_token "165e7148081f263b423722115e2ad40fa5339ecf" \
+                --wandb_name "vit_tiny_patch16_384_nofeat" \
+                --checkpoint_step 2000 \
+                --epoch_train_discriminator 2 \
+                --epoch 10

examples.zip

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bd4531bd3abdec6df90efb0d19fadd54284bdc70d5edfff19752a205159eb4db
+size 6955837

examples/bear/video.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cb4cec5064873a4616f78bdb653830683a4842b2a5cfd0665b395cff4d120d04
+size 1263445

examples/cows/video.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1ac08603d719cd7a8d71fac76c9318d3e8f1e516e9b3c2a06323a0e4e78f6410
+size 2745681

examples/flamingo/video.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5a103fd4991a00e419e5236b885fe9d220704ba0a6ac794c87aaa3f62a4f1561
+size 1239570

gradio_cached_examples/13/log.csv

@@ -0,0 +1,5 @@
+output,flag,username,timestamp
+/content/ViTExCo/gradio_cached_examples/13/output/003c3114319372a78bf2f812ebaf0041afa280fb/output_video.mp4,,,2023-08-15 09:45:37.897615
+/content/ViTExCo/gradio_cached_examples/13/output/e6d6153dedeb9fec586b3241311cc49dbc17bc85/output_video.mp4,,,2023-08-15 09:46:01.048997
+/content/ViTExCo/gradio_cached_examples/13/output/7969adca8ae38cb3b38ff8e7bb54688d942c7bc8/output_video.mp4,,,2023-08-15 09:46:34.503322
+/content/ViTExCo/gradio_cached_examples/13/output/74c76e483235b7e80665e32d7fcdcc3da2be7644/output_video.mp4,,,2023-08-15 09:46:58.088903

gradio_cached_examples/13/output/003c3114319372a78bf2f812ebaf0041afa280fb/output_video.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b5ab666998e14fb00281a90f8801753eca001a432641ae2770007a8336b4c64e
+size 1213824

gradio_cached_examples/13/output/7969adca8ae38cb3b38ff8e7bb54688d942c7bc8/output_video.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:7c367dab34e596f7f0fed34c7e2384525de2ba1824b410d0770bdbd17bc9e72a
+size 1793060