Merge pull request #6 from praxxus11/fix-num-threads

Fix num threads.
praxxus11 · Aug 6, 2024 · 4728e57 · 4728e57
2 parents 0676238 + 00751c9
commit 4728e57
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 4 deletions.
diff --git a/inference/inference.py b/inference/inference.py
@@ -53,6 +53,8 @@ def infer(jpg_filename):
         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']
+        clas_meta.pred_species_3 = classification_results[2]['species']
+        clas_meta.pred_species_3_conf = classification_results[2]['confidence']
 
         current_pitcher_pred["classification"] = classification_results
         current_pitcher_pred["segmentation"] = {

diff --git a/inference/init.py b/inference/init.py
@@ -1,5 +1,6 @@
 import logging
 import os
+import torch
 
 def init():
     logging.basicConfig(
@@ -8,4 +9,5 @@ def init():
         level=logging.DEBUG,
         format='%(asctime)s.%(msecs)03d - %(name)s - %(levelname)s - %(message)s',
         datefmt='%Y-%m-%d %H:%M:%S'
-    )
+    )
+    torch.set_num_threads(4)
diff --git a/inference/storage/db.py b/inference/storage/db.py
@@ -25,13 +25,14 @@ def __init__(self, sqlite_db_path):
             pred_species_1_conf REAL,
             pred_species_2 TEXT,
             pred_species_2_conf REAL,
+            pred_species_3 TEXT,
+            pred_species_3_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):
@@ -44,8 +45,8 @@ def insert_into_meta(self, meta):
 
     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 (?, ?, ?, ?, ?, ?, ?, ?)
+        INSERT INTO ClassificationMeta (img_id, pitcher_id, pred_species_1, pred_species_1_conf, pred_species_2, pred_species_2_conf, pred_species_3, pred_species_3_conf, start_mili, end_mili)
+        VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
         '''
         self.cur.execute(insert_sql, (
             img_id,
@@ -54,6 +55,8 @@ def insert_into_class_meta(self, img_id, classification_meta):
             classification_meta.pred_species_1_conf,
             classification_meta.pred_species_2,
             classification_meta.pred_species_2_conf,
+            classification_meta.pred_species_3,
+            classification_meta.pred_species_3_conf,
             classification_meta.start_mili,
             classification_meta.end_mili,
         ))

diff --git a/inference/storage/meta.py b/inference/storage/meta.py
@@ -13,5 +13,7 @@ def __init__(self):
         self.pred_species_1_conf = None
         self.pred_species_2 = None
         self.pred_sepcies_2_conf = None
+        self.pred_species_3 = None
+        self.pred_sepcies_3_conf = None
         self.start_mili = None
         self.end_mili = None