support anime segmentation model

README.md

@@ -1,19 +1,20 @@
 # Cutie Roto
 
-Cutie Roto is a fork of [Cutie](https://github.com/hkchengrex/Cutie), and is designed to be a user friendly tool for AI assisted rotoscoping of video clips. It serves as a free alternative to commercial solutions such as Adobe’s Roto Brush or DaVinci Resolve Magic Mask. This tool is still in early development.
+Cutie Roto is a fork of [Cutie](https://github.com/hkchengrex/Cutie), and is designed to be a user friendly tool for AI assisted rotoscoping of video clips. It serves as a free alternative to commercial solutions such as Adobe’s Roto Brush or DaVinci Resolve Magic Mask. This tool is still in early development. It is generally less accurate than manual rotoscoping, but can usually give a pretty good result with little effort.
 
 ### Changes from Cutie:
 - Downloadable package for windows users
 - Launcher to select a video file to work on
 - New export video dialog
 - Zoomed minimap inspired from [XMem](https://github.com/hkchengrex/XMem)
 - Ability to import black and white mattes as a mask
+- Additional click segmentation model trained on anime
 - Simplified interface
 
 ### Roadmap (planned features):
-- Re-Implement brush tool from [XMem](https://github.com/hkchengrex/XMem)
-- Investigate implementing a matting model for edge refinement
-- Re-train all models for anime content (original models will also be kept)
+- Re-implement brush tool from [XMem](https://github.com/hkchengrex/XMem)
+- Replace the minimap with the ability to zoom directly on the main canvas
+- Continue to train better models for anime
 
 ### Installation (Windows):
 - Download latest version from [releases](https://github.com/Zarxrax/Cutie-Roto/releases)
@@ -40,14 +41,16 @@ python cutie_roto.py
 
 
 ### How to use:
+You should trim your video down to the specific scene that you are trying to mask prior to loading it into the program. Some sample clips are included in the examples folder.
+
 When you launch the program, you will be prompted to select a video file to load. A workspace folder will also be created for that video. Because the workspace is based on the filename, please avoid working on multiple files that have the same name. At any time, you can delete any workspace folders as long as you don't mind losing any work contained in them.
 
 When the main application launches, you have a timeline where you can view all of the frames of the video. You can create a mask on a frame by left clicking to add areas to the mask (highlighted in red), or right click to remove areas from the mask. If you are satisfied with how a mask looks, press the 'commit to permanent memory' button to store this frame in the application's memory, so it can be used to help mask other frames.
 Use the propagate forward and backward buttons to propagate the mask onto additional frames in the video.
 
 You can import masks created in external software by clicking the 'import mask' button. The mask should be a black and white image file, with white indicating the object to be masked. The mask will be loaded onto the frame that is currently displayed, and it will automatically be committed to memory.
 
-If you want to erase all masked frames and start over, just go to the first frame and press 'Reset frame', then press 'Reset all memory'. Finally, click 'propagate forward' to remove the masks from all remaining frames. 
+If you want to erase all masked frames and start over, just go to the first frame and press 'Reset frame', then press 'Reset all memory'. Finally, click 'full propagate' to remove the masks from all remaining frames. 
 
 Once you have finished masking your clip, press the 'Export as video' button to create a video file that can be brought back into another application. Various image sequences will also be created in the workspace folder.
 

cutie/config/gui_config.yaml

@@ -11,7 +11,7 @@ vis_endpoint: vis
 amp: True
 weights: weights/cutie-base-mega.pth
 ritm_weights: weights/coco_lvis_h18_itermask.pth
-#ritm_weights: weights/last_checkpoint.pth
+ritm_anime_weights: weights/aniclick_h18_itermask.pth
 
 # All "size" parameters represent the length of the shorter edge
 # maximum internal processing size; reducing this speeds up processing

gui/TIPS.md

@@ -5,6 +5,8 @@ Basic workflow: annotate objects on one or more frames and use propagation to co
 Commit good frames to permanent memory for best results. Reset memory if needed.
 
 - Use left-click for foreground annotation and right-click for background annotation.
+- If you do not start getting a good mask result after several clicks, it may be better to reset the frame and try again, clicking in different areas.
+- The standard click segmentation model may work better on some anime scenes than the anime model.
 - Use import mask to import a black and white mask that was created in other software.
 - Use Propogate Forward and Backward buttons to propogate masks across the frames. Use full propagation to run it on the entire clip from the beginning.
 - Use Export as Video to save the output to a video file.

gui/gui.py

@@ -219,6 +219,16 @@ def __init__(self, controller, cfg: DictConfig) -> None:
         self.comboBox_quality.setCurrentText("Normal")
         self.comboBox_quality.currentIndexChanged.connect(controller.on_quality_change)
 
+        self.modelselect_label = QLabel(u"Click Segmentation Model")
+        self.modelselect_label.setAlignment(Qt.AlignRight)
+        self.comboBox_modelselect = QComboBox()
+        self.comboBox_modelselect.setToolTip(u"Select the model you want to use when clicking on the image. Does not affect frame propagation.\nThe anime model is heavily biased towards characters, and wants to select all characters in the image.")
+        self.comboBox_modelselect.setObjectName(u"comboBox_modelselect")
+        self.comboBox_modelselect.addItem(u"Standard")
+        self.comboBox_modelselect.addItem(u"Anime")
+        self.comboBox_modelselect.setCurrentText("Standard")
+        self.comboBox_modelselect.currentIndexChanged.connect(controller.on_modelselect_change)
+
         # import mask/layer
         self.import_mask_button = QPushButton('Import mask')
         self.import_mask_button.clicked.connect(controller.on_import_mask)
@@ -329,7 +339,7 @@ def __init__(self, controller, cfg: DictConfig) -> None:
         right_area.addLayout(minimap_ctrl)
         right_area.addWidget(self.minimap)
         right_area.addWidget(self.tips)
-        right_area.addStretch(1)
+        #right_area.addStretch(1)
 
         # Parameters
         right_area.addLayout(self.perm_mem_gauge_layout)
@@ -347,11 +357,6 @@ def __init__(self, controller, cfg: DictConfig) -> None:
         #right_area.addLayout(self.long_mem_max_layout)
         #right_area.addLayout(self.mem_every_box_layout)
         #right_area.addLayout(self.quality_box_layout)
-        quality_area = QHBoxLayout()
-        quality_area.setAlignment(Qt.AlignmentFlag.AlignBottom)
-        quality_area.addWidget(self.quality_label)
-        quality_area.addWidget(self.comboBox_quality)
-        right_area.addLayout(quality_area)
 
         # import mask/layer/workspace
         import_area = QHBoxLayout()
@@ -361,6 +366,20 @@ def __init__(self, controller, cfg: DictConfig) -> None:
         import_area.addWidget(self.open_workspace_button)
         right_area.addLayout(import_area)
 
+        #quality combobox
+        quality_area = QHBoxLayout()
+        quality_area.setAlignment(Qt.AlignmentFlag.AlignBottom)
+        quality_area.addWidget(self.quality_label)
+        quality_area.addWidget(self.comboBox_quality)
+        right_area.addLayout(quality_area)
+
+        #quality combobox
+        modelselect_area = QHBoxLayout()
+        modelselect_area.setAlignment(Qt.AlignmentFlag.AlignBottom)
+        modelselect_area.addWidget(self.modelselect_label)
+        modelselect_area.addWidget(self.comboBox_modelselect)
+        right_area.addLayout(modelselect_area)
+
         # console
         right_area.addWidget(self.console)
 

gui/main_controller.py

@@ -144,6 +144,14 @@ def initialize_networks(self) -> None:
 
         self.click_ctrl = ClickController(self.cfg.ritm_weights, device=self.device)
 
+    def on_modelselect_change(self):
+        if self.gui.comboBox_modelselect.currentText() == 'Standard':
+            self.click_ctrl = ClickController(self.cfg.ritm_weights, device=self.device)
+            self.gui.text('Standard segmentation model loaded.')
+        else:
+            self.click_ctrl = ClickController(self.cfg.ritm_anime_weights, device=self.device)
+            self.gui.text('Anime segmentation model loaded.')
+
     def hit_number_key(self, number: int):
         if number == self.curr_object:
             return
@@ -580,8 +588,6 @@ def on_quality_change(self):
             elif self.gui.comboBox_quality.currentText() == 'Ultra':
                 self.gui.long_mem_max.setValue(4000)
                 self.gui.quality_box.setValue(720)
-            print(self.gui.long_mem_max.value())
-            print(self.gui.quality_box.value())
             self.update_config()
 
 

scripts/download_models.py

@@ -6,6 +6,7 @@
 
 _links = [
     ('https://github.com/hkchengrex/Cutie/releases/download/v1.0/coco_lvis_h18_itermask.pth', '6fb97de7ea32f4856f2e63d146a09f31'),
+    ('https://github.com/Zarxrax/Cutie-Roto/releases/download/aniclick_v1/aniclick_h18_itermask.pth', '3c5d4a84f5613e088a252f57859e34f3'),
     ('https://github.com/hkchengrex/Cutie/releases/download/v1.0/cutie-base-mega.pth', 'a6071de6136982e396851903ab4c083a'),
 ]