added basic ui

app/app.py

@@ -0,0 +1,81 @@
+import os
+import src
+from src.utils.config_loader import constants,Config
+from src.utils import config_loader
+from src.utils.script_utils import validate_config
+import importlib
+from flask import Flask,request,Response
+import PIL
+import PIL.Image
+import cv2
+import numpy as np
+import base64
+import io
+from flask import render_template
+
+model_config_path = os.path.join(constants.ARTIFACT_MODEL_DIR,"config.yaml")
+config = Config(model_config_path)
+# validate config
+validate_config(config)
+
+config_loader.config = config
+
+
+# now load model
+model_dir = constants.ARTIFACT_MODEL_DIR
+model_save_path = os.path.join(model_dir,"model.weights.h5")
+
+if not os.path.exists(model_save_path):
+    raise Exception("No model found")
+
+Model = importlib.import_module(f"src.{config.task}.model.models.{config.model}").Model
+model = Model(model_save_path)
+
+
+app = Flask(__name__)
+
+
+@app.route("/",methods=["GET"])
+def home():
+    # return "home page"
+    return render_template("index.html")
+
+@app.route("/config",methods=["GET"])
+def read_config():
+    content = open(model_config_path,"r").read()
+    return Response(content,mimetype='text')
+
+@app.route("/colorize",methods=["POST"])
+def colorize():
+    
+    files = request.files
+    file = files.get('image')
+    print(file)
+    img = PIL.Image.open(file)
+    img = img.convert("L")
+    
+    img = img.resize([config.image_size,config.image_size])
+    img = np.array(img)
+    print(img.min(),img.max())
+    print(img.shape)
+
+    # model.predict()
+    L = img[:,:,None]
+    L = (L/255*100).astype("uint8")
+    
+    AB = model.predict(L[None])[0]
+    img = np.concatenate([L, AB], axis=-1)
+    colored_img = cv2.cvtColor(img, cv2.COLOR_LAB2RGB) * 255
+
+    print(colored_img.shape)
+
+    im = PIL.Image.fromarray(colored_img.astype("uint8"))
+    rawBytes = io.BytesIO()
+    im.save(rawBytes, "jpeg")
+    rawBytes.seek(0)
+    base64_img = (base64.b64encode(rawBytes.read())).decode("utf-8")
+    
+    return {"image":base64_img}
+
+
+app.run(debug=True,host="0.0.0.0",port=5000)

app/static/script.js

@@ -0,0 +1,75 @@
+var image = new MarvinImage();
+var gray_scale_image = new MarvinImage();
+
+image_file_input.addEventListener("change",(ev)=>{
+    console.log(ev);
+    console.log();
+    if(ev.target.files.length){
+        show_black_and_white_image(ev.target.files[0]);
+    }
+})
+
+
+function show_black_and_white_image(file){
+    var fr = new FileReader();
+    fr.onload = ()=>{
+        original_img.src = fr.result;
+
+        // var canvas = document.createElement("canvas");
+        // var ctx = canvas.getContext("2d");
+
+        // original_img.onload = function() {
+        //     console.log(original_img.width,original_img.height);
+        //     // console.log(original_img.width,original_img.height);
+        //     canvas.height = original_img.height;
+        //     canvas.width = original_img.width;
+            
+        //     ctx.drawImage(original_img, 0, 0,original_img.width,original_img.height);
+        //     console.log(canvas.width,canvas.height);
+            
+        //     // image grayscale logic
+        //     var imageData = ctx.getImageData(0,0,original_img.width,original_img.height);
+
+        //     for(var i=0;i<imageData.data.length;i+=4){
+        //         var avg = (
+        //             imageData.data[i]
+        //             + imageData.data[i+1]
+        //             + imageData.data[i+2]
+        //         )/3;
+
+        //         imageData.data[i] = avg;
+        //         imageData.data[i+1] = avg;
+        //         imageData.data[i+2] = avg;
+
+        //     }
+        //     ctx.putImageData(imageData,0,0,0,0,imageData.width,imageData.height);
+
+        //     original_img.src = canvas.toDataURL();
+        //     original_img.onload = null;
+        // };
+
+        colorize();
+    }
+    fr.readAsDataURL(file);
+}
+
+function colorize(){
+    fetch(original_img.src)
+    .then((response)=>{
+        return response.blob()
+    })
+    .then((blob)=>{
+        // now send it
+        var formdata = new FormData();
+        formdata.append("image",blob,"image");
+        fetch("/colorize",{
+            method:"POST",
+            body:formdata,
+        })
+        .then(res=>res.json())
+        .then((data)=>{
+            console.log(data);
+            result_img.src = "data:image/png;charset=utf-8;base64,"+data["image"];
+        })
+    })
+}

app/static/style.css

@@ -0,0 +1,103 @@
+/* || RESET CSS */
+* {
+    margin: 0;
+    padding: 0;
+    box-sizing: border-box;
+}
+
+body {
+    font-family: 'Courier New', Courier, monospace;
+    font-size: 1.3em;
+
+    min-height: 100vh;
+
+    display: flex;
+    flex-direction: column;
+
+    background-color: #333;
+
+    margin-inline: 1em;
+}
+
+h1 {
+    color: #bbb;
+}
+
+
+/* || GENERAL STYLES */
+
+header>h1 {
+    text-align: center;
+    margin: 0.5em 0;
+}
+
+main {
+    flex-grow: 1;
+}
+
+.group {
+    width:100%;
+    display: flex;
+    justify-content: center;
+    align-items: center;
+    column-gap: 1em;
+    row-gap: 2em;
+
+    flex-wrap: wrap;
+    position: relative;
+}
+
+.group .container {
+    width: 90%;
+    max-width:400px;
+    height: 300px;
+    overflow: hidden;
+}
+
+.container img {
+    width:100%;
+    height:90%;
+    border-radius: 10px;
+    border: 2px solid #fff;
+    object-fit: contain;
+
+    background-color: #222;
+    /* object-fit: cover; */
+}
+
+.container__label {
+    display: block;
+    color: #fff;
+    font-weight: 500;
+    font-size: 0.9em;
+}
+
+
+#image_file_input {
+    display: None
+}
+
+#original_img {
+    cursor: pointer;
+}
+
+
+
+footer {
+    width: 100%;
+    display: flex;
+    justify-content: center;
+    align-items: center;
+
+    text-align: center;
+    font-size: 0.9em;
+
+    margin-top: 2em;
+    padding: 0.4em 0;
+    color: #ccc;
+}
+
+footer a {
+    text-decoration: underline;
+    color: #fff;
+}

app/templates/index.html

@@ -0,0 +1,38 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8">
+    <meta name="viewport" content="width=device-width, initial-scale=1.0">
+    <title>Image Colorization</title>
+    <link rel="stylesheet" href="static/style.css"/>
+    <script src="static/script.js" defer></script>
+    <script src="https://www.marvinj.org/releases/marvinj-0.8.js"></script>
+</head>
+<body>
+    <header>
+        <h1>Color Black and White Images</h1>
+    </header>
+    <main>
+        <div class="group">
+            <div class="container">
+                <img id="original_img" onclick="image_file_input.click()"/>
+                <input id="image_file_input" type="file" accept="image/*" />
+                <span class="container__label">Original Image</span>
+            </div>
+            <div class="container">
+                <img id="result_img" />
+                <span class="container__label">Predicted Image</span>
+            </div>
+        </div>
+    </main>
+    <footer>
+        <p>
+            made with passion by 
+            <a href="https://github.com/AnujPanthri/Image-Colorization/">Anuj Panthri</a>
+            |
+            trained using 
+            <a href="/config">config.yaml</a>
+        </p>
+    </footer>
+</body>
+</html>

requirements.txt

@@ -1,9 +1,11 @@
 absl-py==2.1.0
 astunparse==1.6.3
 attrs==23.2.0
+blinker==1.8.2
 Cerberus==1.3.5
 certifi==2024.2.2
 charset-normalizer==3.3.2
+click==8.1.7
 comet-ml==3.42.1
 configobj==5.0.8
 contourpy==1.2.1
@@ -11,6 +13,7 @@ cycler==0.12.1
 dulwich==0.22.1
 everett==3.1.0
 filelock==3.14.0
+Flask==3.0.3
 flatbuffers==24.3.25
 fonttools==4.51.0
 fsspec==2024.5.0
@@ -21,6 +24,8 @@ h5py==3.11.0
 huggingface-hub==0.23.1
 idna==3.7
 imageio==2.34.1
+itsdangerous==2.2.0
+Jinja2==3.1.4
 joblib==1.4.2
 jsonschema==4.22.0
 jsonschema-specifications==2023.12.1

src/simple_regression_colorization/model/base_model_interface.py

@@ -58,9 +58,9 @@ class BaseModel(ABC):
 
     def predict_colors(self,L_batch):
         AB_batch = self.predict(L_batch)
-        colored_batch = np.concatenate([L_batch,rescale_AB(AB_batch)],axis=-1)
+        colored_batch = np.concatenate([L_batch,AB_batch],axis=-1)
         colored_batch = lab2rgb(colored_batch) * 255
-        return colored_batch
+        return colored_batch.astype("uint8")
 
     def show_results(self):
         self.prepare_data()