dump meta to sqlite.

docker-compose.yml

@@ -14,6 +14,7 @@ services:
       - ../artifacts/images:/home/artifacts/images
       - ../artifacts/logs:/home/artifacts/logs
       - ../artifacts/models:/home/artifacts/models
+      - ../artifacts/db:/home/artifacts/db
     depends_on:
       - redis
     tty: true

inference/Dockerfile

@@ -9,6 +9,7 @@ ENV WORKING_DIR="/home"
 ENV LOGS_DIR="/home/artifacts/logs/"
 ENV IMAGES_DIR="/home/artifacts/images/"
 ENV MODELS_DIR="/home/artifacts/models/"
+ENV DB_DIR="/home/artifacts/db/"
 ENV SEGMENTATION_MODEL_NAME="pointrend_weights.pth"
 ENV CLASSIFICATION_MODEL_NAME="swinv2_weights.pth"
 ENV SPECIES_LIST="species.json"

inference/inference.py

@@ -3,10 +3,13 @@
 import numpy as np
 from io import BytesIO
 import base64
+import time
 import logging
 
 from segmentation.inference import predict as predict_segmentation
 from classification.inference import predict as predict_classification
+from storage.meta import Meta, ClassificationMeta
+from storage.db import dump_meta
 
 logger = logging.getLogger(__name__)
 
@@ -24,22 +27,40 @@ def infer(jpg_filename):
     jpg_image = Image.open(os.path.join(os.environ["IMAGES_DIR"], jpg_filename))
     numpy_image = np.array(jpg_image)
 
+    seg_meta = Meta()
+    seg_meta.img_id = jpg_filename.split(".")[0]
+
     logger.info(f"Starting segmentation for {jpg_filename}.")
+    seg_meta.start_mili = time.time_ns() // 1_000_000
     numpy_binarymasks = predict_segmentation(numpy_image, 0.95)
-    logger.info(f"Finished segmentation for {jpg_filename} - Predicted {len(numpy_binarymasks)} pitchers.")
+    seg_meta.end_mili = time.time_ns() // 1_000_000
+    logger.info(f"Finished segmentation for {jpg_filename}.")
+
+    seg_meta.num_masks = len(numpy_binarymasks)
 
+    logger.info(f"Starting classification for {jpg_filename}.")
     output = []
     for i, binary_mask in enumerate(numpy_binarymasks):
+        clas_meta = ClassificationMeta()
+        clas_meta.pitcher_id = str(i)
         current_pitcher_pred = {}
 
-        logger.info(f"Starting classification for {jpg_filename} - mask #{i+1}.")
+        clas_meta.start_mili = time.time_ns() // 1_000_000
         classification_results = predict_classification(numpy_image, binary_mask['numpy_mask'], 5)
-        current_pitcher_pred["classification"] = predict_classification(numpy_image, binary_mask['numpy_mask'], 5)
-        logger.info(f"Finished classification for {jpg_filename} - mask #{i+1} - {classification_results}.")
+        clas_meta.end_mili = time.time_ns() // 1_000_000
+
+        clas_meta.pred_species_1 = classification_results[0]['species']
+        clas_meta.pred_species_1_conf = classification_results[0]['confidence']
+        clas_meta.pred_species_2 = classification_results[1]['species']
+        clas_meta.pred_species_2_conf = classification_results[1]['confidence']
 
+        current_pitcher_pred["classification"] = classification_results
         current_pitcher_pred["segmentation"] = {
             'mask': binary_numpy_arr_to_base64png(binary_mask['numpy_mask']),
             'confidence': binary_mask['confidence']
         }
         output.append(current_pitcher_pred)
+        seg_meta.classifications.append(clas_meta)
+    logger.info(f"Finished classification for {jpg_filename}.")
+    dump_meta(seg_meta)
     return output

inference/storage/db.py

@@ -0,0 +1,68 @@
+import sqlite3
+import os
+import logging
+
+logger = logging.getLogger(__name__)
+
+class DB:
+    def __init__(self, sqlite_db_path):
+        self.con = sqlite3.connect(sqlite_db_path)
+        self.cur = self.con.cursor()
+
+        self.cur.execute('''
+        CREATE TABLE IF NOT EXISTS Meta (
+            img_id TEXT PRIMARY KEY,
+            start_mili INTEGER,
+            end_mili INTEGER,
+            num_masks INTEGER
+        )
+        ''')
+        self.cur.execute('''
+        CREATE TABLE IF NOT EXISTS ClassificationMeta (
+            img_id TEXT,
+            pitcher_id TEXT,
+            pred_species_1 TEXT,
+            pred_species_1_conf REAL,
+            pred_species_2 TEXT,
+            pred_species_2_conf REAL,
+            start_mili INTEGER,
+            end_mili INTEGER,
+            PRIMARY KEY (img_id, pitcher_id),
+            FOREIGN KEY (img_id) REFERENCES Meta (img_id)
+        )
+        ''')
+
+        self.con.commit()
+
+    def insert_into_meta(self, meta):
+        insert_sql = '''
+        INSERT INTO Meta (img_id, start_mili, end_mili, num_masks)
+        VALUES (?, ?, ?, ?)
+        '''
+        self.cur.execute(insert_sql, (meta.img_id, meta.start_mili, meta.end_mili, meta.num_masks))
+        self.con.commit()
+
+    def insert_into_class_meta(self, img_id, classification_meta):
+        insert_sql = '''
+        INSERT INTO ClassificationMeta (img_id, pitcher_id, pred_species_1, pred_species_1_conf, pred_species_2, pred_species_2_conf, start_mili, end_mili)
+        VALUES (?, ?, ?, ?, ?, ?, ?, ?)
+        '''
+        self.cur.execute(insert_sql, (
+            img_id,
+            classification_meta.pitcher_id,
+            classification_meta.pred_species_1,
+            classification_meta.pred_species_1_conf,
+            classification_meta.pred_species_2,
+            classification_meta.pred_species_2_conf,
+            classification_meta.start_mili,
+            classification_meta.end_mili,
+        ))
+        self.con.commit()
+
+db = DB(os.path.join(os.environ['DB_DIR'], 'inference.db'))
+
+def dump_meta(meta):
+    db.insert_into_meta(meta)
+    for class_meta in meta.classifications:
+        db.insert_into_class_meta(meta.img_id, class_meta)
+    logger.info(f"Dumped meta for {meta.img_id}")

inference/storage/meta.py

@@ -0,0 +1,17 @@
+class Meta:
+    def __init__(self):
+        self.img_id = None
+        self.start_mili = None
+        self.end_mili = None
+        self.num_masks = None
+        self.classifications = []
+
+class ClassificationMeta:
+    def __init__(self):
+        self.pitcher_id = None
+        self.pred_species_1 = None
+        self.pred_species_1_conf = None
+        self.pred_species_2 = None
+        self.pred_sepcies_2_conf = None
+        self.start_mili = None
+        self.end_mili = None