Changes on scan

app.py

@@ -133,34 +133,20 @@ else:
     print("Model file not found at", model_path)
 
 
-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
-
+def scanmap(satellite_image, model):
+    satellite_image = satellite_image.astype(np.float32) / 255.0
     window_size = (80, 80)
     stride = 10
 
-    height, width, channels = image_np.shape
-
+    height, width, channels = satellite_image.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 = image_np[y:y + window_size[1], x:x + window_size[0]]
-            cropped_window_torch = transforms.ToTensor()(cropped_window).unsqueeze(0)
+            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()
 
             with torch.no_grad():
                 probabilities = model(cropped_window_torch)
@@ -172,18 +158,20 @@ def scanmap(image_np, model):
     probabilities_map = np.array(probabilities_map)
     return probabilities_map
 
+
 def gradio_process_image(input_image):
-    original, heatmap, elapsed_time = process_image(input_image)
-    return original, heatmap, f"Elapsed Time (seconds): {elapsed_time}"
+    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
+
 
 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()
-
-
+iface.launch()