Changes

app.py

@@ -1,6 +1,7 @@
 import cv2
 import numpy as np
 import gradio as gr
+from gradio import Interface, Input, Output, Image
 from PIL import Image, ImageOps
 import matplotlib.pyplot as plt
 import torch
@@ -8,11 +9,63 @@ import torch.nn as nn
 import torch.nn.functional as F
 from torchvision import transforms
 import os
+import time
+import io
+import base64
+
+class Net2(nn.Module):
+    def __init__(self):
+        super(Net2, self).__init__()
+        self.conv1 = nn.Conv2d(3, 64, 3, padding=1)
+        self.bn1 = nn.BatchNorm2d(64)
+        self.pool1 = nn.MaxPool2d(2, 2)
+        self.dropout1 = nn.Dropout(0.25)
+
+        self.conv2 = nn.Conv2d(64, 64, 3, padding=1)
+        self.bn2 = nn.BatchNorm2d(64)
+        self.pool2 = nn.MaxPool2d(2, 2)
+        self.dropout2 = nn.Dropout(0.25)
+
+        self.conv3 = nn.Conv2d(64, 64, 3, padding=1)
+        self.bn3 = nn.BatchNorm2d(64)
+        self.pool3 = nn.MaxPool2d(2, 2)
+        self.dropout3 = nn.Dropout(0.25)
+
+        self.conv4 = nn.Conv2d(64, 64, 3, padding=1)
+        self.bn4 = nn.BatchNorm2d(64)
+        self.pool4 = nn.MaxPool2d(2, 2)
+        self.dropout4 = nn.Dropout(0.25)
+
+        self.flatten = nn.Flatten()
+
+        self.fc1 = nn.Linear(64 * 5 * 5, 200)
+        self.fc2 = nn.Linear(200, 150)
+        self.fc3 = nn.Linear(150, 2)
+
+    def forward(self, x):
+        x = F.relu(self.bn1(self.conv1(x)))
+        x = self.pool1(x)
+        x = self.dropout1(x)
+
+        x = F.relu(self.bn2(self.conv2(x)))
+        x = self.pool2(x)
+        x = self.dropout2(x)
+
+        x = F.relu(self.bn3(self.conv3(x)))
+        x = self.pool3(x)
+        x = self.dropout3(x)
+
+        x = F.relu(self.bn4(self.conv4(x)))
+        x = self.pool4(x)
+        x = self.dropout4(x)
+
+        x = self.flatten(x)
+        x = F.relu(self.fc1(x))
+        x = F.relu(self.fc2(x))
+        x = F.softmax(self.fc3(x), dim=1)
+        return x
 
-# Add your model classes (Net and Net2) here.
 
-# Loading model
-model = None
 model2 = None
 model2_path = "model4.pth"
 
@@ -30,23 +83,57 @@ if os.path.exists(model2_path):
 else:
     print("Model file not found at", model2_path)
 
-# Add the scanmap function here.
 
-def process_image(image: Image.Image):
-    image_np = np.array(image)
+def process_image(input_image):
+    image = Image.open(io.BytesIO(input_image)).convert("RGB")
+
     start_time = time.time()
-    heatmap = scanmap(image_np, model)
+    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 heatmap_img, elapsed_time
+    return image, heatmap_img, int(elapsed_time)
 
-inputs = gr.inputs.Image(label="Upload Image")
+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.outputs.Image(label="Heatmap"),
-    gr.outputs.Textbox(label="Elapsed Time (seconds)")
+    gr.Image(label="Original Image"),
+    gr.Image(label="Heatmap"),
+    gr.Textbox(label="Elapsed Time")
 ]
 
-iface = gr.Interface(fn=process_image, inputs=inputs, outputs=outputs, title="ShipNet Heatmap")
+iface = gr.Interface(fn=gradio_process_image, inputs=inputs, outputs=outputs)
 iface.launch()
+
+

model4.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:90d2dcf5a7c630275f3238f399b59b5de6da5688bc9dd10c95476cffc675e342
+size 1867947