Update Predict.py

Predict.py

@@ -8,6 +8,9 @@ import sys
 import joblib
 from DL_models import CustomResNet
 
+root = '/Users/jianpingye/Desktop/Marketing_Research/XGBoost_Gaze_Prediction_Platform/Gaze-Time-Prediction-for-Advertisement/XGBoost_Prediction_Model'
+sys.path.append(root)
+
 #Ad/Brand Gaze Prediction
 
 #Now the model is only able to process magazine images or images with full-page counterpages
@@ -264,4 +267,53 @@ def CNN_Prediction(adv_imgs, ctpg_imgs, ad_locations, Gaze_Type='AG'): #Gaze_Typ
             pred = torch.exp(pred*a_temp+b_temp) - 1
             gaze += pred/10
 
-    return gaze
+    return gaze
+
+def HeatMap_CNN(adv_imgs, ctpg_imgs, ad_locations, Gaze_Type='AG'):
+    if torch.cuda.is_available():
+        device = 'cuda'
+    elif torch.backends.mps.is_available():
+        device = 'mps'
+    else:
+        device = 'cpu'
+
+    net = CustomResNet()
+    net.load_state_dict(torch.load('CNN_Gaze_Model/Fine-tune_'+Gaze_Type+'/Model_'+str(0)+'.pth',map_location=torch.device('cpu')))
+    net = net.to(device)
+    pred = net(adv_imgs/255.0,ctpg_imgs/255.0,ad_locations)
+
+    pred.backward()
+
+    # pull the gradients out of the model
+    gradients = net.get_activations_gradient()
+
+    # pool the gradients across the channels
+    pooled_gradients = torch.mean(gradients, dim=[0, 2, 3])
+
+    # get the activations of the last convolutional layer
+    activations = net.get_activations(adv_imgs).detach()
+
+    # weight the channels by corresponding gradients
+    for i in range(512):
+        activations[:, i, :, :] *= pooled_gradients[i]
+
+    # average the channels of the activations
+    heatmap = torch.mean(activations, dim=1).squeeze().to('cpu')
+
+    # relu on top of the heatmap
+    # expression (2) in https://arxiv.org/pdf/1610.02391.pdf
+    heatmap = np.maximum(heatmap, 0)
+
+    # normalize the heatmap
+    heatmap /= torch.max(heatmap)
+
+    img = torch.permute(adv_imgs[0],(1,2,0)).to(torch.uint8).numpy()
+    img = cv.cvtColor(img, cv.COLOR_BGR2RGB)
+    heatmap = cv.resize(heatmap.numpy(), (img.shape[1], img.shape[0]))
+    heatmap = np.uint8(255 * heatmap)
+    heatmap = cv.applyColorMap(heatmap, cv.COLORMAP_TURBO)
+    superimposed_img = heatmap * 0.8 + img * 0.5
+    superimposed_img /= np.max(superimposed_img)
+    superimposed_img = np.uint8(255 * superimposed_img)
+
+    return superimposed_img