Spaces:
Sleeping
Sleeping
| import cv2 | |
| import numpy as np | |
| import gradio as gr | |
| from PIL import Image, ImageOps | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from torchvision import transforms | |
| import os | |
| import time | |
| import io | |
| import base64 | |
| import torch | |
| 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): | |
| 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 | |
| class Net(nn.Module): | |
| def __init__(self): | |
| super(Net, self).__init__() | |
| self.conv1 = nn.Conv2d(3, 512, 3, padding=1) | |
| self.bn1 = nn.BatchNorm2d(512) | |
| self.pool1 = nn.MaxPool2d(2, 2) | |
| self.dropout1 = nn.Dropout(0.25) | |
| self.conv2 = nn.Conv2d(512, 256, 3, padding=1) | |
| self.bn2 = nn.BatchNorm2d(256) | |
| self.pool2 = nn.MaxPool2d(2, 2) | |
| self.dropout2 = nn.Dropout(0.25) | |
| self.conv3 = nn.Conv2d(256, 128, 3, padding=1) | |
| self.bn3 = nn.BatchNorm2d(128) | |
| self.pool3 = nn.MaxPool2d(2, 2) | |
| self.dropout3 = nn.Dropout(0.25) | |
| self.conv4 = nn.Conv2d(128, 64, 3, padding=1) | |
| self.bn4 = nn.BatchNorm2d(64) | |
| self.pool4 = nn.MaxPool2d(2, 2) | |
| self.dropout4 = nn.Dropout(0.20) | |
| self.flatten = nn.Flatten() | |
| self.fc1 = nn.Linear(1600, 300) | |
| self.fc2 = nn.Linear(300, 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 | |
| model = None | |
| model_path = "models1.pth" | |
| # model2 = None | |
| # model2_path = "model4.pth" | |
| if os.path.exists(model_path): | |
| state_dict = torch.load(model_path, map_location=torch.device('cpu')) | |
| new_state_dict = {} | |
| for key, value in state_dict.items(): | |
| new_key = key.replace("module.", "") | |
| new_state_dict[new_key] = value | |
| model = Net() | |
| model.load_state_dict(new_state_dict) | |
| model.eval() | |
| 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 | |
| # | |
| # 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 | |
| window_size = (80, 80) | |
| stride = 10 | |
| height, width, channels = satellite_image.shape | |
| # ensure model is in float32 precision | |
| model.float() | |
| fig, ax = plt.subplots(1) | |
| ax.imshow(satellite_image) | |
| ship_images = [] | |
| 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) | |
| 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') | |
| ax.add_patch(rect) | |
| ship_images.append(cropped_window) | |
| 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() | |
| 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) | |
| elapsed_time = time.time() - start_time | |
| # delete temporary file after processing | |
| os.unlink(temp.name) | |
| return output_image_path, f"Elapsed Time (seconds): {elapsed_time}" | |
| inputs = gr.inputs.Image(label="Upload Image") | |
| outputs = [ | |
| gr.outputs.Image(type='filepath', label="Detected Ships"), | |
| gr.outputs.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_partial, inputs=inputs, outputs=outputs) | |
| iface.launch() |