From 3998005334e05114cf7316c25e10b8a8d6f85e65 Mon Sep 17 00:00:00 2001
From: rasmusthog <code@rasmusthog.me>
Date: Sun, 9 Oct 2022 18:39:09 +0200
Subject: [PATCH] Updates to EDS-scripts

---
 nafuma/eds/io.py   | 34 +++++++++++++++++++++++++++++++---
 nafuma/eds/plot.py | 19 ++++++++++++++++---
 2 files changed, 47 insertions(+), 6 deletions(-)

diff --git a/nafuma/eds/io.py b/nafuma/eds/io.py
index 8f269d6..8d1b34f 100644
--- a/nafuma/eds/io.py
+++ b/nafuma/eds/io.py
@@ -2,13 +2,19 @@ from PIL import Image
 import numpy as np
 import cv2
 
-def read_image(path, weight=None, colour=None, resize=None):
+def read_image(path, weight=None, colour=None, crop=None, resize=None, brightness=None):
 
     img = np.array(Image.open(path))
 
     if colour is not None:
         img = change_colour(img, colour)
 
+    if brightness is not None:
+        img = increase_brightness(img, increase=brightness)
+
+    if crop is not None:
+        img = crop_image(img, crop)
+
     if resize is not None:
         img = resize_image(img, resize)
 
@@ -20,7 +26,6 @@ def read_image(path, weight=None, colour=None, resize=None):
 
 def scale_image(image, factor):
 
-
     for i in range(0,image.shape[0]):
         for j in range(0, image.shape[1]):
             image[i][j][0] = image[i][j][0]*factor
@@ -30,7 +35,7 @@ def scale_image(image, factor):
     return image
 
 
-def resize_image(image, factor):
+def crop_image(image, factor):
 
     y, x = image.shape[0:2]
 
@@ -43,6 +48,29 @@ def resize_image(image, factor):
     return res
 
 
+def resize_image(image, factor):
+
+    y, x = image.shape[0:2]
+
+    new_y, new_x = int(y*factor), int(x*factor)
+
+    res = cv2.resize(image, dsize=(new_x, new_y), interpolation=cv2.INTER_CUBIC)
+
+    return res
+
+
+def increase_brightness(image, brightness):
+
+    for i in range(0,image.shape[0]):
+        for j in range(0, image.shape[1]):
+            image[i][j][0] = image[i][j][0]+brightness
+            image[i][j][1] = image[i][j][1]+brightness
+            image[i][j][2] = image[i][j][2]+brightness
+
+
+    return image
+
+
 def add_images(image1, image2):
 
     assert image1.shape == image2.shape
diff --git a/nafuma/eds/plot.py b/nafuma/eds/plot.py
index 9e634e4..7726f75 100644
--- a/nafuma/eds/plot.py
+++ b/nafuma/eds/plot.py
@@ -15,7 +15,12 @@ def show_image(data, options={}):
         'hide_x_ticks': True,
         'hide_y_ticks': True,
         'colours': None,
-        'show_image': True
+        'brightness': None,
+        'show_image': True,
+        'resize': None,
+        'crop': None,
+        'ax': None,
+        'fig': None,
     }
     
     options = aux.update_options(options=options, required_options=default_options.keys(), default_options=default_options)
@@ -37,7 +42,7 @@ def show_image(data, options={}):
             options['colours'] = [None for _ in range(len(data['path']))]
     
         for i, (path, weight, colour) in enumerate(zip(data['path'], data['weights'], options['colours'])):
-            data['image'][i] = io.read_image(path=path, weight=weight, colour=colour, resize=options['resize'])
+            data['image'][i] = io.read_image(path=path, weight=weight, colour=colour, resize=options['resize'], crop=options['crop'])
 
     
     images = []
@@ -45,11 +50,19 @@ def show_image(data, options={}):
         images.append(image)
 #
     final_image = np.mean(images, axis=0) / 255
+    if options['brightness']:
+        final_image = io.increase_brightness(final_image, brightness=options['brightness'])
+
     if len(data['path']) > 1:
         data['image'].append(final_image)
 
+
     if options['show_image']:
-        fig, ax = btp.prepare_plot(options)
+        if not options['fig'] and not options['ax']:
+            fig, ax = btp.prepare_plot(options)
+        else:
+            fig, ax = options['fig'], options['ax']
+
         ax.imshow(final_image)
         btp.adjust_plot(fig=fig, ax=ax, options=options)