Changed scan func

.idea/workspace.xml

@@ -65,7 +65,7 @@
       <workItem from="1683665300392" duration="7649000" />
       <workItem from="1683708398011" duration="1235000" />
       <workItem from="1684437905081" duration="110000" />
-      <workItem from="1686602174110" duration="3739000" />
+      <workItem from="1686602174110" duration="4678000" />
     </task>
     <servers />
   </component>

app.py

@@ -192,7 +192,9 @@ else:
 # 
 # iface = gr.Interface(fn=gradio_process_image, inputs=inputs, outputs=outputs)
 # iface.launch()
-def scanmap(satellite_image, model, threshold=0.5):
+
+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)
@@ -200,6 +202,7 @@ def scanmap(satellite_image, model, threshold=0.5):
 
     height, width, channels = satellite_image.shape
 
+
     fig, ax = plt.subplots(1)
     ax.imshow(satellite_image)
 
@@ -209,19 +212,23 @@ def scanmap(satellite_image, model, threshold=0.5):
         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)
+            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)
 
+            # 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')
+                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
+    output_path = "output.png"
+    plt.savefig(output_path)
+    plt.close()
 
+    return output_path
 
 def process_image(input_image, model, threshold=0.5):
     start_time = time.time()
@@ -231,22 +238,19 @@ def process_image(input_image, model, threshold=0.5):
     return ship_images, int(elapsed_time)
 
 
-def gradio_process_image(input_image, model, threshold=0.5):
-    ship_images, elapsed_time = process_image(input_image, model, threshold)
+def gradio_process_image(input_image_path, model, threshold=0.5):
+    start_time = time.time()
+    output_image_path = scanmap(input_image_path, model, threshold)
+    elapsed_time = time.time() - start_time
 
-    if len(ship_images) > 0:
-        # Convert first image to format compatible with Gradio
-        output_image = Image.fromarray((ship_images[0] * 255).astype(np.uint8))
-    else:
-        output_image = None
+    output_image = Image.open(output_image_path) if output_image_path else None
 
     return output_image, f"Elapsed Time (seconds): {elapsed_time}"
 
-
-inputs = gr.Image(label="Upload Image")
+inputs = gr.inputs.Image(label="Upload Image")
 outputs = [
-    gr.Image(label="Detected Ships"),
-    gr.Textbox(label="Elapsed Time")
+    gr.outputs.Image(label="Detected Ships"),
+    gr.outputs.Textbox(label="Elapsed Time")
 ]
 
 # Use 0.5 as the threshold, but adjust according to your needs