Update app.py

app.py

@@ -36,18 +36,33 @@ from models.psp import pSp
 # Huggingface login
 login(token=os.getenv("TOKENKEY"))
 
+########################################################################
+############## tensorflow model for age calculation #######################
+
 # If 'mse' is a custom function needed, 
 #custom_objects = {'mse': MeanSquaredError()}
 #new_age_model = load_model("age_prediction_model.h5")
+########################################################################
+
+
+########################################################################
+############## pytorch model for age calculation #######################
+age_calc_model = torch.load('Custom_Age_prediction_model.pth')
+
+########################################################################
 
 # Download models from Huggingface
 age_prototxt = hf_hub_download(repo_id="AshanGimhana/Age_Detection_caffe", filename="age.prototxt")
 caffe_model = hf_hub_download(repo_id="AshanGimhana/Age_Detection_caffe", filename="dex_imdb_wiki.caffemodel")
 sam_ffhq_aging = hf_hub_download(repo_id="AshanGimhana/Face_Agin_model", filename="sam_ffhq_aging.pt")
 
+########################################################################
+############## caffe model for age calculation #######################
 
 # Age prediction model setup
-age_net = cv2.dnn.readNetFromCaffe(age_prototxt, caffe_model)
+#age_net = cv2.dnn.readNetFromCaffe(age_prototxt, caffe_model)
+########################################################################
+
 
 # Face detection and landmarks predictor setup
 detector = dlib.get_frontal_face_detector()
@@ -116,16 +131,37 @@ def get_mouth_region(image):
 
 # Function to predict age
 
+# old tensorflow function for age predict
+
+#def predict_age(image):
+    #image = np.array(image.resize((64, 64)))
+    #image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) 
+    #image = image / 255.0  
+    #image = np.expand_dims(image, axis=0)
+
+    ##### Predict age
+    #val = new_age_model.predict(np.array(image))
+    #age = val[0][0]
+    #return int(age)
+
 def predict_age(image):
-    image = np.array(image.resize((64, 64)))
-    image = cv2.cvtColor(image, cv2.COLOR_RGB2GRAY) 
-    image = image / 255.0  
-    image = np.expand_dims(image, axis=0)
+    age_calc_model.eval()
+    # Load and preprocess the image
+    image = cv2.imread(image, cv2.IMREAD_GRAYSCALE)  # Load as grayscale
+    image = cv2.resize(image, (64, 64))  # Resize to 64x64
+    image = image / 255.0  # Normalize pixel values to [0, 1]
+    image = np.expand_dims(image, axis=0)  # Add batch dimension
+    image = np.expand_dims(image, axis=0)  # Add channel dimension
+    image = torch.tensor(image, dtype=torch.float32).to(device)
+
+    # Convert to tensor
+    image_tensor = torch.tensor(image, dtype=torch.float32)
 
     # Predict age
-    val = new_age_model.predict(np.array(image))
-    age = val[0][0]
-    return int(age)
+    with torch.no_grad():
+        predicted_age = age_calc_model(image_tensor)
+    
+    return int(predicted_age.item()) 
 
 # Function for color correction
 def color_correct(source, target):