Trace error in CLI and UI (#137)

* Trace error in CLI and UI

* Update CHANGELOG

* Update README

* Update resnet18 model path

* Update dot url

* Update checkpoints url

* Replace the function load_url

* Update executable url

CHANGELOG.md

@@ -8,6 +8,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ## [Unreleased]
 
 ## What's Changed
+* Trace error in CLI and UI by @Ghassen-Chaabouni in https://github.com/sensity-ai/dot/pull/137
 * Update Windows executable by @Ghassen-Chaabouni in https://github.com/sensity-ai/dot/pull/133
 * Update colab notebook by @Ghassen-Chaabouni in https://github.com/sensity-ai/dot/pull/128
 * Add a Docker container for dot by @Ghassen-Chaabouni in https://github.com/sensity-ai/dot/pull/95

README.md

@@ -50,7 +50,7 @@ Supported methods:
 Download and run the dot executable for your OS:
 
 - Windows (Tested on Windows 10 and 11):
-  - Download `dot.zip` from [here](https://drive.google.com/file/d/1dSwxNFYafV4UmRoof59ohVWEwbTj7LfW/view), unzip it and then run `dot.exe`
+  - Download `dot.zip` from [here](https://drive.google.com/file/d/10LXgtE721YPXdHfzcDUp6ba-9klueocR/view), unzip it and then run `dot.exe`
 - Ubuntu:
   - ToDo
 - Mac (Tested on Apple M2 Sonoma 14.0):
@@ -147,7 +147,7 @@ pip install -e .
 
 ##### Download Models
 
-- Download dot model checkpoints from [here](https://drive.google.com/file/d/1Qaf9hE62XSvgmxR43dfiwEPWWS_dXSCE/view?usp=sharing)
+- Download dot model checkpoints from [here](https://drive.google.com/file/d/1nL3WkntTrVBZVQvOF2i7clY8eFRqSr8L/view)
 - Unzip the downloaded file in the root of this project
 
 #### CLI Usage

src/dot/__main__.py

@@ -4,6 +4,7 @@
 licensed under the BSD 3-Clause "New" or "Revised" License.
 """
 
+import traceback
 from typing import Union
 
 import click
@@ -51,32 +52,35 @@ def run(
         use_image (bool, optional): Pass flag to use image-swap pipeline. Defaults to False.
         limit (int, optional): The number of frames to process. Defaults to None.
     """
-    # initialize dot
-    _dot = DOT(use_video=use_video, use_image=use_image, save_folder=save_folder)
+    try:
+        # initialize dot
+        _dot = DOT(use_video=use_video, use_image=use_image, save_folder=save_folder)
 
-    # build dot
-    option = _dot.build_option(
-        swap_type=swap_type,
-        use_gpu=use_gpu,
-        gpen_type=gpen_type,
-        gpen_path=gpen_path,
-        crop_size=crop_size,
-    )
+        # build dot
+        option = _dot.build_option(
+            swap_type=swap_type,
+            use_gpu=use_gpu,
+            gpen_type=gpen_type,
+            gpen_path=gpen_path,
+            crop_size=crop_size,
+        )
 
-    # run dot
-    _dot.generate(
-        option=option,
-        source=source,
-        target=target,
-        show_fps=show_fps,
-        model_path=model_path,
-        limit=limit,
-        parsing_model_path=parsing_model_path,
-        arcface_model_path=arcface_model_path,
-        checkpoints_dir=checkpoints_dir,
-        opt_crop_size=crop_size,
-        head_pose=head_pose,
-    )
+        # run dot
+        _dot.generate(
+            option=option,
+            source=source,
+            target=target,
+            show_fps=show_fps,
+            model_path=model_path,
+            limit=limit,
+            parsing_model_path=parsing_model_path,
+            arcface_model_path=arcface_model_path,
+            checkpoints_dir=checkpoints_dir,
+            opt_crop_size=crop_size,
+            head_pose=head_pose,
+        )
+    except:  # noqa
+        print(traceback.format_exc())
 
 
 @click.command()

src/dot/simswap/parsing_model/resnet.py

@@ -4,9 +4,8 @@
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
-import torch.utils.model_zoo as modelzoo
 
-resnet18_url = "https://download.pytorch.org/models/resnet18-5c106cde.pth"
+resnet18_url = "saved_models/simswap/resnet18-5c106cde.pth"
 
 
 def conv3x3(in_planes, out_planes, stride=1):
@@ -77,7 +76,7 @@ def forward(self, x):
         return feat8, feat16, feat32
 
     def init_weight(self):
-        state_dict = modelzoo.load_url(resnet18_url)
+        state_dict = torch.load(resnet18_url, map_location=None, weights_only=False)
         self_state_dict = self.state_dict()
         for k, v in state_dict.items():
             if "fc" in k:

src/dot/ui/ui.py

@@ -7,6 +7,7 @@
 import os
 import sys
 import tkinter
+import traceback
 from pathlib import Path
 
 import click
@@ -115,10 +116,10 @@ def __init__(self, *args, **kwargs):
         self.textbox.insert(
             "0.0",
             """
-            dot (aka Deepfake Offensive Toolkit) makes real-time, controllable deepfakes ready for virtual cameras injection. \n
-            dot is created for performing penetration testing against e.g. identity verification and video conferencing systems, \n
-            for the use by security analysts, Red Team members, and biometrics researchers. \n
-            dot is developed for research and demonstration purposes. \n
+            DOT (aka Deepfake Offensive Toolkit) makes real-time, controllable deepfakes ready for virtual \n
+            cameras injection. DOT is created for performing penetration testing against e.g. identity \n
+            verification and video conferencing systems, for the use by security analysts, \n
+            Red Team members, and biometrics researchers. DOT is developed for research and demonstration purposes. \n
             As an end user, you have the responsibility to obey all applicable laws when using this program. \n
             Authors and contributing developers assume no liability and are not responsible for any misuse \n
             or damage caused by the use of this program.
@@ -821,6 +822,7 @@ def start_button_event(self, error_label):
             )
         except Exception as e:
             print(e)
+            print(traceback.format_exc())
             error_label.configure(text=e)