face recognition

app/Hackathon_setup/face_recognition.py

@@ -104,11 +104,41 @@ def get_similarity(img1, img2):
 ##Caution: Don't change the definition or function name; for loading the model use the current_path for path example is given in comments to the function
 def get_face_class(img1):
     device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
-    
-    det_img1 = detected_face(img1)
-    if(det_img1 == 0):
+    BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+
+    # 1 Load the Decision Tree classifier
+    clf_path = os.path.join(BASE_DIR, "siamese_decision_tree.pkl")
+    clf = joblib.load(clf_path)
+
+    # 2 Load the Siamese feature extractor
+    myModel = Siamese().to(device)
+    ckpt_path = os.path.join(BASE_DIR, "siamese_model.t7")
+    ckpt = torch.load(ckpt_path, map_location=device)
+    myModel.load_state_dict(ckpt['net_dict'])
+    myModel.eval()
+
+    # 3 Face detection (if available)
+    det_img1 = detected_face(img1)   # returns cropped face or 0 if not detected
+    if det_img1 == 0:
+        # fallback: use original image
         det_img1 = Image.fromarray(cv2.cvtColor(img1, cv2.COLOR_BGR2GRAY))
-    ##YOUR CODE HERE, return face class here
-    ##Hint: you need a classifier finetuned for your classes, it takes o/p of siamese as i/p to it
-    ##Better Hint: Siamese experiment is covered in one of the labs
-    return "YET TO BE CODED"
+    else:
+        det_img1 = Image.fromarray(cv2.cvtColor(det_img1, cv2.COLOR_BGR2GRAY))
+
+    # 4 Transform the face
+    img_tensor = trnscm(det_img1).unsqueeze(0)
+
+    # 5 Extract embeddings
+    with torch.no_grad():
+        embedding = myModel.forward_once(img_tensor)
+        embedding = embedding.view(embedding.size(0), -1).cpu().numpy()  # shape (1, embedding_dim)
+
+    # 6 Predict class using Decision Tree
+    pred_label = clf.predict(embedding)[0]
+
+    # 7 Optional: return class name (if available)
+    # If you have the dataset available:
+    # class_names = finalClassifierDset.classes
+    # return class_names[pred_label]
+    class_names = ['Vikram', 'Vikram', 'Vikram']
+    return pred_label