Allow filtering of marking confidence.

Also rename the `x` and `y` filter to `width` and `height` as it is better describing their behaviour

Author: StableLlama

README.md

@@ -144,6 +144,10 @@ apply:
       tag `catastrophe`.
 - `marking`: Images that contain at least one marking with this label. It
   doesn't matter whether it is a _hint_, _include_ or _exclude_ marking.
+  - `marking:cat:>0.789` when the label is followed by a colon with a
+    relational operator and a number then only those markings are matched 
+    where the confidence fits to the specified number, in this example case
+    a `cat` marking must have a confidence number higher than 0.789.
 - `name`: Images that contain the filter term in the file name
     - `name:cat` will match images such as `cat-1.jpg` or `large_cat.png`.
 - `path`: Images that contain the filter term in the full file path
@@ -174,9 +178,9 @@ comparison.
       caption.
     - `tokens:<=50` will match images that have 50 or fewer tokens in the
       caption.
-- `x` and `y`: will match images with the specified x or y dimension.
-    - `x:>512` will match images where the width is greater than 512 pixels.
-    - `y:=1024` will match images where the height is exactly 1024 pixels.
+- `width` and `height`: will match images with the specified width or height.
+    - `width:>512` will match images where the width is greater than 512 pixels.
+    - `height:=1024` will match images where the height is exactly 1024 pixels.
 
 ### Spaces and quotes
 

taggui/auto_marking/marking_thread.py

@@ -58,6 +58,6 @@ def generate_output(self, image_index, image: Image, image_prompt, model_inputs)
                     markings.append({'box': box,
                                      'label': marking[0],
                                      'type': marking[1],
-                                     'confidence': confidence})
+                                     'confidence': round(confidence, 3)})
         self.marking_generated.emit(image_index, markings)
         return f'Found {len(markings)} marking(s).'

taggui/models/proxy_image_list_model.py

@@ -1,4 +1,5 @@
 import operator
+import re
 from fnmatch import fnmatchcase
 
 from PySide6.QtCore import QModelIndex, QSortFilterProxyModel, Qt
@@ -8,6 +9,17 @@
 from utils.image import Image
 import utils.target_dimension as target_dimension
 
+comparison_operators = {
+    '=': operator.eq,
+    '==': operator.eq,
+    '!=': operator.ne,
+    '<': operator.lt,
+    '>': operator.gt,
+    '<=': operator.le,
+    '>=': operator.ge
+}
+
+
 class ProxyImageListModel(QSortFilterProxyModel):
     def __init__(self, image_list_model: ImageListModel,
                  tokenizer: PreTrainedTokenizerBase, tag_separator: str):
@@ -34,7 +46,23 @@ def does_image_match_filter(self, image: Image,
                 caption = self.tag_separator.join(image.tags)
                 return fnmatchcase(caption, f'*{filter_[1]}*')
             if filter_[0] == 'marking':
-                return any(fnmatchcase(markings.label, filter_[1]) for markings in image.markings)
+                last_colon_index = filter_[1].rfind(':')
+                if last_colon_index < 0:
+                    return any(fnmatchcase(marking.label, filter_[1])
+                               for marking in image.markings)
+                else:
+                    label = filter_[1][:last_colon_index]
+                    confidence = filter_[1][last_colon_index + 1:]
+                    pattern =r'^(<=|>=|==|<|>|=)\s*(0?[.,][0-9]+)'
+                    match = re.match(pattern, confidence)
+                    if not match or len(match.group(2)) == 0:
+                        return False
+                    comparison_operator = comparison_operators[match.group(1)]
+                    confidence_target = float(match.group(2).replace(',', '.'))
+                    return any((fnmatchcase(marking.label, label) and
+                               comparison_operator(marking.confidence,
+                                                   confidence_target))
+                               for marking in image.markings)
             if filter_[0] == 'name':
                 return fnmatchcase(image.path.name, f'*{filter_[1]}*')
             if filter_[0] == 'path':
@@ -48,7 +76,7 @@ def does_image_match_filter(self, image: Image,
             if filter_[0] == 'target':
                 # accept any dimension separator of [x:]
                 dimension = (filter_[1]).replace(':', 'x').split('x')
-                if image.target_dimension == None:
+                if image.target_dimension is None:
                     image.target_dimension = target_dimension.get(image.dimensions)
                 return (len(dimension) == 2
                         and dimension[0] == str(image.target_dimension.width())
@@ -59,15 +87,6 @@ def does_image_match_filter(self, image: Image,
         if filter_[1] == 'OR':
             return (self.does_image_match_filter(image, filter_[0])
                     or self.does_image_match_filter(image, filter_[2:]))
-        comparison_operators = {
-            '=': operator.eq,
-            '==': operator.eq,
-            '!=': operator.ne,
-            '<': operator.lt,
-            '>': operator.gt,
-            '<=': operator.le,
-            '>=': operator.ge
-        }
         comparison_operator = comparison_operators[filter_[1]]
         number_to_compare = None
         if filter_[0] == 'tags':
@@ -79,9 +98,9 @@ def does_image_match_filter(self, image: Image,
             caption = self.tag_separator.join(image.tags)
             # Subtract 2 for the `<|startoftext|>` and `<|endoftext|>` tokens.
             number_to_compare = len(self.tokenizer(caption).input_ids) - 2
-        elif filter_[0] == 'x':
+        elif filter_[0] == 'width':
             number_to_compare = image.dimensions[0]
-        elif filter_[0] == 'y':
+        elif filter_[0] == 'height':
             number_to_compare = image.dimensions[1]
         return comparison_operator(number_to_compare, int(filter_[2]))
 

taggui/widgets/image_list.py

@@ -56,7 +56,7 @@ def __init__(self):
                                            + optionally_quoted_string)
                                      for key in string_filter_keys]
         comparison_operator = one_of('= == != < > <= >=')
-        number_filter_keys = ['tags', 'chars', 'tokens', 'x', 'y']
+        number_filter_keys = ['tags', 'chars', 'tokens', 'width', 'height']
         number_filter_expressions = [Group(CaselessLiteral(key) + Suppress(':')
                                            + comparison_operator + Word(nums))
                                      for key in number_filter_keys]