ostris/ai-toolkit
1
0
Fork 0
mirror of https://github.com/ostris/ai-toolkit.git synced 2026-04-30 11:11:37 +00:00
Code Issues Packages Projects Releases Wiki Activity

Fixed Dora implementation. Still highly experimental

Browse Source
This commit is contained in:
Jaret Burkett
2024-02-24 10:26:01 -07:00
parent 1bd94f0f01
commit f965a1299f
5 changed files with 128 additions and 33 deletions
Download Patch File Download Diff File Expand all files Collapse all files

19
toolkit/dataloader_mixins.py
View File

@@ -629,8 +629,23 @@ class ClipImageFileItemDTOMixin:
# Convert RGB to BGR
open_cv_image = open_cv_image[:, :, ::-1].copy()
# apply augmentations
augmented = self.clip_image_aug_transform(image=open_cv_image)["image"]
if self.clip_vision_is_quad:
# image is in a 2x2 gris. split, run augs, and recombine
# split
img1, img2 = np.hsplit(open_cv_image, 2)
img1_1, img1_2 = np.vsplit(img1, 2)
img2_1, img2_2 = np.vsplit(img2, 2)
# apply augmentations
img1_1 = self.clip_image_aug_transform(image=img1_1)["image"]
img1_2 = self.clip_image_aug_transform(image=img1_2)["image"]
img2_1 = self.clip_image_aug_transform(image=img2_1)["image"]
img2_2 = self.clip_image_aug_transform(image=img2_2)["image"]
# recombine
augmented = np.vstack((np.hstack((img1_1, img1_2)), np.hstack((img2_1, img2_2))))
else:
# apply augmentations
augmented = self.clip_image_aug_transform(image=open_cv_image)["image"]
# convert back to RGB tensor
augmented = cv2.cvtColor(augmented, cv2.COLOR_BGR2RGB)