Skip to content

Commit

Permalink
Added mask code and input file
Browse files Browse the repository at this point in the history
  • Loading branch information
@lcipolina
lcipolina committed Apr 19, 2022
1 parent fdb2b11 commit 264b242
Showing 1 changed file with 192 additions and 16 deletions.
208 changes: 192 additions & 16 deletions ML/Colab-co-mod-gan.ipynb
View file
Original file line number Diff line number Diff line change
Expand Up @@ -53,6 +53,8 @@
"id": "cRlqyRpO2fai"
},
"source": [
"import cv2\n",
"import numpy as np\n",
"!nvidia-smi"
],
"execution_count": null,
Expand Down Expand Up @@ -100,6 +102,187 @@
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"source": [
"### Upload image and save name\n",
" This is a little bit slower than just drag and drop - the only advantage is the name saving"
],
"metadata": {
"id": "fG99T5EjeCRx"
}
},
{
"cell_type": "code",
"source": [
"from google.colab import files\n",
"uploaded = files.upload()\n",
"# Get the name of uploaded image:\n",
"for key in uploaded.keys():\n",
" uploaded_image = key"
],
"metadata": {
"id": "WvIrdSCKeNm2"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"source": [
"uploaded_image"
],
"metadata": {
"id": "ua5fsYLreiff"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"source": [
"### Mask drawing code"
],
"metadata": {
"id": "MTO8223mdVTq"
}
},
{
"cell_type": "code",
"source": [
"from IPython.display import HTML\n",
"from google.colab.output import eval_js\n",
"from base64 import b64decode\n",
"\n",
"canvas_html = \"\"\"\n",
"<style>\n",
".button {\n",
" background-color: #4CAF50;\n",
" border: none;\n",
" color: white;\n",
" padding: 15px 32px;\n",
" text-align: center;\n",
" text-decoration: none;\n",
" display: inline-block;\n",
" font-size: 16px;\n",
" margin: 4px 2px;\n",
" cursor: pointer;\n",
"}\n",
"</style>\n",
"<canvas1 width=%d height=%d>\n",
"</canvas1>\n",
"<canvas width=%d height=%d>\n",
"</canvas>\n",
"\n",
"<button class=\"button\">Finish</button>\n",
"<script>\n",
"var canvas = document.querySelector('canvas')\n",
"var ctx = canvas.getContext('2d')\n",
"\n",
"var canvas1 = document.querySelector('canvas1')\n",
"var ctx1 = canvas.getContext('2d')\n",
"\n",
"\n",
"ctx.strokeStyle = 'red';\n",
"\n",
"var img = new Image();\n",
"img.src = \"data:image/%s;charset=utf-8;base64,%s\";\n",
"console.log(img)\n",
"img.onload = function() {\n",
" ctx1.drawImage(img, 0, 0);\n",
"};\n",
"img.crossOrigin = 'Anonymous';\n",
"\n",
"ctx.clearRect(0, 0, canvas.width, canvas.height);\n",
"\n",
"ctx.lineWidth = %d\n",
"var button = document.querySelector('button')\n",
"var mouse = {x: 0, y: 0}\n",
"\n",
"canvas.addEventListener('mousemove', function(e) {\n",
" mouse.x = e.pageX - this.offsetLeft\n",
" mouse.y = e.pageY - this.offsetTop\n",
"})\n",
"canvas.onmousedown = ()=>{\n",
" ctx.beginPath()\n",
" ctx.moveTo(mouse.x, mouse.y)\n",
" canvas.addEventListener('mousemove', onPaint)\n",
"}\n",
"canvas.onmouseup = ()=>{\n",
" canvas.removeEventListener('mousemove', onPaint)\n",
"}\n",
"var onPaint = ()=>{\n",
" ctx.lineTo(mouse.x, mouse.y)\n",
" ctx.stroke()\n",
"}\n",
"\n",
"var data = new Promise(resolve=>{\n",
" button.onclick = ()=>{\n",
" resolve(canvas.toDataURL('image/png'))\n",
" }\n",
"})\n",
"</script>\n",
"\"\"\"\n",
"\n",
"def draw(imgm, filename='mask.png', w=400, h=200, line_width=1):\n",
" display(HTML(canvas_html % (w, h, w,h, filename.split('.')[-1], imgm, line_width)))\n",
" data = eval_js(\"data\")\n",
" binary = b64decode(data.split(',')[1])\n",
" with open(filename, 'wb') as f:\n",
" f.write(binary)\n"
],
"metadata": {
"id": "UTD9kg1WdiOf"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"source": [
"width = 6"
],
"metadata": {
"id": "je8UaAZHdpdK"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"source": [
"import base64, os\n",
"from base64 import b64decode\n",
"#import cv2\n",
"#convert to np array\n",
"img = cv2.imread(uploaded_image)\n",
"image64 = base64.b64encode(open(uploaded_image, 'rb').read()) #encoding images as text, this is a trick to manipulate the image better.\n",
"image64 = image64.decode('utf-8')\n",
"\n",
"pencil_width = (width/100)*img.shape[1] #sets width of brush\n",
"draw(image64, w=img.shape[1], h=img.shape[0], line_width=pencil_width)\n",
"\n",
"#Save the masked image and convert the mask\n",
"import matplotlib.pyplot as plt #it uses this garbage to read images and convert to arrays, instead of CV2, wrong choice of library\n",
"with_mask = np.array(plt.imread(\"mask.png\")[:,:,:3])\n",
"mask = (with_mask[:,:,0]==1)*(with_mask[:,:,1]==0)*(with_mask[:,:,2]==0)\n",
"plt.imsave(\"mask.png\",mask, cmap='gray') #save mask"
],
"metadata": {
"id": "e9xKy4c5d9ed"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "markdown",
"source": [
"### Resize and convert images if needed"
],
"metadata": {
"id": "yxtpYdQMdX2R"
}
},
{
"cell_type": "code",
"metadata": {
Expand All @@ -110,8 +293,11 @@
"#@title (optional) read / write / resize file with opencv / invert mask\n",
"#@markdown You need masks where ``white = original`` and ``black = inpaint``. Both images are rgb and have the needed dimensions.\n",
"import cv2\n",
"path_image = '/content/input.png'\n",
"path_mask = '/content/mask.png'\n",
"\n",
"input_name = uploaded_image\n",
"\n",
"path_image = '/content/co-mod-gan/' + input_name\n",
"path_mask = '/content/co-mod-gan/' + 'mask.png'\n",
"image = cv2.imread(path_image)\n",
"image = cv2.resize(image, (512,512), cv2.INTER_NEAREST)\n",
"cv2.imwrite(path_image, image)\n",
Expand All @@ -127,10 +313,12 @@
{
"cell_type": "code",
"source": [
"# CHECK mask color AND APPLY THIS ONLY IF NEEDED\n",
"\n",
"# ***** mask color inversion ***************************************************************************\n",
"# The part of the image that has to be predicted by the model should be in black and the rest in white\n",
"# The code above inverts the colors of the mask!\n",
"src = '/content/'\n",
"src = '/content/co-mod-gan'\n",
"#convert mask\n",
"img = cv2.imread(src + 'input.png') # bitwise function needs img as array # use previously saved name \n",
"cv2.imwrite(src + 'mask.png', cv2.bitwise_not(img)) #operations are not done -in place - we need to go back from array to image"
Expand All @@ -154,18 +342,6 @@
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"source": [
"import cv2\n",
"import numpy as np"
],
"metadata": {
"id": "6jzYeR3m5oLh"
},
"execution_count": null,
"outputs": []
},
{
"cell_type": "code",
"source": [
Expand Down Expand Up @@ -201,7 +377,7 @@
"# change name of input and mask files\n",
"\n",
"%cd /content/co-mod-gan\n",
"!python run_generator.py -c /content/co-mod-gan/co-mod-gan-places2-050000.pkl -i /content/input.png -m /content/mask.png -o /content/output.png\n",
"!python run_generator.py -c /content/co-mod-gan/co-mod-gan-places2-050000.pkl -i /content/co-mod-gan/Ecce_Homo.png -m /content/co-mod-gan/mask.png -o /content/output.png\n",
"\n",
"print('DONE')\n",
"\n",
Expand Down

0 comments on commit 264b242

Please sign in to comment.