Changed scan func

.idea/workspace.xml

@@ -34,20 +34,20 @@
     <option name="hideEmptyMiddlePackages" value="true" />
     <option name="showLibraryContents" value="true" />
   </component>
-  <component name="PropertiesComponent">{
-  &quot;keyToString&quot;: {
-    &quot;RunOnceActivity.ShowReadmeOnStart&quot;: &quot;true&quot;,
-    &quot;WebServerToolWindowFactoryState&quot;: &quot;false&quot;,
-    &quot;git-widget-placeholder&quot;: &quot;main&quot;,
-    &quot;last_opened_file_path&quot;: &quot;D:/fastApiProjects/pythonProject1/shipnet&quot;,
-    &quot;node.js.detected.package.eslint&quot;: &quot;true&quot;,
-    &quot;node.js.detected.package.tslint&quot;: &quot;true&quot;,
-    &quot;node.js.selected.package.eslint&quot;: &quot;(autodetect)&quot;,
-    &quot;node.js.selected.package.tslint&quot;: &quot;(autodetect)&quot;,
-    &quot;settings.editor.selected.configurable&quot;: &quot;shared-indexes&quot;,
-    &quot;vue.rearranger.settings.migration&quot;: &quot;true&quot;
+  <component name="PropertiesComponent"><![CDATA[{
+  "keyToString": {
+    "RunOnceActivity.ShowReadmeOnStart": "true",
+    "WebServerToolWindowFactoryState": "false",
+    "git-widget-placeholder": "main",
+    "last_opened_file_path": "D:/fastApiProjects/pythonProject1",
+    "node.js.detected.package.eslint": "true",
+    "node.js.detected.package.tslint": "true",
+    "node.js.selected.package.eslint": "(autodetect)",
+    "node.js.selected.package.tslint": "(autodetect)",
+    "settings.editor.selected.configurable": "shared-indexes",
+    "vue.rearranger.settings.migration": "true"
   }
-}</component>
+}]]></component>
   <component name="RecentsManager">
     <key name="CopyFile.RECENT_KEYS">
       <recent name="D:\fastApiProjects\pythonProject1" />
@@ -65,7 +65,7 @@
       <workItem from="1683665300392" duration="7649000" />
       <workItem from="1683708398011" duration="1235000" />
       <workItem from="1684437905081" duration="110000" />
-      <workItem from="1686602174110" duration="1036000" />
+      <workItem from="1686602174110" duration="2045000" />
     </task>
     <servers />
   </component>

app.py

@@ -2,7 +2,6 @@ import cv2
 import numpy as np
 import gradio as gr
 from PIL import Image, ImageOps
-import matplotlib.pyplot as plt
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -11,7 +10,10 @@ import os
 import time
 import io
 import base64
-
+import torch
+import cv2
+import matplotlib.pyplot as plt
+import matplotlib.patches as patches
 
 class Net2(nn.Module):
     def __init__(self):
@@ -137,53 +139,112 @@ 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 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 = 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 = 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)
+# 
+#             row_probabilities.append(probabilities[0, 1].item())
+# 
+#         probabilities_map.append(row_probabilities)
+# 
+#     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}"
+# 
+# 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()
+def scanmap(image_path, model, device, threshold=0.5):
+    satellite_image = cv2.imread(image_path)
+    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
+
+    model.to(device)  # ensure model is on correct device
+
+    fig, ax = plt.subplots(1)
+    ax.imshow(satellite_image)
 
-    probabilities_map = []
+    ship_images = []
 
     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)
+            cropped_window_torch = cropped_window_torch.to(device)  # move data to the same device as model
 
             with torch.no_grad():
                 probabilities = model(cropped_window_torch)
 
-            row_probabilities.append(probabilities[0, 1].item())
+            # if probability is greater than threshold, draw a bounding box and add to ship_images
+            if probabilities[0, 1].item() > threshold:
+                rect = patches.Rectangle((x, y), window_size[0], window_size[1], linewidth=1, edgecolor='r',
+                                         facecolor='none')
+                ax.add_patch(rect)
+                ship_images.append(cropped_window)
+
+    plt.show()
+
+    return ship_images
+
+
+def process_image(input_image, model, device, threshold=0.5):
+    start_time = time.time()
+    ship_images = scanmap(input_image, model, device, threshold)
+    elapsed_time = time.time() - start_time
+
+    return ship_images, int(elapsed_time)
+
 
-        probabilities_map.append(row_probabilities)
+def gradio_process_image(input_image, model, device, threshold=0.5):
+    ship_images, elapsed_time = process_image(input_image, model, device, threshold)
 
-    probabilities_map = np.array(probabilities_map)
-    return probabilities_map
+    return ship_images, f"Elapsed Time (seconds): {elapsed_time}"
 
-def gradio_process_image(input_image):
-    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.Image(label="Detected Ships"),
     gr.Textbox(label="Elapsed Time")
 ]