Update app.py

app.py

@@ -13,6 +13,21 @@ print(f"CUDA device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
 
 ###################################################
 
+gpus = tf.config.list_physical_devices('GPU')
+print("Available GPUs TF:", gpus)
+
+if gpus:
+    try:
+        # Allow TensorFlow to allocate memory as needed
+        for gpu in gpus:
+            tf.config.experimental.set_memory_growth(gpu, True)
+    except RuntimeError as e:
+        print(e)  # Print error if unable to set up GPU
+else:
+    print("No GPUs available.")
+
+###################################################
+
 
 from argparse import Namespace
 import pprint
@@ -41,13 +56,14 @@ login(token=os.getenv("TOKENKEY"))
 
 # If 'mse' is a custom function needed, 
 #custom_objects = {'mse': MeanSquaredError()}
-#new_age_model = load_model("age_prediction_model.h5")
+new_age_model = load_model("age_prediction_model.h5")
 ########################################################################
 
 
 ########################################################################
 ############## pytorch model for age calculation #######################
-age_calc_model = torch.load('Custom_Age_prediction_model.pth')
+
+#age_calc_model = torch.load('Custom_Age_prediction_model.pth')
 
 ########################################################################
 
@@ -140,35 +156,35 @@ def get_mouth_region(image):
 
 # 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)
+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)
+    val = new_age_model.predict(np.array(image))
+    age = val[0][0]
+    return int(age)
 
-def predict_age(image):
-    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)
+#def predict_age(image):
+    #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)
+    #image_tensor = torch.tensor(image, dtype=torch.float32)
 
-    # Predict age
-    with torch.no_grad():
-        predicted_age = age_calc_model(image_tensor)
+    #### Predict age
+    #with torch.no_grad():
+        #predicted_age = age_calc_model(image_tensor)
     
-    return int(predicted_age.item()) 
+    #return int(predicted_age.item()) 
 
 # Function for color correction
 def color_correct(source, target):