# -- Code Cell --
import pandas as pd
from PIL import Image
import cv2 as cv
import numpy as np
import matplotlib.pyplot as plt
img = cv.imread("./starter_kit/data/images/train_data/image_10.jpg")
img =cv.cvtColor(img, cv.COLOR_BGR2RGB)
mask = np.zeros((256,256,1), dtype=np.uint8)
for i in range(256):
    if i%2 != 0:
        cv.line(mask,(0,i),(256,i),(255,255,255),1)
dst = cv.inpaint(img,mask,3,cv.INPAINT_TELEA)
plt.imshow(dst,cmap='gray')

# -- Code Cell --
img = cv.imread("./starter_kit/data/images/train_data/image_10.jpg")
img =cv.cvtColor(img, cv.COLOR_BGR2RGB)
plt.imshow(dst,cmap='gray')

# -- Code Cell --
img = cv.imread("./starter_kit/data/images/train_data/image_66.jpg")
img_gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
img_color = cv.cvtColor(img, cv.COLOR_BGR2RGB)
ret3,th3 = cv.threshold(img_gray,10,255,cv.THRESH_BINARY_INV)
mask = th3
dst = cv.inpaint(img_color,mask,3,cv.INPAINT_TELEA)

b, g, r = cv.split(dst)
def shift_channel(channel, shift):
        h, w = channel.shape
        M = np.float32([[1, 0, shift], [0, 1, 0]])
        return cv.warpAffine(channel, M, (w, h))
b_fixed = shift_channel(b, -5) 
r_fixed = shift_channel(r, 5)   
fixed = cv.merge([b_fixed, g, r_fixed])
fixed = cv.medianBlur(fixed,3)

mask = np.zeros((256,256,1), dtype=np.uint8)
for i in range(256):
    if i%2 != 0:
        cv.line(mask,(0,i),(256,i),(255,255,255),1)
fixed = cv.inpaint(fixed,mask,3,cv.INPAINT_TELEA)

lower_yellow = np.array([0, 150, 150])    
upper_yellow = np.array([100, 255, 255]) 
mask_yellow = cv.inRange(fixed, lower_yellow, upper_yellow)

fixed = cv.inpaint(fixed, mask_yellow, 3, cv.INPAINT_TELEA)

lower_cyan = np.array([150, 150, 0])      
upper_cyan = np.array([255, 255, 100]) 
mask_cyan = cv.inRange(fixed, lower_cyan, upper_cyan)

fixed = cv.inpaint(fixed, mask_cyan, 3, cv.INPAINT_TELEA)
plt.imshow(fixed)

# -- Code Cell --
