Public/automatic
1
0
Fork 0
mirror of https://github.com/vladmandic/automatic synced 2026-09-20 01:31:13 +02:00
Code Issues Packages Projects Releases Wiki Activity

fix yolo size checks

Browse Source 
Signed-off-by: Vladimir Mandic <mandic00@live.com>
This commit is contained in:
Vladimir Mandic
2024-11-02 07:02:03 -04:00
parent 16857dc856
commit c709110c59
1 changed files with 5 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
modules/postprocess/yolo.py
+5 -6
View File
@@ -25,7 +25,6 @@ class YoloResult:
self.box = box
self.mask = mask
self.item = item
self.size = size
self.width = width
self.height = height
self.args = args
@@ -127,17 +126,17 @@ class YoloRestorer(Detailer):
box = box.tolist()
mask_image = None
w, h = box[2] - box[0], box[3] - box[1]
size = w * h / (image.width * image.height)
min_size = (shared.opts.detailer_min_size if shared.opts.detailer_min_size > 0 and shared.opts.detailer_min_size < 1 else 0) * min(w, h)
max_size = (shared.opts.detailer_max_size if shared.opts.detailer_max_size > 0 and shared.opts.detailer_min_size < 1 else 1) * max(w, h)
if (min(w, h) >= int(min_size)) and (max(w, h) <= int(max_size)):
x_size, y_size = w/image.width, h/image.height
min_size = shared.opts.detailer_min_size if shared.opts.detailer_min_size > 0 and shared.opts.detailer_min_size < 1 else 0
max_size = shared.opts.detailer_max_size if shared.opts.detailer_max_size > 0 and shared.opts.detailer_max_size < 1 else 1
if x_size >= min_size and y_size >=min_size and x_size <= max_size and y_size <= max_size:
if mask:
mask_image = image.copy()
mask_image = Image.new('L', image.size, 0)
draw = ImageDraw.Draw(mask_image)
draw.rectangle(box, fill="white", outline=None, width=0)
cropped = image.crop(box)
result.append(YoloResult(cls=cls, label=label, score=round(score, 2), box=box, mask=mask_image, item=cropped, size=size, width=w, height=h, args=args))
result.append(YoloResult(cls=cls, label=label, score=round(score, 2), box=box, mask=mask_image, item=cropped, width=w, height=h, args=args))
if len(result) >= shared.opts.detailer_max:
break
return result