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="5060000" />
+      <workItem from="1686602174110" duration="6264000" />
     </task>
     <servers />
   </component>

app.py

@@ -15,7 +15,6 @@ import cv2
 import matplotlib.pyplot as plt
 import matplotlib.patches as patches
 from functools import partial
-import tempfile
 
 class Net2(nn.Module):
     def __init__(self):
@@ -123,8 +122,8 @@ class Net(nn.Module):
 model = None
 model_path = "models1.pth"
 
-# model2 = None
-# model2_path = "model4.pth"
+model2 = None
+model2_path = "model4.pth"
 
 if os.path.exists(model_path):
     state_dict = torch.load(model_path, map_location=torch.device('cpu'))
@@ -143,130 +142,107 @@ else:
 
 # 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, threshold=0.5):
-    satellite_image = cv2.imread(image_path)
-    satellite_image = satellite_image.astype(np.float32) / 255.0
+def scanmap(image_np, model, threshold=0.5):
+    image_np = image_np.astype(np.float32) / 255.0
 
     window_size = (80, 80)
     stride = 10
 
-    height, width, channels = satellite_image.shape
-
-    # ensure model is in float32 precision
-    model.float()
+    height, width, channels = image_np.shape
 
     fig, ax = plt.subplots(1)
-    ax.imshow(satellite_image)
-
-    ship_images = []
+    ax.imshow(image_np)
 
     for y in range(0, height - window_size[1] + 1, stride):
         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 = 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)
 
-            # if probability is greater than threshold, draw a bounding box and add to ship_images
+            # if probability is greater than threshold, draw a bounding box
             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)
 
-    output_path = "output.png"
-    plt.savefig(output_path)
-    plt.close()
+    # Save the image to a byte buffer
+    buf = io.BytesIO()
+    plt.savefig(buf, format='png')
+    buf.seek(0)
 
-    return output_path
+    return Image.open(buf)  # return PIL Image
 
+def process_image(input_image):
+    image = Image.fromarray(input_image).convert("RGB")
 
-def process_image(input_image, model, threshold=0.5):
     start_time = time.time()
-    ship_images = scanmap(input_image, model, threshold)
-    elapsed_time = time.time() - start_time
-
-    return ship_images, int(elapsed_time)
-
-
-def gradio_process_image(input_image, model, threshold=0.5):
-    start_time = time.time()
-
-    # save numpy array to a temporary file
-    temp = tempfile.NamedTemporaryFile(delete=False, suffix=".png")
-    temp.close()
-    cv2.imwrite(temp.name, cv2.cvtColor(input_image * 255, cv2.COLOR_RGB2BGR))
-
-    # pass file path to scanmap
-    output_image_path = scanmap(temp.name, model, threshold)
-
+    detected_ships_image = scanmap(np.array(image), model)
     elapsed_time = time.time() - start_time
 
-    # delete temporary file after processing
-    os.unlink(temp.name)
+    return image, detected_ships_image, int(elapsed_time)
 
-    return output_image_path, f"Elapsed Time (seconds): {elapsed_time}"
+def gradio_process_image(input_image):
+    original, detected_ships_image, elapsed_time = process_image(input_image)
+    return original, detected_ships_image, f"Elapsed Time (seconds): {elapsed_time}"
 
-inputs = gr.inputs.Image(label="Upload Image")
+inputs = gr.Image(label="Upload Image")
 outputs = [
-    gr.outputs.Image(type='filepath', label="Detected Ships"),
-    gr.outputs.Textbox(label="Elapsed Time")
+    gr.Image(label="Original Image"),
+    gr.Image(label="Heatmap"),
+    gr.Textbox(label="Elapsed Time")
 ]
 
-# Use 0.5 as the threshold, but adjust according to your needs
-gradio_process_image_partial = partial(gradio_process_image, model=model, threshold=0.5)
+iface = gr.Interface(fn=gradio_process_image, inputs=inputs, outputs=outputs)
+iface.launch()
 
-iface = gr.Interface(fn=gradio_process_image_partial, inputs=inputs, outputs=outputs)
-iface.launch()