Mariusz wielozadaniowy

.gitignore

@@ -141,3 +141,4 @@ fabric.properties
 ### own
 damaged_*
 ok_*
+/dataset/

dataset.py

@@ -1,11 +1,10 @@
-
 import os
 from torch.utils import data
 from torchvision import transforms, utils
 from PIL import Image
 
 class dataset(data.Dataset):
-    def __init__(self, root_path, ok_path = "ok_dataset", damaged_path = "damaged_dataset"):
+    def __init__(self, root_path, ok_path="ok_dataset", damaged_path="damaged_dataset"):
         super(dataset, self).__init__()
         self.ok_images_path = os.path.join(root_path, ok_path)
         self.damaged_images_path = os.path.join(root_path, damaged_path)
@@ -13,7 +12,6 @@ class dataset(data.Dataset):
         self.damaged_images_paths = [x for x in os.listdir(self.damaged_images_path)]
         self.transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])
 
-        
     def __getitem__(self, i):
         ok_image = Image.open(os.path.join(self.ok_images_path, self.ok_images_paths[i])).convert("RGB")
         ok_image = self.transform(ok_image)
@@ -25,10 +23,8 @@ class dataset(data.Dataset):
     def __len__(self):
         return len(self.ok_images_paths)
 
+
 def get_dataloader():
     trainset = dataset(".")
     trainloader = data.DataLoader(trainset)
     return trainloader
-
-
-    

prepare_dataset.py

@@ -21,13 +21,13 @@ if __name__ == '__main__':
     
     if not os.path.exists(damaged_path):
         os.makedirs(damaged_path)
+
     if not os.path.exists(ok_path):
-        os.makedirs(ok_path);
+        os.makedirs(ok_path)
     
     dataset_images_paths = [os.path.join(dataset_path, f) for f in os.listdir(dataset_path)]
     ok_images_paths = [os.path.join(ok_path, f) for f in os.listdir(dataset_path)]
     damaged_images_paths = [os.path.join(damaged_path, f) for f in os.listdir(dataset_path)]
 
     for dataset_image_path, ok_image_path, damaged_image_path in zip(dataset_images_paths, ok_images_paths, damaged_images_paths):
-        call(['./uss-mariusz.py', dataset_image_path, ok_image_path, damaged_image_path, '-d', str(args.decals), '-n', str(args.noise)])
-        
+        call(['python', './uss-mariusz.py', dataset_image_path, ok_image_path, damaged_image_path, '-d', str(args.decals), '-n', str(args.noise)])

uss-mariusz.py

@@ -1,9 +1,11 @@
 from random import random, choice
-from PIL import Image, ImageFilter
+from PIL import Image, ImageFilter, ImageEnhance
 from functools import reduce
 
 import argparse
 import os
+import builtins
+
 
 def color_ramp(white=(255, 255, 255), black=(0, 0, 0)):
     wr, wg, wb = white
@@ -14,42 +16,70 @@ def color_ramp(white=(255, 255, 255), black=(0, 0, 0)):
     return list(map(int, reduce(list.__add__, colors)))
 
 
+def clamp(min, value, max):
+    return builtins.max(min, builtins.min(value, max))
+
+
 sepia = color_ramp((255 - int(random()*25), 230 + int(random()*25), 179+int(random()*25)))
 
-decals = [os.path.join("decals", f) for f in os.listdir("decals") if os.path.isfile(os.path.join("decals", f))]
 
+class Mariusz:
+    def __init__(self, colors=sepia, decals=5, noise=.4, brightness=.9, blur=1.3):
+        self.colors = colors
+        self.decal_count = decals
+        self.noise = noise
+        self.brightness = brightness
+        self.blur = blur
 
-def clamp(x, value, y):
-    return max(x, min(value, y))
+        self.decals = [os.path.join("decals", f) for f in os.listdir("decals") if os.path.isfile(os.path.join("decals", f))]
 
+    def fuck(self, image):
+        print("Fucking {}...".format(image))
+        im = Image.open(image)
+        im = im.convert("L")
+        im.putpalette(self.colors)
+        im = im.convert("RGBA")
+        im = ImageEnhance.Brightness(im).enhance(self.brightness)
 
-def add_decal(image, decal):
-    decal = Image.open(decal)
-    w, h = image.size
+        ok = im.copy()
 
-    scale = .5 + random()
-    scale = scale * min(w/decal.size[0], 1)
+        Mariusz.add_noise(im, self.noise)
+        for _ in range(self.decal_count):
+            Mariusz.add_decal(im, choice(self.decals))
 
-    decal = decal.rotate(random() * 360, expand=True)
-    decal = decal.resize(map(int, (scale * decal.size[0], scale * decal.size[1])), Image.ANTIALIAS)
+        im = im.filter(ImageFilter.GaussianBlur(self.blur))
+        damaged = im.convert("RGB")
 
-    dw, dh = decal.size
-    pos = tuple(map(int, (random() * w - .5 * dw, random() * h - .5 * dh)))
+        return ok, damaged
 
-    image.paste(decal, pos, decal)
+    @staticmethod
+    def add_decal(image, decal):
+        decal = Image.open(decal)
+        w, h = image.size
 
+        scale = .5 + random()
+        scale = scale * min(w/decal.size[0], 1)
 
-def add_noise(image, level):
-    data = image.getdata()
+        decal = decal.rotate(random() * 360, expand=True)
+        decal = decal.resize(map(int, (scale * decal.size[0], scale * decal.size[1])), Image.ANTIALIAS)
 
-    def noise(t):
-        w = int(random() * 255 * level)
-        r, g, b, a = t
+        dw, dh = decal.size
+        pos = tuple(map(int, (random() * w - .5 * dw, random() * h - .5 * dh)))
 
-        return clamp(0, r - w, 255), clamp(0, g - w, 255), clamp(0, b - w, 255), a
+        image.paste(decal, pos, decal)
 
-    data = list(map(noise, data))
-    image.putdata(data)
+    @staticmethod
+    def add_noise(image, level):
+        data = image.getdata()
+
+        def noise(t):
+            w = int(random() * 255 * level)
+            r, g, b, a = t
+
+            return clamp(0, r - w, 255), clamp(0, g - w, 255), clamp(0, b - w, 255), a
+
+        data = list(map(noise, data))
+        image.putdata(data)
 
 
 if __name__ == '__main__':
@@ -60,6 +90,8 @@ if __name__ == '__main__':
 
     parser.add_argument('-d', help='Decals count', type=int, default=8, dest="decals")
     parser.add_argument('-n', help='Noise level', type=float, default=.4, dest="noise")
+    parser.add_argument('-g', help='Blur level', type=float, default=1.5, dest="blur")
+    parser.add_argument('-b', help='Noise level', type=float, default=.9, dest="brightness")
 
     args = parser.parse_args()
 
@@ -71,15 +103,15 @@ if __name__ == '__main__':
         dir, name = os.path.split(args.input)
         args.damaged = os.path.join(dir, "damaged_" + name)
 
-    im = Image.open(args.input)
-    im = im.convert("L")
-    im.putpalette(sepia)
-    im = im.convert("RGBA")
+    mariusz = Mariusz(sepia, args.decals, args.noise, args.brightness, args.blur)
+    if os.path.isfile(args.input):
+        ok, damaged = mariusz.fuck(args.input)
+        ok.convert('RGB').save(args.output)
+        damaged.convert('RGB').save(args.damaged)
+    else:
+        for f in os.listdir(args.input):
+            path = os.path.join(args.input, f)
 
-    im.convert("RGB").save(args.output)
-    add_noise(im, args.noise)
-    for _ in range(args.decals):
-        add_decal(im, choice(decals))
-
-    im = im.filter(ImageFilter.GaussianBlur(1.5))
-    im.convert("RGB").save(args.damaged)
+            ok, damaged = mariusz.fuck(os.path.join(args.input, f))
+            ok.convert('RGB').save(os.path.join(args.output, f))
+            damaged.convert('RGB').save(os.path.join(args.damaged, f))