Add Database Functionality

app.py

@@ -4,31 +4,64 @@ import gradio as gr
 from helpers import *
 import shutil
 from PIL import Image
+import sqlite3
+import pathlib
+
+conn = sqlite3.connect('database.db', check_same_thread=False)
+c = conn.cursor()
+create(c)
 
 detector = load_detector()
 model = load_model()
 
-source_imgs = []
-for r, _, f in os.walk(os.getcwd() + "/images"):
-    for file in f:
-        if (
-            (".jpg" in file.lower())
-            or (".jpeg" in file.lower())
-            or (".png" in file.lower())
-        ):
-            exact_path = r + "/" + file
-            source_imgs.append(exact_path)
-
-source_faces = []
-for img in source_imgs:
-  try:
-    faces, id = extract_faces(detector, img)
-    source_faces.append(faces[id])
-    # source_faces.append(Image.open(img))
-  except Exception as e:
-    print(f"Skipping {img}, {e}")
-
-source_embeddings = get_embeddings(model, source_faces)
+# source_imgs = []
+# for r, _, f in os.walk(os.getcwd() + "/images"):
+#     for file in f:
+#         if (
+#             (".jpg" in file.lower())
+#             or (".jpeg" in file.lower())
+#             or (".png" in file.lower())
+#         ):
+#             exact_path = r + "/" + file
+#             source_imgs.append(exact_path)
+
+# source_faces = []
+# for img in source_imgs:
+#   try:
+#     faces, id = extract_faces(detector, img)
+#     source_faces.append(faces[id])
+#     # source_faces.append(Image.open(img))
+#   except Exception as e:
+#     print(f"Skipping {img}, {e}")
+
+# source_embeddings = get_embeddings(model, source_faces)
+
+def init():
+    source_imgs = c.execute("SELECT image FROM students").fetchall()
+    cwd = os.getcwd()
+    source_imgs = [os.path.join(cwd, "images", s[0]) for s in source_imgs]
+    source_faces = [Image.open(s) for s in source_imgs]
+    # source_embeddings = get_embeddings(model, source_faces)
+    source_embeddings = c.execute("SELECT embeddings FROM students").fetchall()
+    source_embeddings = [os.path.join(cwd, "embeds", s[0]) for s in source_embeddings]
+    source_embeddings = [np.load(s) for s in source_embeddings]
+    names = c.execute("SELECT name FROM students").fetchall()
+    names = [n[0] for n in names]
+    return names, source_faces, source_embeddings
+
+
+def init():
+    source_imgs = c.execute("SELECT image FROM students").fetchall()
+    cwd = os.getcwd()
+    source_imgs = [os.path.join(cwd, "images", s[0]) for s in source_imgs]
+    source_faces = [Image.open(s) for s in source_imgs]
+    # source_embeddings = get_embeddings(model, source_faces)
+    source_embeddings = c.execute("SELECT embeddings FROM students").fetchall()
+    source_embeddings = [os.path.join(cwd, "embeds", s[0]) for s in source_embeddings]
+    source_embeddings = [np.load(s) for s in source_embeddings]
+    names = c.execute("SELECT name FROM students").fetchall()
+    names = [n[0] for n in names]
+    return names, source_faces, source_embeddings
 
 def find_names(image, minSize, minConf):
   imgs, _ = extract_faces(detector, image)
@@ -37,12 +70,11 @@ def find_names(image, minSize, minConf):
     if((face.size[0] * face.size[1]) > minSize):
       ims.append(face)
   imgs = ims
-
   embeds = get_embeddings(model, imgs)
   d = np.zeros((len(source_embeddings), len(embeds)))
   for i, s in enumerate(source_embeddings):
     for j, t in enumerate(embeds):
-      d[i][j] = findCosineDistance(s, t)
+      d[i][j] = findCosineDistance(s.squeeze(), t)
   ids = np.argmin(d, axis = 0)
   names = []
   for j, i in enumerate(ids):
@@ -53,10 +85,12 @@ def find_names(image, minSize, minConf):
   recognition(imgs, ids, names, source_faces, d, source_imgs)
   return ",".join(names), "Recognition.jpg"
 
+source_imgs, source_faces, source_embeddings = init()
+
 detect = gr.Interface(
     find_names,
     [gr.Image(type="filepath", label="Class Photo"), gr.Number(label = "Minimum Size"), gr.Number(label = "Minimum Confidence")],
-    ["text" ,gr.Image(type = "filepath", label="Matching")],
+    [gr.Textbox(label = "Roll No") ,gr.Image(type = "filepath", label="Matching")],
     examples = [
         [os.path.join(os.path.dirname(__file__), "examples/group1.jpg"), 1000, 0.3],
         [os.path.join(os.path.dirname(__file__), "examples/group2.jpg"), 1000, 0.3]
@@ -65,20 +99,16 @@ detect = gr.Interface(
 
 
 def upload_files(files):
+    global i, imgs
     if not os.path.exists(os.path.join(os.getcwd(), "temp")):
         os.mkdir(os.path.join(os.getcwd(), "temp"))
     for file in files:
+        # faces, id = extract_faces(detector, os.path.join(os.getcwd(), "temp", file))
+        # imgs.append(faces)
         shutil.move(file.name, os.path.join(os.getcwd(), "temp", file.name.split('\\')[-1]))
+
     return None, "Uploaded!"
 
-with gr.Blocks() as upload:
-    gr.Markdown("# Select Images to Upload and click Upload")
-    with gr.Row():
-        input = gr.Files(file_types=[".jpg", ".jpeg", ".png"], label="Upload images")
-        # input = gr.Image(type="filepath")
-        output = gr.Textbox()
-    upload_btn = gr.Button(value="Upload")
-    upload_btn.click(upload_files, inputs=[input], outputs=[input, output])
 
 def load_image():
     global i, imgs
@@ -87,45 +117,119 @@ def load_image():
     for image in images:
         faces, id = extract_faces(detector, os.path.join(os.getcwd(), "temp", image))
         # imgs.append(Image.open(os.path.join(os.getcwd(), "temp", image)))
-        imgs.append(faces[id])
-    return imgs[0], "Loaded!"
-
-def save_img(label):
-    global i, imgs
-    imgs[i].save(os.path.join(os.getcwd(), "images", f"{label}.jpg"))
+        imgs.append(faces)
     i+=1
+    shutil.rmtree(os.path.join(os.getcwd(), "temp"))
+    return imgs[i], "Loaded!", None, None
+
+def save_img(label, email, roll, selected):
+    global i, imgs, source_imgs, source_faces, source_embeddings
+    if label:
+        imgs[i][int(selected)].save(os.path.join(os.getcwd(), "images", f"{roll}.jpg"))
+        np.save(os.path.join(os.getcwd(), "embeds", f"{roll}.npy"),get_embeddings(model, [imgs[i][int(selected)]]))
+        insert(c, label, email, roll, f"{roll}.jpg", f"{roll}.npy")
+        conn.commit()
+        source_imgs, source_faces, source_embeddings = init()
+
+    imgs[i].pop(int(selected))
+    if(len(imgs[i]) == 0):
+        i+=1
     if i < len(imgs):
-        return imgs[i], "Saved!"
+        if not label:
+            return imgs[i], None, None, None
+        return imgs[i], "Saved!", None, None
     else:
         clear()
-        return None, "Finished!"
+        return None, "Finished!", None, None
 
 def clear():
     global i, imgs
     i = 0
     imgs = None
-    shutil.rmtree(os.path.join(os.getcwd(), "temp"))
-    return None, None
+    # shutil.rmtree(os.path.join(os.getcwd(), "temp"))
+    return None, None, None, None
 
-i = 0
+i = -1
 imgs = None
 
-with gr.Blocks() as annotate:
+with gr.Blocks() as upload:
+    gr.Markdown("# Select Images to Upload and click Upload")
     with gr.Row():
-        input = gr.Textbox(label = "Enter Label")
-        output = gr.Image(type="pil", label="Image").style(height=400)
+        files = gr.Files(file_types=[".jpg", ".jpeg", ".png"], label="Upload images")
+        # input = gr.Image(type="filepath")
+        alert = gr.Textbox()
+    upload_btn = gr.Button(value="Upload")
+    upload_btn.click(upload_files, inputs=[files], outputs=[files, alert])
+    with gr.Accordion("Annotate", open=False):
+        with gr.Row():
+            with gr.Column():
+                input = gr.Textbox(label = "Name")
+                email = gr.Textbox(label = "Email")
+                roll = gr.Textbox(label = "Roll No")
+            selected = gr.Number(visible=False)
+            def get_select_index(evt: gr.SelectData):
+                return evt.index
+            output = gr.Gallery(label="Found Faces", height = 400)
+            output.select(get_select_index, None, selected)
+        with gr.Row():
+            next_btn = gr.Button(value="Next")
+            next_btn.click(load_image, inputs=[], outputs=[output, input, email, roll])
+            save_btn = gr.Button(value="Save")
+            save_btn.click(save_img, inputs=[input, email, roll, selected], outputs=[output, input, email, roll])
+            clear_btn = gr.Button(value="Clear")
+            clear_btn.click(clear, inputs=[], outputs=[output, input, email, roll])
+# with gr.Blocks() as annotate:
+
+def match(name, roll):
+    if roll:
+        results = c.execute(f"SELECT name, roll_no, image from students WHERE roll_no = {roll}").fetchall()
+    elif name:
+        results = c.execute(f'SELECT name, roll_no, image from students WHERE name = "{name}"').fetchall()
+    else:
+        results = c.execute(f"SELECT name, roll_no, image from students").fetchall()
+    names = []
+    rolls = []
+    images = []
+    cwd = os.getcwd()
+    for r in results:
+        names.append(r[0])
+        rolls.append(r[1])
+        images.append(os.path.join(cwd, "images", r[2]))
+    return names, rolls, images, images
+
+def update(name, roll, rolls, img, selected):
+    global source_imgs, source_faces, source_embeddings
+    c.execute(f'UPDATE students SET name = "{name}", roll_no = {roll} where roll_no = {rolls[int(selected)]}')
+    conn.commit()
+    source_imgs, source_faces, source_embeddings = init()
+    result = c.execute("SELECT name, roll_no from students").fetchall()
+    names = [r[0] for r in result]
+    rolls = [r[1] for r in result]
+    return None, None, img, img, names, rolls
+
+with gr.Blocks() as find:
     with gr.Row():
-        next_btn = gr.Button(value="Next")
-        next_btn.click(load_image, inputs=[], outputs=[output, input])
-        save_btn = gr.Button(value="Save")
-        save_btn.click(save_img, inputs=[input], outputs=[output, input])
-        clear_btn = gr.Button(value="Clear")
-        clear_btn.click(clear, inputs=[], outputs=[output, input])
+        with gr.Column():
+            name = gr.Textbox(label = "Name")
+            roll = gr.Textbox(label = "Roll No")
+        image = gr.Gallery(label = "Matches", height = 400)
+    with gr.Row():
+        names = gr.State()
+        rolls = gr.State()
+        images = gr.State()
+        selected = gr.Number(visible=False)
+        find_btn = gr.Button(value = "Find")
+        find_btn.click(match, inputs = [name, roll], outputs = [names, rolls, image, images])
+        update_btn = gr.Button(value = "Update")
+        update_btn.click(update, inputs = [name, roll, rolls, images, selected], outputs = [name, roll, image, images, names, rolls])
+        clear_btn = gr.ClearButton([name, roll, image])
+        def get_select_index(evt: gr.SelectData, names, rolls):
+            return names[evt.index], rolls[evt.index], evt.index
+        image.select(get_select_index, [names, rolls], [name, roll, selected])
 
 tabbed_interface = gr.TabbedInterface(
-    [detect, upload, annotate],
-    ["Attendance", "Upload", "Annotate"],
+    [detect, upload, find],
+    ["Attendance", "Register", "Find"],
 )
 if __name__ == "__main__":
-    tabbed_interface.launch()
-
+    tabbed_interface.launch()

helpers.py

@@ -3,6 +3,8 @@ from PIL import Image
 import numpy as np
 from swin import load_model, get_embeddings
 import matplotlib.pyplot as plt
+import sqlite3
+import pathlib
 
 def load_detector():
   # load model
@@ -52,4 +54,16 @@ def recognition(imgs, ids, names, source_faces, d, source_imgs):
                 axes[i, j].set_title(f"Roll No.:{source_imgs[ids[img_count]].split('/')[-1].split('.')[0]}")
                 axes[i, j].imshow(source_faces[ids[img_count]])
               axes[i, j].set_axis_off()
-  plt.savefig("Recognition.jpg")
+  plt.savefig("Recognition.jpg")
+
+def create(c):
+  pathlib.Path("images").mkdir(parents=True, exist_ok=True)
+  pathlib.Path("embeds").mkdir(parents=True, exist_ok=True)
+  c.execute('''CREATE TABLE IF NOT EXISTS students
+              (name text, email text, roll_no text, image text, embeddings text)''')
+
+  c.execute('''CREATE TABLE IF NOT EXISTS attendance
+                  (date text, roll_no text, present text)''')
+  
+def insert(c, name, email, roll_no, image, embeddings):
+    c.execute("INSERT INTO students VALUES (?, ?, ?, ?, ?)", (name, email, roll_no, image, embeddings))