Changes on scan

app.py

@@ -133,20 +133,34 @@ else:
     print("Model file not found at", model_path)
 
 
-def scanmap(satellite_image, model):
-    satellite_image = satellite_image.astype(np.float32) / 255.0
+def process_image(input_image):
+    image = Image.fromarray(input_image).convert("RGB")
+
+    start_time = time.time()
+    heatmap = scanmap(np.array(image), model)
+    elapsed_time = time.time() - start_time
+    heatmap_img = Image.fromarray(np.uint8(plt.cm.hot(heatmap) * 255)).convert('RGB')
+
+    heatmap_img = heatmap_img.resize(image.size)
+
+    return image, heatmap_img, int(elapsed_time)
+
+
+def scanmap(image_np, model):
+    image_np = image_np.astype(np.float32) / 255.0
+
     window_size = (80, 80)
     stride = 10
 
-    height, width, channels = satellite_image.shape
+    height, width, channels = image_np.shape
+
     probabilities_map = []
 
     for y in range(0, height - window_size[1] + 1, stride):
         row_probabilities = []
         for x in range(0, width - window_size[0] + 1, stride):
-            cropped_window = satellite_image[y:y + window_size[1], x:x + window_size[0]]
-            cropped_window_torch = torch.tensor(cropped_window.transpose(2, 0, 1), dtype=torch.float32).unsqueeze(
-                0).cpu()
+            cropped_window = image_np[y:y + window_size[1], x:x + window_size[0]]
+            cropped_window_torch = transforms.ToTensor()(cropped_window).unsqueeze(0)
 
             with torch.no_grad():
                 probabilities = model(cropped_window_torch)
@@ -158,26 +172,18 @@ def scanmap(satellite_image, model):
     probabilities_map = np.array(probabilities_map)
     return probabilities_map
 
-
 def gradio_process_image(input_image):
-    if input_image is None:
-        return None, None
-
-    if isinstance(input_image, np.ndarray):
-        input_image = Image.fromarray(input_image)
-
-    image_np = np.array(input_image.convert("RGB"))
-    heatmap = scanmap(image_np, model)
-    heatmap_normalized = (heatmap - np.min(heatmap)) / (np.max(heatmap) - np.min(heatmap))
-    heatmap_img = Image.fromarray(np.uint8(plt.cm.hot(heatmap_normalized) * 255)).convert('RGB')
-    return input_image, heatmap_img
-
+    original, heatmap, elapsed_time = process_image(input_image)
+    return original, heatmap, f"Elapsed Time (seconds): {elapsed_time}"
 
 inputs = gr.Image(label="Upload Image")
 outputs = [
     gr.Image(label="Original Image"),
     gr.Image(label="Heatmap"),
+    gr.Textbox(label="Elapsed Time")
 ]
 
 iface = gr.Interface(fn=gradio_process_image, inputs=inputs, outputs=outputs)
 iface.launch()
+
+