Update app.py

app.py

@@ -81,23 +81,16 @@ transform = transforms.Compose([
 
 # Define the prediction function
 def predict_count(input_image):
-    # Preprocess the input image
     image = transform(input_image).unsqueeze(0).cpu()
-
-    # Perform the forward pass
     output = csrmodel(image)
-
-    # Calculate the predicted count
     predicted_count = int(output.detach().cpu().sum().numpy())
+    density_map = output.detach().cpu().numpy().reshape(output.shape[2], output.shape[3])
+    density_map_color = plt.cm.jet(density_map / np.max(density_map))
+    return predicted_count, density_map_color
 
-    return predicted_count
-
-# Define the input and output interfaces for Gradio
-input_interface = gr.inputs.Image()
-output_interface = gr.outputs.Textbox()
-
-# Create the Gradio app
-grapp = gr.Interface(fn=predict_count, inputs=input_interface, outputs=output_interface)
+output_interface = gr.outputs.Textbox(label="Predicted Count")
+density_map_interface = gr.outputs.Image(label="Density Map")
 
+grapp = gr.Interface(fn=predict_count, inputs=input_interface, outputs=[output_interface, density_map_interface])
 # Launch the app
 grapp.launch()